mtfyun/ds4
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
4f806387f71d3f0e8517650d5cd13b559003cdfd
ds4/ds4_metal.m
T

52092 lines
2.3 MiB
Plaintext
Raw Blame History

#import <Foundation/Foundation.h>
#import <Metal/Metal.h>
#include <execinfo.h>
#include <stdint.h>
#include <inttypes.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <math.h>
#include <float.h>
#include <fcntl.h>
#include <limits.h>
#include <time.h>
#include <pthread.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/sysctl.h>
#include <mach/mach.h>
#include <mach-o/dyld.h>
#include <objc/runtime.h>
#include "ds4.h"
#include "ds4_gpu.h"
#include "ds4_image.h"
/*
* Objective-C Metal glue for the C engine.
*
* The C code owns model semantics and graph scheduling. This file owns only
* Metal objects: device/queue/library setup, mmap-backed weight views, command
* batching, persistent tensors, scratch buffers, and thin wrappers around the
* kernel files in the metal directory. Keeping this boundary narrow makes the
* inference path readable from C while still using Objective-C where Metal
* requires it.
*/
enum {
DS4_METAL_TENSOR_Q4_0 = 2,
DS4_METAL_TENSOR_Q8_0 = 8,
DS4_METAL_TENSOR_Q2_K = 10,
DS4_METAL_TENSOR_Q4_K = 12,
DS4_METAL_TENSOR_Q5_K = 13,
DS4_METAL_TENSOR_Q6_K = 14,
DS4_METAL_TENSOR_Q8_K = 15,
DS4_METAL_TENSOR_IQ2_XXS = 16,
DS4_METAL_TENSOR_MXFP4 = 39,
};
@class DS4MetalQ4ExpertTable;
static id<MTLDevice> g_device;
static id<MTLCommandQueue> g_queue;
static id<MTLLibrary> g_library;
static id<MTLCommandBuffer> g_batch_cb;
static id<MTLComputeCommandEncoder> g_batch_enc;
static BOOL g_batch_encoder_concurrent;
static BOOL g_batch_has_work;
/* ---- Encoder timeline diagnostic (DS4_METAL_ENCODER_TIMELINE=<file>) ----
* Gives every batch dispatch group its own compute pass with stage-boundary
* GPU timestamps, so one decode token can be read as a per-kernel GPU
* timeline (durations and the gaps between them) without ending command
* buffers. Pass boundaries add a little serialization, so absolute numbers
* are an upper bound; the anatomy is what this exists for. Off by default. */
#define DS4_TIMELINE_NAME_MAX 80
#define DS4_TIMELINE_SAMPLES_PER_BUFFER 4096
typedef struct {
uintptr_t caller;
uint32_t n_dispatch;
uint32_t tg[3];
uint32_t tpt[3];
char kernel[DS4_TIMELINE_NAME_MAX];
} ds4_timeline_rec;
@interface DS4TimelineBatch : NSObject {
@public
NSMutableArray<id<MTLCounterSampleBuffer>> *samples;
ds4_timeline_rec *recs;
uint32_t count;
uint32_t cap;
uint64_t seq;
}
@end
@implementation DS4TimelineBatch
- (void)dealloc { free(recs); }
@end
static BOOL g_timeline_enabled;
static FILE *g_timeline_file;
static pthread_mutex_t g_timeline_mutex = PTHREAD_MUTEX_INITIALIZER;
static DS4TimelineBatch *g_timeline_batch;
static uint64_t g_timeline_seq;
static NSMutableDictionary<NSNumber *, NSString *> *g_timeline_pso_names;
static id<MTLCounterSet> g_timeline_counter_set;
static BOOL g_timeline_encoder_hooked;
static void ds4_gpu_timeline_note_pso_name(id pso, NSString *name) {
if (!g_timeline_enabled || !pso || !name) return;
pthread_mutex_lock(&g_timeline_mutex);
if (!g_timeline_pso_names) g_timeline_pso_names = [NSMutableDictionary dictionary];
g_timeline_pso_names[@((uintptr_t)pso)] = name;
pthread_mutex_unlock(&g_timeline_mutex);
}
@interface NSObject (DS4TimelineSwizzle)
- (id)ds4_tl_newComputePipelineStateWithFunction:(id<MTLFunction>)fn
error:(NSError * __autoreleasing *)error NS_RETURNS_RETAINED;
- (id)ds4_tl_newComputePipelineStateWithDescriptor:(MTLComputePipelineDescriptor *)descriptor
options:(MTLPipelineOption)options
reflection:(MTLAutoreleasedComputePipelineReflection *)reflection
error:(NSError * __autoreleasing *)error NS_RETURNS_RETAINED;
- (void)ds4_tl_setComputePipelineState:(id<MTLComputePipelineState>)pso;
- (void)ds4_tl_dispatchThreadgroups:(MTLSize)tg threadsPerThreadgroup:(MTLSize)tpt;
- (void)ds4_tl_dispatchThreads:(MTLSize)threads threadsPerThreadgroup:(MTLSize)tpt;
@end
static ds4_timeline_rec *ds4_gpu_timeline_current_rec(void) {
DS4TimelineBatch *b = g_timeline_batch;
if (!b || b->count == 0) return NULL;
return &b->recs[b->count - 1];
}
@implementation NSObject (DS4TimelineSwizzle)
- (id)ds4_tl_newComputePipelineStateWithFunction:(id<MTLFunction>)fn
error:(NSError * __autoreleasing *)error {
id pso = [self ds4_tl_newComputePipelineStateWithFunction:fn error:error];
ds4_gpu_timeline_note_pso_name(pso, [fn name]);
return pso;
}
- (id)ds4_tl_newComputePipelineStateWithDescriptor:(MTLComputePipelineDescriptor *)descriptor
options:(MTLPipelineOption)options
reflection:(MTLAutoreleasedComputePipelineReflection *)reflection
error:(NSError * __autoreleasing *)error {
id pso = [self ds4_tl_newComputePipelineStateWithDescriptor:descriptor
options:options
reflection:reflection
error:error];
ds4_gpu_timeline_note_pso_name(pso, [[descriptor computeFunction] name]);
return pso;
}
- (void)ds4_tl_setComputePipelineState:(id<MTLComputePipelineState>)pso {
[self ds4_tl_setComputePipelineState:pso];
ds4_timeline_rec *rec = ds4_gpu_timeline_current_rec();
if (!rec) return;
NSString *name = nil;
pthread_mutex_lock(&g_timeline_mutex);
if (g_timeline_pso_names) name = g_timeline_pso_names[@((uintptr_t)pso)];
pthread_mutex_unlock(&g_timeline_mutex);
const char *cname = name ? [name UTF8String] : "?";
if (rec->kernel[0] == '\0') {
snprintf(rec->kernel, sizeof(rec->kernel), "%s", cname);
} else if (strstr(rec->kernel, cname) == NULL) {
size_t used = strlen(rec->kernel);
if (used + 1 < sizeof(rec->kernel))
snprintf(rec->kernel + used, sizeof(rec->kernel) - used, "+%s", cname);
}
}
- (void)ds4_tl_dispatchThreadgroups:(MTLSize)tg threadsPerThreadgroup:(MTLSize)tpt {
[self ds4_tl_dispatchThreadgroups:tg threadsPerThreadgroup:tpt];
ds4_timeline_rec *rec = ds4_gpu_timeline_current_rec();
if (!rec) return;
if (rec->n_dispatch == 0) {
rec->tg[0] = (uint32_t)tg.width; rec->tg[1] = (uint32_t)tg.height; rec->tg[2] = (uint32_t)tg.depth;
rec->tpt[0] = (uint32_t)tpt.width; rec->tpt[1] = (uint32_t)tpt.height; rec->tpt[2] = (uint32_t)tpt.depth;
}
rec->n_dispatch++;
}
- (void)ds4_tl_dispatchThreads:(MTLSize)threads threadsPerThreadgroup:(MTLSize)tpt {
[self ds4_tl_dispatchThreads:threads threadsPerThreadgroup:tpt];
ds4_timeline_rec *rec = ds4_gpu_timeline_current_rec();
if (!rec) return;
if (rec->n_dispatch == 0) {
rec->tg[0] = (uint32_t)((threads.width + tpt.width - 1) / (tpt.width ? tpt.width : 1));
rec->tg[1] = (uint32_t)((threads.height + tpt.height - 1) / (tpt.height ? tpt.height : 1));
rec->tg[2] = (uint32_t)((threads.depth + tpt.depth - 1) / (tpt.depth ? tpt.depth : 1));
rec->tpt[0] = (uint32_t)tpt.width; rec->tpt[1] = (uint32_t)tpt.height; rec->tpt[2] = (uint32_t)tpt.depth;
}
rec->n_dispatch++;
}
@end
static void ds4_gpu_timeline_swizzle(Class cls, SEL orig, SEL repl) {
Method m_orig = class_getInstanceMethod(cls, orig);
Method m_repl = class_getInstanceMethod(cls, repl);
if (!m_orig || !m_repl) {
fprintf(stderr, "ds4: timeline hook missing for %s\n", sel_getName(orig));
return;
}
if (class_addMethod(cls, orig, method_getImplementation(m_repl),
method_getTypeEncoding(m_orig))) {
class_replaceMethod(cls, repl, method_getImplementation(m_orig),
method_getTypeEncoding(m_orig));
} else {
method_exchangeImplementations(m_orig, m_repl);
}
}
/* Called once the device exists: opens the log and hooks pipeline creation so
* every pipeline made afterwards has a kernel name. */
static void ds4_gpu_timeline_probe(id<MTLDevice> device) {
static int probed;
if (probed) return;
probed = 1;
const char *path = getenv("DS4_METAL_ENCODER_TIMELINE");
if (!path || !path[0] || !device) return;
g_timeline_file = fopen(path, "a");
if (!g_timeline_file) {
fprintf(stderr, "ds4: cannot open encoder timeline %s\n", path);
return;
}
for (id<MTLCounterSet> cs in [device counterSets]) {
if ([[cs name] isEqualToString:MTLCommonCounterSetTimestamp]) g_timeline_counter_set = cs;
}
if (!g_timeline_counter_set ||
![device supportsCounterSampling:MTLCounterSamplingPointAtStageBoundary]) {
fprintf(stderr, "ds4: encoder timeline unsupported on this device\n");
fclose(g_timeline_file);
g_timeline_file = NULL;
return;
}
g_timeline_enabled = YES;
Class dev_cls = object_getClass(device);
ds4_gpu_timeline_swizzle(dev_cls,
@selector(newComputePipelineStateWithFunction:error:),
@selector(ds4_tl_newComputePipelineStateWithFunction:error:));
ds4_gpu_timeline_swizzle(dev_cls,
@selector(newComputePipelineStateWithDescriptor:options:reflection:error:),
@selector(ds4_tl_newComputePipelineStateWithDescriptor:options:reflection:error:));
fprintf(g_timeline_file, "# ds4 encoder timeline; slide=0x%llx pid=%d\n",
(unsigned long long)_dyld_get_image_vmaddr_slide(0), (int)getpid());
fprintf(g_timeline_file, "# B <seq> <n_encoders> <gpu_start_ns> <gpu_end_ns>\n");
fprintf(g_timeline_file, "# E <seq> <idx> <start_ns> <end_ns> <dur_us> <gap_us> <caller_unslid> <n_dispatch> <tg> <tpt> <kernel>\n");
fflush(g_timeline_file);
fprintf(stderr, "ds4: encoder timeline -> %s\n", path);
}
static void ds4_gpu_timeline_hook_encoder(id<MTLComputeCommandEncoder> enc) {
if (g_timeline_encoder_hooked || !enc) return;
g_timeline_encoder_hooked = YES;
Class cls = object_getClass(enc);
ds4_gpu_timeline_swizzle(cls, @selector(setComputePipelineState:),
@selector(ds4_tl_setComputePipelineState:));
ds4_gpu_timeline_swizzle(cls, @selector(dispatchThreadgroups:threadsPerThreadgroup:),
@selector(ds4_tl_dispatchThreadgroups:threadsPerThreadgroup:));
ds4_gpu_timeline_swizzle(cls, @selector(dispatchThreads:threadsPerThreadgroup:),
@selector(ds4_tl_dispatchThreads:threadsPerThreadgroup:));
}
static void ds4_gpu_timeline_resolve(DS4TimelineBatch *b, id<MTLCommandBuffer> cb) {
if (!b || !g_timeline_file) return;
pthread_mutex_lock(&g_timeline_mutex);
fprintf(g_timeline_file, "B %llu %u %.0f %.0f\n", (unsigned long long)b->seq, b->count,
cb.GPUStartTime * 1e9, cb.GPUEndTime * 1e9);
uint64_t prev_end = 0;
for (uint32_t i = 0; i < b->count; i++) {
const uint32_t sb_index = i / (DS4_TIMELINE_SAMPLES_PER_BUFFER / 2);
const uint32_t slot = (i % (DS4_TIMELINE_SAMPLES_PER_BUFFER / 2)) * 2u;
uint64_t start = 0, end = 0;
if (sb_index < [b->samples count]) {
NSData *d = [b->samples[sb_index] resolveCounterRange:NSMakeRange(slot, 2)];
if (d && [d length] >= 2 * sizeof(uint64_t)) {
const uint64_t *ts = (const uint64_t *)[d bytes];
start = ts[0];
end = ts[1];
}
}
const ds4_timeline_rec *r = &b->recs[i];
const double dur_us = (end >= start) ? (double)(end - start) / 1000.0 : -1.0;
const double gap_us = (prev_end && start >= prev_end) ? (double)(start - prev_end) / 1000.0 : 0.0;
fprintf(g_timeline_file, "E %llu %u %llu %llu %.2f %.2f 0x%llx %u %ux%ux%u %ux%ux%u %s\n",
(unsigned long long)b->seq, i,
(unsigned long long)start, (unsigned long long)end, dur_us, gap_us,
(unsigned long long)(r->caller - (uintptr_t)_dyld_get_image_vmaddr_slide(0)),
r->n_dispatch, r->tg[0], r->tg[1], r->tg[2], r->tpt[0], r->tpt[1], r->tpt[2],
r->kernel[0] ? r->kernel : "?");
if (end) prev_end = end;
}
fflush(g_timeline_file);
pthread_mutex_unlock(&g_timeline_mutex);
}
/* Attach to a freshly created batch command buffer. */
static void ds4_gpu_timeline_attach(id<MTLCommandBuffer> cb) {
if (!g_timeline_enabled || !cb) return;
DS4TimelineBatch *b = [DS4TimelineBatch new];
b->samples = [NSMutableArray array];
b->cap = 1024;
b->recs = calloc(b->cap, sizeof(ds4_timeline_rec));
b->count = 0;
b->seq = ++g_timeline_seq;
g_timeline_batch = b;
[cb addCompletedHandler:^(id<MTLCommandBuffer> done) {
ds4_gpu_timeline_resolve(b, done);
}];
}
static id<MTLComputeCommandEncoder> ds4_gpu_timeline_new_encoder(
id<MTLCommandBuffer> cb, BOOL concurrent, uintptr_t caller) {
DS4TimelineBatch *b = g_timeline_batch;
if (!b) return nil;
const uint32_t per_buffer = DS4_TIMELINE_SAMPLES_PER_BUFFER / 2;
const uint32_t sb_index = b->count / per_buffer;
if (sb_index >= [b->samples count]) {
MTLCounterSampleBufferDescriptor *d = [MTLCounterSampleBufferDescriptor new];
d.counterSet = g_timeline_counter_set;
d.storageMode = MTLStorageModeShared;
d.sampleCount = DS4_TIMELINE_SAMPLES_PER_BUFFER;
NSError *error = nil;
id<MTLCounterSampleBuffer> sb = [g_device newCounterSampleBufferWithDescriptor:d error:&error];
if (!sb) {
fprintf(stderr, "ds4: timeline sample buffer failed: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
[b->samples addObject:sb];
}
if (b->count == b->cap) {
b->cap *= 2;
b->recs = realloc(b->recs, (size_t)b->cap * sizeof(ds4_timeline_rec));
}
const uint32_t slot = (b->count % per_buffer) * 2u;
MTLComputePassDescriptor *pd = [MTLComputePassDescriptor computePassDescriptor];
pd.dispatchType = concurrent ? MTLDispatchTypeConcurrent : MTLDispatchTypeSerial;
MTLComputePassSampleBufferAttachmentDescriptor *att = pd.sampleBufferAttachments[0];
att.sampleBuffer = b->samples[sb_index];
att.startOfEncoderSampleIndex = slot;
att.endOfEncoderSampleIndex = slot + 1u;
id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoderWithDescriptor:pd];
if (!enc) return nil;
ds4_timeline_rec *rec = &b->recs[b->count++];
memset(rec, 0, sizeof(*rec));
rec->caller = caller;
ds4_gpu_timeline_hook_encoder(enc);
return enc;
}
/* Blit encoder with timeline sampling (blit passes are otherwise invisible
* to the encoder timeline and show up as gaps). */
static id<MTLBlitCommandEncoder> ds4_gpu_blit_encoder(id<MTLCommandBuffer> cb, const char *label,
uint32_t n_ops) {
DS4TimelineBatch *b = g_timeline_batch;
if (!g_timeline_enabled || !b || cb != g_batch_cb) return [cb blitCommandEncoder];
const uint32_t per_buffer = DS4_TIMELINE_SAMPLES_PER_BUFFER / 2;
const uint32_t sb_index = b->count / per_buffer;
if (sb_index >= [b->samples count]) {
MTLCounterSampleBufferDescriptor *d = [MTLCounterSampleBufferDescriptor new];
d.counterSet = g_timeline_counter_set;
d.storageMode = MTLStorageModeShared;
d.sampleCount = DS4_TIMELINE_SAMPLES_PER_BUFFER;
NSError *error = nil;
id<MTLCounterSampleBuffer> sb = [g_device newCounterSampleBufferWithDescriptor:d error:&error];
if (!sb) return [cb blitCommandEncoder];
[b->samples addObject:sb];
}
if (b->count == b->cap) {
b->cap *= 2;
b->recs = realloc(b->recs, (size_t)b->cap * sizeof(ds4_timeline_rec));
}
const uint32_t slot = (b->count % per_buffer) * 2u;
MTLBlitPassDescriptor *pd = [MTLBlitPassDescriptor blitPassDescriptor];
MTLBlitPassSampleBufferAttachmentDescriptor *att = pd.sampleBufferAttachments[0];
att.sampleBuffer = b->samples[sb_index];
att.startOfEncoderSampleIndex = slot;
att.endOfEncoderSampleIndex = slot + 1u;
id<MTLBlitCommandEncoder> enc = [cb blitCommandEncoderWithDescriptor:pd];
if (!enc) return [cb blitCommandEncoder];
ds4_timeline_rec *rec = &b->recs[b->count++];
memset(rec, 0, sizeof(*rec));
rec->caller = (uintptr_t)__builtin_return_address(0);
rec->n_dispatch = n_ops;
snprintf(rec->kernel, sizeof(rec->kernel), "blit:%s", label ? label : "?");
return enc;
}
static void ds4_gpu_parallel_ffn_reset_state(BOOL close_encoder);
static NSMutableArray<id<MTLCommandBuffer>> *g_pending_cbs;
static id<MTLSharedEvent> g_selected_readback_event;
static uint64_t g_selected_readback_event_value;
static id<MTLComputePipelineState> g_set_rows_f32_i32_pipeline;
static id<MTLComputePipelineState> g_get_rows_f32_pipeline;
static id<MTLComputePipelineState> g_get_rows_f16_pipeline;
static id<MTLComputePipelineState> g_get_rows_i32_pipeline;
static id<MTLComputePipelineState> g_get_rows_q8_0_pipeline;
static id<MTLComputePipelineState> g_get_rows_q4_0_pipeline;
static id<MTLComputePipelineState> g_get_rows_q4_K_pipeline;
static id<MTLComputePipelineState> g_repeat_f32_pipeline;
static id<MTLComputePipelineState> g_concat_pipeline;
static id<MTLComputePipelineState> g_cpy_f32_f32_pipeline;
static id<MTLComputePipelineState> g_cpy_f32_f16_pipeline;
static id<MTLComputePipelineState> g_cpy_contig_f32_f16_pipeline;
static id<MTLComputePipelineState> g_cpy_f16_f32_pipeline;
static id<MTLComputePipelineState> g_cpy_f16_f16_pipeline;
static id<MTLComputePipelineState> g_cpy_contig_f16_f32_pipeline;
static id<MTLComputePipelineState> g_cpy_contig_f16_f16_pipeline;
static id<MTLComputePipelineState> g_flash_kv_stage_f16_pipeline;
static id<MTLComputePipelineState> g_swiglu_pipeline;
static id<MTLComputePipelineState> g_swiglu_flat_pipeline;
static id<MTLComputePipelineState> g_add_pipeline;
static id<MTLComputePipelineState> g_add2_pipeline;
static id<MTLComputePipelineState> g_add3_pipeline;
static id<MTLComputePipelineState> g_moe_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_sum8_pipeline;
static id<MTLComputePipelineState> g_mul_pipeline;
static id<MTLComputePipelineState> g_rms_norm_pipeline;
static id<MTLComputePipelineState> g_rms_norm_plain_pipeline;
static id<MTLComputePipelineState> g_add_rms_norm_pipeline;
static id<MTLComputePipelineState> g_rms_norm_scale_pipeline;
static id<MTLComputePipelineState> g_dsv4_qkv_rms_norm_pipeline;
static id<MTLComputePipelineState> g_dsv4_head_rms_norm_rope_tail_pipeline;
static bool g_use_dsv4_head_rms_norm_rope_tail_pipeline;
static id<MTLComputePipelineState> g_hc_split_sinkhorn_pipeline;
static id<MTLComputePipelineState> g_hc_split_weighted_sum_pipeline;
static id<MTLComputePipelineState> g_hc_split_weighted_sum_norm_pipeline;
static id<MTLComputePipelineState> g_dsv4_hc_producer_pre_norm_pipeline;
static id<MTLComputePipelineState> g_dsv4_hc_expand_producer_pre_norm_pipeline;
static NSMutableDictionary<NSString *, id<MTLBuffer>> *g_dsv4_hc_barrier_cache;
static NSMutableDictionary<NSString *, NSNumber *> *g_dsv4_hc_barrier_gen;
static id<MTLComputePipelineState> g_hc_weighted_sum_pipeline;
static id<MTLComputePipelineState> g_output_hc_weights4_pipeline;
static uint32_t g_test_flags;
static id<MTLComputePipelineState> g_hc_expand_pipeline;
static id<MTLComputePipelineState> g_unary_sigmoid_pipeline;
static id<MTLComputePipelineState> g_unary_silu_pipeline;
static id<MTLComputePipelineState> g_unary_softplus_pipeline;
static id<MTLComputePipelineState> g_unary_sqrt_pipeline;
static id<MTLComputePipelineState> g_unary_clamp_pipeline;
static id<MTLComputePipelineState> g_unary_scale_pipeline;
static id<MTLComputePipelineState> g_unary_fill_pipeline;
static id<MTLComputePipelineState> g_unary_fill_f16_pipeline;
static id<MTLComputePipelineState> g_bin_mul_scalar_pipeline;
static id<MTLComputePipelineState> g_bin_div_row_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_pair_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q2_k_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q2_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_iq2_xxs_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q4_k_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q4_k_pair_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group6_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group8_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group24_q4_k_id_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_group24_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_q2_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_pair_swiglu_tp_static_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_id_mxfp4_sum6_tp_full_rows_static_pipeline_nsg1;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_slots6_mxfp4_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_iq2_xxs_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_q2_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline;
static id<MTLComputePipelineState> g_moe_stream_expert_cache_validate_pipeline;
static id<MTLComputePipelineState> g_moe_q4_gather_slots6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_table_q4_k_sum6_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline;
static id<MTLComputePipelineState> g_moe_mul_mv_addr_q4_k_sum6_pipeline;
static id<MTLArgumentEncoder> g_moe_table_q4_pair_gate_encoder;
static id<MTLArgumentEncoder> g_moe_table_q4_pair_up_encoder;
static id<MTLArgumentEncoder> g_moe_table_q4_sum_down_encoder;
static id<MTLComputePipelineState> g_rope_tail_batch_pipeline;
static id<MTLComputePipelineState> g_rope_tail_inplace_pair_pipeline;
static id<MTLComputePipelineState> g_rope_tail_inplace_pair_shared4_pipeline;
static id<MTLComputePipelineState> g_rope_tail_inplace_pair_affine_pipeline;
static id<MTLComputePipelineState> g_dsv4_fp8_kv_quantize_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexer_qat_pipeline;
static id<MTLComputePipelineState> g_dsv4_kv_fp8_store_pipeline;
static id<MTLComputePipelineState> g_dsv4_kv_rope_fp8_store_pipeline;
static id<MTLComputePipelineState> g_dsv4_ratio4_shift_pipeline;
static id<MTLComputePipelineState> g_dsv4_compressor_pack_ratio4_pipeline;
static id<MTLComputePipelineState> g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline;
static id<MTLComputePipelineState> g_dsv4_compressor_exact_softmax_product_pipeline;
static id<MTLComputePipelineState> g_dsv4_compressor_exact_pool_ratio4_pipeline;
static id<MTLComputePipelineState> g_dsv4_softmax_pool_ratio4_direct_pipeline;
static id<MTLComputePipelineState> g_dsv4_softmax_pool_pipeline;
static id<MTLComputePipelineState> g_soft_max_f32_pipeline;
static id<MTLComputePipelineState> g_soft_max_f32_4_pipeline;
static id<MTLComputePipelineState> g_argsort_f32_i32_desc_pipeline;
static id<MTLComputePipelineState> g_argsort_merge_f32_i32_desc_pipeline;
static id<MTLComputePipelineState> g_sum_rows_f32_f32_pipeline;
static id<MTLComputePipelineState> g_dsv4_topk_mask_pipeline;
static id<MTLComputePipelineState> g_dsv4_topk_mask_scatter_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexer_weighted_sum_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexer_score_one_direct_pipeline;
static id<MTLComputePipelineState> g_dsv4_compressor_store_one_pipeline;
static id<MTLComputePipelineState> g_dsv4_sort_i32_rows_asc_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_rb16_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads16_dual_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_split_pipeline;
static id<MTLComputePipelineState> g_dsv4_indexed_attention_heads8_split_reduce_pipeline;
static id<MTLComputePipelineState> g_dsv4_softplus_sqrt_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_finalize_one_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_finalize_one_simd_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_finalize_weights_one_simd_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_transform_finalize_weights_one_simd_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_select_visual_batch_pipeline;
static NSCache<NSString *, id<MTLBuffer>> *g_dsv4_completion_cache;
static __weak id<MTLBuffer> g_dsv4_hc_producer_last_mix_buffer;
static NSUInteger g_dsv4_hc_producer_last_mix_offset;
static id<MTLBuffer> g_dsv4_hc_producer_last_completion;
static id<MTLComputePipelineState> g_dsv4_router_weights_one_pipeline;
static id<MTLComputePipelineState> g_glm_router_select_one_pipeline;
static id<MTLComputePipelineState> g_glm_kv_lora_rms_norm_pipeline;
static id<MTLComputePipelineState> g_glm_k_b_project_pipeline;
static id<MTLComputePipelineState> g_glm_store_compact_kv_pipeline;
static id<MTLComputePipelineState> g_glm_qkv_norm_store_compact_kv_pipeline;
static id<MTLComputePipelineState> g_glm_store_indexer_k_pipeline;
static id<MTLComputePipelineState> g_glm53_indexer_pool_update_pipeline;
static id<MTLComputePipelineState> g_glm_build_kv_cache_pipeline;
static id<MTLComputePipelineState> g_glm_build_kv_cache_decode_group4_pipeline;
static id<MTLComputePipelineState> g_glm_build_kv_cache_flash_pipeline;
static id<MTLComputePipelineState> g_glm_attention_full_pipeline;
static id<MTLComputePipelineState> g_glm_fill_selected_range_pipeline;
static id<MTLComputePipelineState> g_glm_fill_selected_range_batch_pipeline;
static id<MTLComputePipelineState> g_glm53_expand_pool_selection_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_rope_tail_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_score_one_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_score_one_direct_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_score_one_wide_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_scores_batch_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_scores_tiled_pipeline;
static id<MTLComputePipelineState> g_glm_indexer_scores_tiled_f32_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_glm52_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_glm52_sg_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_batch_pipeline;
static id<MTLComputePipelineState> g_glm_qk_lowrank_batch_glm52_t4_pipeline;
static id<MTLComputePipelineState> g_glm_value_project_q8_0_pipeline;
static id<MTLComputePipelineState> g_glm_value_project_q8_0_batch_heads_pipeline;
static id<MTLComputePipelineState> g_glm_value_project_q8_0_batch_heads_mma_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_split_group8_partial_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_split_group8_reduce_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_decode_split_group8_reduce16_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_group2_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_q2_group4_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_group8_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline;
static id<MTLComputePipelineState> g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu2_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu4_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_pair_swiglu_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_addr_pair_swiglu2_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_addr_pair_swiglu_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q2_k_addr_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q4_k_addr_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q5_k_pair_swiglu_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q5_k_pair_swiglu_mapped_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q5_k_down_f32_pipeline;
static id<MTLComputePipelineState> g_glm_q6_k_down_f32_pipeline;
static id<MTLComputePipelineState> g_dsv4_router_weights_batch_pipeline;
static id<MTLComputePipelineState> g_dsv4_hc_expand4_pipeline;
static NSMutableDictionary<NSString *, id<MTLComputePipelineState>> *g_pipeline_cache;
enum {
DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS = 64,
DS4_METAL_DECODE_PIPELINE_FAST_NAME_BYTES = 96,
};
/*
* Sentinel nxpsg stored by the single-constant ds4_gpu_get_mul_mv_pipeline
* path, which has no second SIMD-group constant. The extended
* ds4_gpu_get_mul_mv_ext_pipeline path derives nxpsg as a positive divisor of
* the 32-row output tile, so INT16_MIN can never collide with a real value.
*/
#define DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE INT16_MIN
_Static_assert(
(DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS &
(DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS - 1u)) == 0,
"decode pipeline fast cache slot count must be a power of two");
typedef struct {
id<MTLComputePipelineState> __strong pipeline;
uint64_t hash;
int16_t nsg;
int16_t nxpsg;
uint16_t name_len;
bool used;
char name[DS4_METAL_DECODE_PIPELINE_FAST_NAME_BYTES];
} ds4_gpu_decode_pipeline_fast_cache_entry;
static ds4_gpu_decode_pipeline_fast_cache_entry
g_decode_pipeline_fast_cache[DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS];
/* Command encoding and the existing NSMutableDictionary pipeline cache are
* backend-serialized; this hot-path mirror intentionally shares that model. */
static bool g_decode_pipeline_fast_lookup_active;
static uint32_t g_decode_pipeline_fast_cache_entries;
static NSMutableDictionary<NSString *, id<MTLBuffer>> *g_model_buffer_cache;
static NSMutableDictionary<NSString *, DS4MetalQ4ExpertTable *> *g_q4_expert_table_cache;
static NSMutableDictionary<NSString *, id> *g_q4_expert_layer_residency_cache;
static NSMutableArray<id<MTLBuffer>> *g_transient_buffers;
static id g_model_residency_set;
typedef struct {
id<MTLBuffer> __strong mask;
id<MTLBuffer> __strong blk;
NSUInteger mask_bytes;
NSUInteger blk_bytes;
uint32_t kind;
uint32_t n_tokens;
uint32_t n_comp;
uint32_t n_keys;
uint32_t window;
uint32_t ratio;
uint32_t nqptg;
uint32_t ncpsg;
bool has_kvpad;
bool bc_mask;
bool valid;
bool blk_ready;
} ds4_gpu_zero_prefix_prefill_mask_cache_entry;
enum {
DS4_GPU_PREFILL_MASK_CACHE_RAW = 1,
DS4_GPU_PREFILL_MASK_CACHE_RATIO4 = 2,
DS4_GPU_PREFILL_MASK_CACHE_RATIO128 = 3,
DS4_GPU_PREFILL_MASK_CACHE_SLOTS = 3,
};
static ds4_gpu_zero_prefix_prefill_mask_cache_entry
g_zero_prefix_prefill_mask_cache[DS4_GPU_PREFILL_MASK_CACHE_SLOTS];
static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void);
static id<MTLBuffer> g_flash_attn_mask_buffer;
static id<MTLBuffer> g_flash_attn_zero_mask_buffer;
static id<MTLBuffer> g_flash_attn_pad_buffer;
static id<MTLBuffer> g_flash_attn_tmp_buffer;
static id<MTLBuffer> g_flash_attn_blk_buffer;
static id<MTLBuffer> g_flash_attn_ring_buffer;
static id<MTLBuffer> g_flash_attn_kv_buffer;
static id<MTLBuffer> g_glm_flash_attn_mask_buffer;
static id<MTLBuffer> g_compressor_pool_kv_buffer;
static id<MTLBuffer> g_compressor_pool_score_buffer;
static id<MTLBuffer> g_compressor_pool_score_cont_buffer;
static id<MTLBuffer> g_compressor_pool_softmax_buffer;
static id<MTLBuffer> g_compressor_pool_product_buffer;
static id<MTLBuffer> g_compressor_store_ape_buffer;
static id<MTLBuffer> g_compressor_store_score_buffer;
static id<MTLBuffer> g_embed_rows_buffer;
static id<MTLBuffer> g_router_selection_buffer;
static id<MTLBuffer> g_router_weight_sum_buffer;
static id<MTLBuffer> g_indexer_head_scores_buffer;
static id<MTLBuffer> g_indexer_topk_buffer;
static id<MTLBuffer> g_topk_select_buffer;
static NSUInteger g_topk_select_bytes;
static uint32_t g_topk_select_parity;
static id<MTLBuffer> g_indexed_topk_buffer;
static id<MTLBuffer> g_f16_round_scratch_buffer;
static id<MTLBuffer> g_q8_prefill_scratch_buffer;
static id<MTLBuffer> g_qwen4_conv_halo_buffer;
static NSUInteger g_qwen4_conv_halo_bytes;
static id<MTLBuffer> g_raw_store_round_buffer;
static id<MTLBuffer> g_moe_gate_scratch_buffer;
static id<MTLBuffer> g_moe_down_scratch_buffer;
static id<MTLBuffer> g_moe_id_map_buffer;
static id<MTLBuffer> g_moe_packed_rhs_buffer;
static NSUInteger g_moe_packed_rhs_bytes;
static id<MTLBuffer> g_moe_q4_gate_slots_buffer;
static id<MTLBuffer> g_moe_q4_up_slots_buffer;
static id<MTLBuffer> g_moe_q4_down_slots_buffer;
static id<MTLBuffer> g_attn_out_group_ids_buffer;
static int g_model_fd = -1;
static const void *g_model_map_ptr;
static uint64_t g_model_map_size;
static uint64_t g_model_mapped_offset;
static uint64_t g_model_mapped_size;
static uint64_t g_model_mapped_max_tensor_bytes;
static uint64_t g_tensor_alloc_live_bytes;
static uint64_t g_tensor_alloc_peak_bytes;
static pthread_mutex_t g_tensor_mu = PTHREAD_MUTEX_INITIALIZER;
static uintptr_t *g_tensor_live_slots;
static size_t g_tensor_live_cap;
static size_t g_tensor_live_count;
static size_t g_tensor_live_tombs;
static uint64_t g_model_wrap_count;
static uint64_t g_model_wrap_bytes;
static uint64_t g_model_wrap_max_bytes;
static uint64_t g_model_buffer_cache_bytes;
static uint64_t g_model_buffer_cache_evictions;
static int g_model_buffer_cache_over_limit;
static uint64_t g_stream_expert_cache_bytes;
static uint64_t g_stream_expert_cache_expert_bytes;
static uint32_t g_stream_expert_cache_entry_count;
static uint32_t g_stream_expert_cache_budget_override;
static uint64_t g_stream_expert_cache_hits;
static uint64_t g_stream_expert_cache_misses;
static uint64_t g_stream_expert_cache_evictions;
static uint64_t g_stream_expert_cache_wraps;
static uint64_t g_stream_expert_cache_clock;
static uint64_t g_stream_expert_cache_evict_advise_bytes;
static uint64_t g_stream_expert_cache_willneed_advise_bytes;
static uint64_t g_stream_expert_cache_pread_bytes;
static double g_stream_expert_cache_pread_ms;
static uint64_t g_stream_expert_cache_buffer_allocs;
static uint64_t g_stream_expert_cache_buffer_reuses;
static uint64_t g_stream_expert_cache_decode_tokens;
static uint32_t g_stream_expert_cache_first_decode_layer = UINT32_MAX;
static uint64_t g_stream_expert_cache_hotness_decay_token;
static uint64_t g_stream_expert_timing_selected_calls;
static double g_stream_expert_timing_selected_read_ms;
static double g_stream_expert_timing_selected_sync_ms;
static double g_stream_expert_timing_selected_copy_ms;
static double g_stream_expert_timing_selected_bind_ms;
static uint64_t g_stream_expert_timing_split_layers;
static uint64_t g_stream_expert_timing_split_resident_experts;
static uint64_t g_stream_expert_timing_split_missing_experts;
static double g_stream_expert_timing_split_resident_ms;
static double g_stream_expert_timing_split_missing_ms;
static double g_stream_expert_timing_split_missing_load_ms;
static double g_stream_expert_timing_split_missing_slot_ms;
static double g_stream_expert_timing_split_missing_prune_ms;
static double g_stream_expert_timing_split_missing_addr_ms;
static double g_stream_expert_timing_split_missing_wait_ms;
static uint64_t g_stream_expert_timing_load_calls;
static double g_stream_expert_timing_load_prepare_ms;
static double g_stream_expert_timing_load_pread_ms;
static double g_stream_expert_timing_load_install_ms;
static uint64_t g_stream_expert_timing_prepare_batch_reuse_calls;
static double g_stream_expert_timing_prepare_batch_reuse_ms;
static uint64_t g_stream_expert_timing_prepare_buffer_calls;
static double g_stream_expert_timing_prepare_buffer_ms;
static uint64_t g_stream_expert_timing_prepare_task_experts;
static double g_stream_expert_timing_prepare_task_ms;
static uint64_t g_stream_expert_timing_reuse_scan_calls;
static uint64_t g_stream_expert_timing_reuse_scan_entries;
static double g_stream_expert_timing_reuse_scan_ms;
static double g_stream_expert_timing_reuse_clear_ms;
static uint64_t g_stream_expert_timing_readahead_calls;
static uint64_t g_stream_expert_timing_readahead_bytes;
static double g_stream_expert_timing_readahead_ms;
static uint64_t g_stream_expert_timing_cache_all_resident_layers;
static uint64_t g_stream_expert_timing_cache_all_missing_layers;
static uint64_t g_stream_expert_timing_cache_mixed_layers;
static uint64_t g_stream_expert_timing_cache_resident_experts;
static uint64_t g_stream_expert_timing_cache_missing_experts;
typedef struct {
uint64_t selected_calls;
double selected_read_ms;
double selected_sync_ms;
double selected_copy_ms;
double selected_bind_ms;
uint64_t split_layers;
uint64_t split_resident_experts;
uint64_t split_missing_experts;
double split_resident_ms;
double split_missing_ms;
double split_missing_load_ms;
double split_missing_slot_ms;
double split_missing_prune_ms;
double split_missing_addr_ms;
double split_missing_wait_ms;
uint64_t load_calls;
double load_prepare_ms;
double load_pread_ms;
double load_install_ms;
uint64_t prepare_batch_reuse_calls;
double prepare_batch_reuse_ms;
uint64_t prepare_buffer_calls;
double prepare_buffer_ms;
uint64_t prepare_task_experts;
double prepare_task_ms;
uint64_t reuse_scan_calls;
uint64_t reuse_scan_entries;
double reuse_scan_ms;
double reuse_clear_ms;
uint64_t readahead_calls;
uint64_t readahead_bytes;
double readahead_ms;
uint64_t cache_all_resident_layers;
uint64_t cache_all_missing_layers;
uint64_t cache_mixed_layers;
uint64_t cache_resident_experts;
uint64_t cache_missing_experts;
} ds4_gpu_stream_expert_timing_snapshot;
static ds4_gpu_stream_expert_timing_snapshot g_stream_expert_timing_last_report;
static int g_stream_prefill_batch_selected_addr_building;
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_ssd_streaming_mode;
static int g_glm_streaming_prefill_full_layer_runtime;
static int g_metal4_runtime_available;
static int g_metal4_family_supported;
static int g_metal4_queue_supported;
static int g_metal4_m5_neural_accelerators_hint;
static int g_metal4_tensor_api_enabled;
static int g_metal4_tensor_api_compile_supported;
static char g_metal_device_name[128];
static int ds4_gpu_model_map_log_enabled(void);
static int ds4_gpu_stream_expert_cache_note_expert_size(
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes);
static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void);
static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats);
static void ds4_gpu_stream_expert_pending_load_clear(void);
static void ds4_gpu_stream_expert_pread_pool_shutdown(void);
static int ds4_gpu_stream_expert_timing_summary_enabled(void);
static int ds4_gpu_stream_expert_cache_entry_protected(
uint32_t layer,
uint32_t expert,
uint32_t protect_layer,
const int32_t *protect_ids,
uint32_t n_protect);
/* The async selected-load worker registers itself so cache paths that would
* flush/wait on command buffers (a race against the encoding thread) fail
* the load instead; the caller then retries on the main thread. */
static pthread_t g_stream_expert_service_thread;
static int g_stream_expert_service_thread_set;
void ds4_gpu_stream_expert_cache_note_service_thread(void) {
g_stream_expert_service_thread = pthread_self();
g_stream_expert_service_thread_set = 1;
}
static int ds4_gpu_stream_expert_cache_on_service_thread(void) {
return g_stream_expert_service_thread_set &&
pthread_equal(pthread_self(), g_stream_expert_service_thread);
}
static NSUInteger g_flash_attn_mask_bytes;
static NSUInteger g_flash_attn_zero_mask_bytes;
static NSUInteger g_flash_attn_pad_bytes;
static NSUInteger g_flash_attn_tmp_bytes;
static NSUInteger g_flash_attn_blk_bytes;
static NSUInteger g_flash_attn_ring_bytes;
static NSUInteger g_flash_attn_kv_bytes;
static NSUInteger g_glm_flash_attn_mask_bytes;
static uint32_t g_glm_flash_attn_mask_pos0;
static uint32_t g_glm_flash_attn_mask_tokens;
static uint32_t g_glm_flash_attn_mask_cache_len;
static int g_glm_flash_attn_mask_valid;
static NSUInteger g_compressor_pool_kv_bytes;
static NSUInteger g_compressor_pool_score_bytes;
static NSUInteger g_compressor_pool_score_cont_bytes;
static NSUInteger g_compressor_pool_softmax_bytes;
static NSUInteger g_compressor_pool_product_bytes;
static NSUInteger g_compressor_store_ape_bytes;
static NSUInteger g_compressor_store_score_bytes;
static NSUInteger g_embed_rows_bytes;
static NSUInteger g_router_selection_bytes;
static NSUInteger g_router_weight_sum_bytes;
static NSUInteger g_indexer_head_scores_bytes;
static NSUInteger g_indexer_topk_bytes;
static NSUInteger g_indexed_topk_bytes;
static NSUInteger g_f16_round_scratch_bytes;
static NSUInteger g_q8_prefill_scratch_bytes;
static NSUInteger g_raw_store_round_bytes;
static NSUInteger g_moe_gate_scratch_bytes;
static NSUInteger g_moe_down_scratch_bytes;
static NSUInteger g_moe_id_map_bytes;
static NSUInteger g_moe_q4_gate_slots_bytes;
static NSUInteger g_moe_q4_up_slots_bytes;
static NSUInteger g_moe_q4_down_slots_bytes;
static NSUInteger g_attn_out_group_ids_bytes;
static int g_initialized;
static int g_quality_mode;
static int g_mpp_invalid_env_reported;
#define DS4_METAL_MAX_ROUTED_EXPERT_USED 8
static int32_t g_routed_moe_selected_override[DS4_METAL_MAX_ROUTED_EXPERT_USED];
static uint32_t g_routed_moe_selected_override_n;
static int g_moe_selected_trace_record_initialized;
static FILE *g_moe_selected_trace_record_fp;
static uint64_t g_moe_selected_trace_record_count;
static int g_moe_selected_trace_replay_initialized;
static int32_t *g_moe_selected_trace_replay_ids;
static uint64_t g_moe_selected_trace_replay_count;
static uint64_t g_moe_selected_trace_replay_pos;
static double ds4_gpu_gib(uint64_t bytes);
static uint64_t ds4_gpu_system_memory_bytes(void) {
uint64_t bytes = 0;
size_t len = sizeof(bytes);
if (sysctlbyname("hw.memsize", &bytes, &len, NULL, 0) != 0) return 0;
return len == sizeof(bytes) ? bytes : 0;
}
static void ds4_gpu_print_device_summary(void) {
const char *name = g_device.name ? [g_device.name UTF8String] : "unknown Metal device";
uint64_t mem = ds4_gpu_system_memory_bytes();
if (mem) {
double gib = (double)mem / 1024.0 / 1024.0 / 1024.0;
fprintf(stderr, "ds4: Metal device %s, %.2f GiB RAM\n", name, gib);
} else {
fprintf(stderr, "ds4: Metal device %s\n", name);
}
}
#define DS4_METAL_MAX_MODEL_VIEWS 4096
/* Compatibility fallback for callers that cannot provide a parsed GGUF tensor
* span. The normal DS4 engine passes the exact maximum tensor byte size. */
#define DS4_METAL_FALLBACK_MAX_TENSOR_BYTES (4ull * 1024ull * 1024ull * 1024ull)
typedef struct {
__strong id<MTLBuffer> buffer;
const void *model_map;
uint64_t model_size;
uint64_t model_offset;
uint64_t bytes;
} ds4_gpu_model_view;
static ds4_gpu_model_view g_model_views[DS4_METAL_MAX_MODEL_VIEWS];
static uint32_t g_model_view_count;
enum {
DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER = 80,
DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT = 384,
DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED = DS4_METAL_MAX_ROUTED_EXPERT_USED,
DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES =
DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER *
DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT,
DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS = 256,
DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS = 16,
DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS = 16,
};
typedef struct {
uint32_t layer;
uint32_t expert;
uint64_t hits;
} ds4_gpu_moe_selected_hotlist_entry;
static int g_moe_selected_hotlist_initialized;
static uint64_t
g_moe_selected_hotlist_counts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
static uint64_t g_moe_selected_hotlist_records;
static uint64_t g_moe_selected_hotlist_selections;
typedef struct {
__strong id<MTLBuffer> gate_buffer;
__strong id<MTLBuffer> up_buffer;
__strong id<MTLBuffer> down_buffer;
const void *model_map;
uint64_t model_size;
uint64_t gate_abs_offset;
uint64_t up_abs_offset;
uint64_t down_abs_offset;
uint64_t gate_expert_bytes;
uint64_t down_expert_bytes;
uint64_t logical_bytes;
uint64_t last_used;
uint64_t use_count;
NSUInteger gate_inner;
NSUInteger up_inner;
NSUInteger down_inner;
uint64_t inflight_seq;
uint32_t slab_slot;
uint8_t valid;
uint8_t slab_backed;
} ds4_gpu_stream_expert_cache_entry;
typedef struct {
__strong id<MTLBuffer> gate_buffer;
__strong id<MTLBuffer> up_buffer;
__strong id<MTLBuffer> down_buffer;
NSUInteger gate_inner;
NSUInteger up_inner;
NSUInteger down_inner;
} ds4_gpu_stream_expert_reusable_buffers;
static ds4_gpu_stream_expert_cache_entry
g_stream_expert_cache[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
static ds4_gpu_stream_expert_cache_entry
g_stream_full_expert_addr_entry[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint32_t g_stream_expert_cache_layer_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static double g_stream_expert_cache_layer_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_last_hits[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_last_misses[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_last_evictions[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint64_t g_stream_expert_cache_layer_last_pread_bytes[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static double g_stream_expert_cache_layer_last_pread_ms[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static uint32_t
g_stream_expert_cache_route_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER][DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
static id<MTLBuffer> g_stream_expert_cache_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_expert_cache_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_expert_cache_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_expert_cache_slabs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS];
static uint32_t g_stream_expert_cache_slab_start_slot[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS];
static uint32_t g_stream_expert_cache_slab_slot_count[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS];
static uint32_t g_stream_expert_cache_slab_slots_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS];
static uint32_t g_stream_expert_cache_slab_count;
static uint32_t g_stream_expert_cache_slab_total_slots;
static uint32_t g_stream_expert_cache_free_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES];
static uint32_t g_stream_expert_cache_free_slot_count;
static uint8_t g_stream_expert_cache_slab_slot_locked[DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES];
static uint32_t g_stream_expert_cache_mlock_budget_cap;
static uint8_t g_stream_expert_cache_mlock_relief_applied;
static uint64_t g_stream_expert_cache_mlock_bytes;
static uint64_t g_stream_expert_cache_mlock_fail_bytes;
static uint64_t g_stream_expert_cache_mlock_failures;
static double g_stream_expert_cache_mlock_ms;
static int g_stream_expert_cache_mlock_warned;
static uint64_t g_stream_expert_cache_slab_slot_bytes;
static uint64_t g_stream_expert_cache_cb_seq;
static uint64_t g_stream_expert_cache_done_seq;
static uint64_t g_stream_expert_cache_batch_seq;
static uint64_t g_stream_expert_cache_owned_seq;
static uint64_t g_stream_expert_cache_pending_max_seq;
static id<MTLBuffer> g_stream_compact_gate_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_compact_up_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_compact_down_addr_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_compact_selected_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_selected_id_buffers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER];
static id<MTLBuffer> g_stream_expert_validate_status_buffer;
@interface DS4MetalTensor : NSObject
@property(nonatomic, strong) id<MTLBuffer> buffer;
@property(nonatomic, assign) uint64_t offset;
@property(nonatomic, assign) uint64_t bytes;
@property(nonatomic, assign) uint8_t owner;
@end
@implementation DS4MetalTensor
@end
@interface DS4MetalQ4ExpertTable : NSObject
@property(nonatomic, strong) id<MTLBuffer> argumentBuffer;
@property(nonatomic, strong) id<MTLBuffer> addressBuffer;
@property(nonatomic, strong) NSMutableArray<id<MTLBuffer>> *expertBuffers;
@property(nonatomic, strong) id residencySet;
@property(nonatomic, assign) BOOL residencySetAddedToQueue;
@property(nonatomic, assign) uint32_t nExpert;
@property(nonatomic, assign) uint64_t expertBytes;
@end
@implementation DS4MetalQ4ExpertTable
- (void)dealloc {
#if TARGET_OS_OSX
if (@available(macOS 15.0, *)) {
if (_residencySet) {
if (_residencySetAddedToQueue &&
g_queue &&
[g_queue respondsToSelector:@selector(removeResidencySet:)]) {
[g_queue removeResidencySet:_residencySet];
}
[_residencySet endResidency];
}
}
#endif
}
@end
@interface DS4MetalQ4LayerResidency : NSObject
@property(nonatomic, strong) id residencySet;
@property(nonatomic, assign) BOOL addedToQueue;
@end
@implementation DS4MetalQ4LayerResidency
- (void)dealloc {
#if TARGET_OS_OSX
if (@available(macOS 15.0, *)) {
if (_residencySet) {
if (_addedToQueue &&
g_queue &&
[g_queue respondsToSelector:@selector(removeResidencySet:)]) {
[g_queue removeResidencySet:_residencySet];
}
[_residencySet endResidency];
}
}
#endif
}
@end
static DS4MetalTensor *ds4_gpu_tensor_obj(ds4_gpu_tensor *tensor) {
return (__bridge DS4MetalTensor *)tensor;
}
static const DS4MetalTensor *ds4_gpu_tensor_const_obj(const ds4_gpu_tensor *tensor) {
return (__bridge const DS4MetalTensor *)tensor;
}
/* C code owns ds4_gpu_tensor handles as retained Objective-C objects. Freeing
* the same opaque handle twice would make the second __bridge_transfer release
* an already-deallocated object, which macOS reports as malloc corruption. The
* live table lets free validate a handle before touching Objective-C state; the
* same mutex also serializes the diagnostic allocation counters. */
static uint64_t ds4_gpu_tensor_ptr_hash(uintptr_t ptr) {
uint64_t x = (uint64_t)(ptr >> 4);
x ^= x >> 33;
x *= UINT64_C(0xff51afd7ed558ccd);
x ^= x >> 33;
x *= UINT64_C(0xc4ceb9fe1a85ec53);
x ^= x >> 33;
return x;
}
static int ds4_gpu_tensor_live_resize_locked(size_t min_cap) {
size_t new_cap = 1024;
while (new_cap < min_cap) new_cap <<= 1;
uintptr_t *new_slots = calloc(new_cap, sizeof(new_slots[0]));
if (!new_slots) return 0;
for (size_t i = 0; i < g_tensor_live_cap; i++) {
const uintptr_t key = g_tensor_live_slots[i];
if (key == 0 || key == UINTPTR_MAX) continue;
size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (new_cap - 1);
while (new_slots[idx] != 0) idx = (idx + 1) & (new_cap - 1);
new_slots[idx] = key;
}
free(g_tensor_live_slots);
g_tensor_live_slots = new_slots;
g_tensor_live_cap = new_cap;
g_tensor_live_tombs = 0;
return 1;
}
static int ds4_gpu_tensor_live_insert_locked(const void *ptr) {
if (!ptr || (uintptr_t)ptr == UINTPTR_MAX) return 0;
if ((g_tensor_live_count + g_tensor_live_tombs + 1) * 10 >=
g_tensor_live_cap * 7)
{
const size_t min_cap = g_tensor_live_cap ? g_tensor_live_cap * 2 : 1024;
if (!ds4_gpu_tensor_live_resize_locked(min_cap)) return 0;
}
const uintptr_t key = (uintptr_t)ptr;
size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1);
size_t tomb = (size_t)-1;
for (;;) {
const uintptr_t cur = g_tensor_live_slots[idx];
if (cur == key) return 0;
if (cur == UINTPTR_MAX) {
if (tomb == (size_t)-1) tomb = idx;
} else if (cur == 0) {
if (tomb != (size_t)-1) {
idx = tomb;
g_tensor_live_tombs--;
}
g_tensor_live_slots[idx] = key;
g_tensor_live_count++;
return 1;
}
idx = (idx + 1) & (g_tensor_live_cap - 1);
}
}
static int ds4_gpu_tensor_live_remove_locked(const void *ptr) {
if (!ptr || g_tensor_live_cap == 0) return 0;
const uintptr_t key = (uintptr_t)ptr;
size_t idx = (size_t)ds4_gpu_tensor_ptr_hash(key) & (g_tensor_live_cap - 1);
for (;;) {
const uintptr_t cur = g_tensor_live_slots[idx];
if (cur == 0) return 0;
if (cur == key) {
g_tensor_live_slots[idx] = UINTPTR_MAX;
g_tensor_live_count--;
g_tensor_live_tombs++;
return 1;
}
idx = (idx + 1) & (g_tensor_live_cap - 1);
}
}
static int ds4_gpu_tensor_track_alloc_locked(
const void *ptr,
uint64_t bytes,
uint64_t *live_snap,
uint64_t *peak_snap)
{
if (!ds4_gpu_tensor_live_insert_locked(ptr)) return 0;
g_tensor_alloc_live_bytes += bytes;
if (g_tensor_alloc_live_bytes > g_tensor_alloc_peak_bytes) {
g_tensor_alloc_peak_bytes = g_tensor_alloc_live_bytes;
}
if (live_snap) *live_snap = g_tensor_alloc_live_bytes;
if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes;
return 1;
}
static int ds4_gpu_tensor_track_view_locked(const void *ptr) {
return ds4_gpu_tensor_live_insert_locked(ptr);
}
static int ds4_gpu_tensor_prepare_free(
ds4_gpu_tensor *tensor,
uint8_t *owner,
uint64_t *bytes,
uint64_t *live_snap,
uint64_t *peak_snap)
{
pthread_mutex_lock(&g_tensor_mu);
if (!ds4_gpu_tensor_live_remove_locked(tensor)) {
pthread_mutex_unlock(&g_tensor_mu);
fprintf(stderr,
"ds4: Metal tensor free ignored for unknown handle %p\n",
(void *)tensor);
return 0;
}
DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor);
const uint8_t obj_owner = obj.owner;
const uint64_t obj_bytes = obj.bytes;
if (obj_owner) {
if (obj_bytes <= g_tensor_alloc_live_bytes) {
g_tensor_alloc_live_bytes -= obj_bytes;
} else {
g_tensor_alloc_live_bytes = 0;
}
}
if (owner) *owner = obj_owner;
if (bytes) *bytes = obj_bytes;
if (live_snap) *live_snap = g_tensor_alloc_live_bytes;
if (peak_snap) *peak_snap = g_tensor_alloc_peak_bytes;
pthread_mutex_unlock(&g_tensor_mu);
return 1;
}
static void ds4_gpu_tensor_tracking_reset(void) {
pthread_mutex_lock(&g_tensor_mu);
if (g_tensor_live_count != 0) {
fprintf(stderr,
"ds4: Metal cleanup discarded %zu live tensor handles\n",
g_tensor_live_count);
}
free(g_tensor_live_slots);
g_tensor_live_slots = NULL;
g_tensor_live_cap = 0;
g_tensor_live_count = 0;
g_tensor_live_tombs = 0;
g_tensor_alloc_live_bytes = 0;
g_tensor_alloc_peak_bytes = 0;
pthread_mutex_unlock(&g_tensor_mu);
}
static id<MTLBuffer> ds4_gpu_tensor_buffer(const ds4_gpu_tensor *tensor) {
if (!tensor) return nil;
const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor);
return obj.buffer;
}
static NSUInteger ds4_gpu_tensor_offset(const ds4_gpu_tensor *tensor) {
if (!tensor) return 0;
const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor);
return (NSUInteger)obj.offset;
}
static id<MTLCommandBuffer> ds4_gpu_new_command_buffer(void);
static void ds4_gpu_stream_expert_cache_note_owned_created(void);
static id<MTLCommandBuffer> ds4_gpu_command_buffer(int *owned) {
if (g_batch_cb) {
*owned = 0;
return g_batch_cb;
}
*owned = 1;
id<MTLCommandBuffer> cb = ds4_gpu_new_command_buffer();
if (cb) ds4_gpu_stream_expert_cache_note_owned_created();
return cb;
}
static id<MTLComputeCommandEncoder> ds4_gpu_compute_encoder(id<MTLCommandBuffer> cb) {
if (g_batch_cb && cb == g_batch_cb) {
g_batch_has_work = YES;
if (g_timeline_enabled && g_timeline_batch) {
/* Diagnostic: one timestamped pass per dispatch group. A
* concurrent section keeps its single encoder. */
if (g_batch_enc && !g_batch_encoder_concurrent) {
[g_batch_enc endEncoding];
g_batch_enc = nil;
}
if (!g_batch_enc) {
g_batch_enc = ds4_gpu_timeline_new_encoder(
cb, g_batch_encoder_concurrent,
(uintptr_t)__builtin_return_address(0));
}
if (g_batch_enc) return g_batch_enc;
}
if (!g_batch_enc) {
g_batch_enc = g_batch_encoder_concurrent
? [cb computeCommandEncoderWithDispatchType:MTLDispatchTypeConcurrent]
: [cb computeCommandEncoder];
}
return g_batch_enc;
}
return [cb computeCommandEncoder];
}
static void ds4_gpu_end_compute_encoder(id<MTLCommandBuffer> cb, id<MTLComputeCommandEncoder> enc) {
if (!enc) return;
if (g_batch_cb && cb == g_batch_cb && enc == g_batch_enc) return;
[enc endEncoding];
}
static void ds4_gpu_close_batch_encoder(void) {
if (!g_batch_enc) return;
[g_batch_enc endEncoding];
g_batch_enc = nil;
}
static double g_gpu_busy_accum;
static uint64_t g_gpu_busy_cbs;
/* A failed command buffer can leave a cross-threadgroup arrival counter at an
* arbitrary partial value. Drop cached ownership instead of CPU-resetting
* buffers that another in-flight command buffer might still reference; bound
* resources remain retained by their command buffers/transient list. */
static void ds4_gpu_invalidate_completion_counters(void) {
[g_dsv4_completion_cache removeAllObjects];
g_dsv4_hc_producer_last_mix_buffer = nil;
g_dsv4_hc_producer_last_mix_offset = 0;
g_dsv4_hc_producer_last_completion = nil;
}
static int ds4_gpu_wait_command_buffer(id<MTLCommandBuffer> cb, const char *label) {
[cb waitUntilCompleted];
if (getenv("DS4_METAL_CB_TIMES")) {
static double prev_gpu_end;
static uint64_t n_printed;
if (n_printed < 400) {
n_printed++;
fprintf(stderr,
"ds4: cb %s: driver %.0f us, queue-wait %.0f us, gpu %.0f us, gap-from-prev-gpu-end %.0f us\n",
label,
(cb.kernelEndTime - cb.kernelStartTime) * 1e6,
(cb.GPUStartTime - cb.kernelEndTime) * 1e6,
(cb.GPUEndTime - cb.GPUStartTime) * 1e6,
prev_gpu_end ? (cb.GPUStartTime - prev_gpu_end) * 1e6 : 0.0);
}
prev_gpu_end = cb.GPUEndTime;
}
if (getenv("DS4_METAL_GPU_BUSY_PROFILE")) {
const double busy = cb.GPUEndTime - cb.GPUStartTime;
if (busy > 0) g_gpu_busy_accum += busy;
if ((++g_gpu_busy_cbs % 64u) == 0u) {
fprintf(stderr, "ds4: gpu busy accum %.1f ms over %llu cbs\n",
g_gpu_busy_accum * 1000.0,
(unsigned long long)g_gpu_busy_cbs);
}
}
if (cb.status == MTLCommandBufferStatusError) {
fprintf(stderr, "ds4: Metal %s failed: %s\n",
label, [[cb.error localizedDescription] UTF8String]);
ds4_gpu_invalidate_completion_counters();
return 0;
}
return 1;
}
static id<MTLCommandBuffer> ds4_gpu_new_command_buffer(void) {
static int initialized;
static int use_unretained;
if (!initialized) {
use_unretained = getenv("DS4_METAL_UNRETAINED_COMMAND_BUFFERS") != NULL;
initialized = 1;
}
if (use_unretained) {
return [g_queue commandBufferWithUnretainedReferences];
}
return [g_queue commandBuffer];
}
static uint64_t ds4_gpu_exact_view_cache_limit_bytes(void) {
static int initialized;
static uint64_t limit_bytes;
if (initialized) return limit_bytes;
const uint64_t mib = 1024ull * 1024ull;
const uint64_t gib = 1024ull * mib;
limit_bytes = 64ull * gib;
const char *gib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_GIB");
if (gib_env && gib_env[0]) {
char *end = NULL;
unsigned long long v = strtoull(gib_env, &end, 10);
if (end != gib_env && *end == '\0') {
limit_bytes = v > UINT64_MAX / gib ? UINT64_MAX : (uint64_t)v * gib;
}
}
const char *mib_env = getenv("DS4_METAL_EXACT_VIEW_CACHE_MIB");
if (mib_env && mib_env[0]) {
char *end = NULL;
unsigned long long v = strtoull(mib_env, &end, 10);
if (end != mib_env && *end == '\0') {
limit_bytes = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib;
}
}
initialized = 1;
return limit_bytes;
}
static void ds4_gpu_model_buffer_cache_note_insert(uint64_t bytes) {
if (g_model_buffer_cache_bytes > UINT64_MAX - bytes) {
g_model_buffer_cache_bytes = UINT64_MAX;
} else {
g_model_buffer_cache_bytes += bytes;
}
const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes();
if (limit != 0 && g_model_buffer_cache_bytes > limit) {
g_model_buffer_cache_over_limit = 1;
}
}
static void ds4_gpu_model_buffer_cache_clear(const char *reason) {
if (!g_model_buffer_cache) {
g_model_buffer_cache_bytes = 0;
g_model_buffer_cache_over_limit = 0;
return;
}
const NSUInteger entries = [g_model_buffer_cache count];
if (entries != 0) {
if (getenv("DS4_METAL_EXACT_VIEW_CACHE_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal exact model view cache evict reason=%s entries=%lu bytes=%.2f GiB limit=%.2f GiB\n",
reason ? reason : "unknown",
(unsigned long)entries,
ds4_gpu_gib(g_model_buffer_cache_bytes),
ds4_gpu_gib(ds4_gpu_exact_view_cache_limit_bytes()));
}
[g_model_buffer_cache removeAllObjects];
g_model_buffer_cache_evictions++;
}
g_model_buffer_cache_bytes = 0;
g_model_buffer_cache_over_limit = 0;
}
static void ds4_gpu_model_buffer_cache_maybe_evict(const char *reason) {
if (g_model_buffer_cache_over_limit) {
ds4_gpu_model_buffer_cache_clear(reason);
}
}
static uint64_t ds4_gpu_stream_expert_cache_next_cb_seq(void) {
if (g_stream_expert_cache_cb_seq == UINT64_MAX) {
/*
* A real wrap would require an astronomical number of command buffers.
* Resetting the epoch space is still safer than letting zero become a
* valid in-flight marker.
*/
g_stream_expert_cache_cb_seq = 0;
g_stream_expert_cache_done_seq = 0;
g_stream_expert_cache_batch_seq = 0;
g_stream_expert_cache_owned_seq = 0;
g_stream_expert_cache_pending_max_seq = 0;
for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) {
for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) {
g_stream_expert_cache[layer][expert].inflight_seq = 0;
}
}
}
return ++g_stream_expert_cache_cb_seq;
}
static void ds4_gpu_stream_expert_cache_note_batch_created(void) {
g_stream_expert_cache_batch_seq =
ds4_gpu_stream_expert_cache_next_cb_seq();
}
static void ds4_gpu_stream_expert_cache_note_batch_committed(void) {
if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_pending_max_seq) {
g_stream_expert_cache_pending_max_seq = g_stream_expert_cache_batch_seq;
}
g_stream_expert_cache_batch_seq = 0;
}
static void ds4_gpu_stream_expert_cache_note_owned_created(void) {
g_stream_expert_cache_owned_seq =
ds4_gpu_stream_expert_cache_next_cb_seq();
}
static void ds4_gpu_stream_expert_cache_note_pending_completed(void) {
if (g_stream_expert_cache_pending_max_seq > g_stream_expert_cache_done_seq) {
g_stream_expert_cache_done_seq = g_stream_expert_cache_pending_max_seq;
}
g_stream_expert_cache_pending_max_seq = 0;
}
static void ds4_gpu_stream_expert_cache_note_owned_completed(void) {
if (g_stream_expert_cache_owned_seq > g_stream_expert_cache_done_seq) {
g_stream_expert_cache_done_seq = g_stream_expert_cache_owned_seq;
}
g_stream_expert_cache_owned_seq = 0;
}
static int ds4_gpu_stream_expert_cache_entry_inflight(
const ds4_gpu_stream_expert_cache_entry *e) {
return e && e->valid && e->inflight_seq > g_stream_expert_cache_done_seq;
}
static int ds4_gpu_stream_expert_cache_mark_inflight(
ds4_gpu_stream_expert_cache_entry *e) {
if (!e || !e->valid) return 0;
const uint64_t seq = g_stream_expert_cache_batch_seq ?
g_stream_expert_cache_batch_seq :
g_stream_expert_cache_owned_seq;
if (seq == 0) return 0;
e->inflight_seq = seq;
return 1;
}
static int ds4_gpu_stream_expert_cache_mark_entries_inflight(
ds4_gpu_stream_expert_cache_entry * const *entries,
uint32_t n_entries,
uint32_t active_mask) {
if (!entries || n_entries == 0) return 0;
for (uint32_t i = 0; i < n_entries; i++) {
if (active_mask != 0 && (active_mask & (1u << i)) == 0) continue;
if (!ds4_gpu_stream_expert_cache_mark_inflight(entries[i])) return 0;
}
return 1;
}
static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label);
static int ds4_gpu_wait_pending_command_buffers(const char *label) {
int ok = 1;
for (id<MTLCommandBuffer> pending in g_pending_cbs) {
if (!ds4_gpu_wait_command_buffer(pending, label)) ok = 0;
}
[g_pending_cbs removeAllObjects];
ds4_gpu_stream_expert_cache_note_pending_completed();
if (!ok) ds4_gpu_invalidate_zero_prefix_prefill_block_maps();
return ok;
}
static double g_batch_cb_created_ms;
static double ds4_gpu_now_ms(void);
static void ds4_gpu_queue_keepalive_start(void);
static void ds4_gpu_queue_keepalive_stop_thread(void);
static int ds4_gpu_finish_command_buffer(id<MTLCommandBuffer> cb, int owned, const char *label) {
if (!owned) return 1;
const double t_commit = ds4_gpu_now_ms();
[cb commit];
int ok = ds4_gpu_wait_pending_command_buffers(label);
if (!ds4_gpu_wait_command_buffer(cb, label)) {
ok = 0;
ds4_gpu_invalidate_zero_prefix_prefill_block_maps();
}
if (getenv("DS4_METAL_CB_TIMES")) {
const double t_done = ds4_gpu_now_ms();
struct timespec ts_mono;
clock_gettime(CLOCK_MONOTONIC, &ts_mono);
const double mono_ms = (double)ts_mono.tv_sec * 1e3 + (double)ts_mono.tv_nsec / 1e6;
fprintf(stderr, "ds4: cb '%s': mono %.1f ms: ", label ? label : "?", mono_ms);
fprintf(stderr, "encode %.1f ms, commit->done %.1f ms, gpu span %.1f ms, gpu start +%.1f ms, caller 0x%llx\n",
t_commit - g_batch_cb_created_ms, t_done - t_commit,
(cb.GPUEndTime - cb.GPUStartTime) * 1e3,
cb.GPUStartTime * 1e3 - (mono_ms - (t_done - t_commit)),
(unsigned long long)((uintptr_t)__builtin_return_address(0) - (uintptr_t)_dyld_get_image_vmaddr_slide(0)));
if (t_done - t_commit > 100.0 && (cb.GPUEndTime - cb.GPUStartTime) * 1e3 < 10.0) {
void *frames[12];
const int n = backtrace(frames, 12);
fprintf(stderr, "ds4: cb stall backtrace:");
for (int i = 1; i < n; i++)
fprintf(stderr, " 0x%llx", (unsigned long long)((uintptr_t)frames[i] - (uintptr_t)_dyld_get_image_vmaddr_slide(0)));
fprintf(stderr, "\n");
}
}
ds4_gpu_stream_expert_cache_note_owned_completed();
if (getenv("DS4_METAL_CB_TIMES")) {
uint64_t bytes = 0;
for (id<MTLBuffer> b in g_transient_buffers) bytes += [b length];
fprintf(stderr, "ds4: transient buffers this batch: %lu (%.1f MiB)\n",
(unsigned long)[g_transient_buffers count], bytes / 1048576.0);
}
[g_transient_buffers removeAllObjects];
ds4_gpu_model_buffer_cache_maybe_evict(label);
return ok;
}
static int ds4_gpu_device_name_contains(const char *needle);
static int ds4_gpu_use_m5_private_scratch(void) {
static int initialized;
static int enabled;
if (!initialized) {
enabled = ds4_gpu_device_name_contains("M5");
initialized = 1;
}
return enabled;
}
static int ds4_gpu_scratch_needs_cpu_access(const char *label) {
if (!label) return 0;
return strstr(label, "mask") != NULL ||
strcmp(label, "ds4_attention_output_group_ids") == 0;
}
static MTLResourceOptions ds4_gpu_model_resource_options(void) {
MTLResourceOptions options = MTLResourceStorageModeShared;
if (getenv("DS4_METAL_MODEL_UNTRACKED") != NULL) {
options |= MTLResourceHazardTrackingModeUntracked;
}
return options;
}
static int ds4_gpu_ensure_scratch_buffer(
id<MTLBuffer> __strong *buffer,
NSUInteger *capacity,
NSUInteger bytes,
const char *label) {
if (*buffer && *capacity >= bytes) return 1;
if (bytes == 0) bytes = 1;
if (bytes > NSUIntegerMax) return 0;
MTLResourceOptions options = MTLResourceStorageModeShared;
if (ds4_gpu_use_m5_private_scratch() &&
!ds4_gpu_scratch_needs_cpu_access(label)) {
/*
* M5 scratch buffers that only flow between Metal kernels do not need
* CPU-visible shared storage. This reduces shared-memory traffic and
* residency pressure for the long prefill scratch pools without
* changing the public buffer lifetime model. Keep default hazard
* tracking because the graph reuses these buffers across dependent
* compute encoders.
*/
options = MTLResourceStorageModePrivate;
}
*buffer = [g_device newBufferWithLength:bytes options:options];
if (!*buffer && options != MTLResourceStorageModeShared) {
*buffer = [g_device newBufferWithLength:bytes options:MTLResourceStorageModeShared];
}
if (!*buffer) {
fprintf(stderr, "ds4: failed to allocate Metal scratch buffer %s (%llu bytes)\n",
label, (unsigned long long)bytes);
*capacity = 0;
return 0;
}
(*buffer).label = [NSString stringWithUTF8String:label];
*capacity = bytes;
return 1;
}
static int ds4_gpu_ensure_zero_attention_mask(NSUInteger bytes) {
const NSUInteger capacity = 8192u * sizeof(uint16_t);
if (bytes > capacity) return 0;
if (g_flash_attn_zero_mask_buffer &&
g_flash_attn_zero_mask_bytes >= capacity) {
return 1;
}
if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_zero_mask_buffer,
&g_flash_attn_zero_mask_bytes,
capacity,
"ds4_flash_attn_zero_mask")) {
return 0;
}
void *contents = [g_flash_attn_zero_mask_buffer contents];
if (!contents) return 0;
memset(contents, 0, g_flash_attn_zero_mask_bytes);
return 1;
}
static uint64_t round_up_u64(uint64_t v, uint64_t align) {
return (v + align - 1) & ~(align - 1);
}
static uint64_t ds4_gpu_effective_model_max_tensor_bytes(uint64_t map_size, uint64_t max_tensor_bytes) {
if (max_tensor_bytes != 0) return max_tensor_bytes;
return map_size < DS4_METAL_FALLBACK_MAX_TENSOR_BYTES ?
map_size : DS4_METAL_FALLBACK_MAX_TENSOR_BYTES;
}
static id<MTLComputePipelineState> ds4_gpu_get_pipeline(const char *function_name);
static int ds4_gpu_warm_model_views(void);
static double ds4_gpu_gib(uint64_t bytes);
static double ds4_gpu_now_ms(void) {
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return ts.tv_sec * 1000.0 + ts.tv_nsec / 1000000.0;
}
static int ds4_gpu_moe_selected_hotlist_cmp(const void *a, const void *b) {
const ds4_gpu_moe_selected_hotlist_entry *ea = a;
const ds4_gpu_moe_selected_hotlist_entry *eb = b;
if (ea->hits < eb->hits) return 1;
if (ea->hits > eb->hits) return -1;
if (ea->layer != eb->layer) return ea->layer < eb->layer ? -1 : 1;
if (ea->expert != eb->expert) return ea->expert < eb->expert ? -1 : 1;
return 0;
}
static int ds4_gpu_moe_selected_hotlist_merge_requested(void) {
return getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_MERGE") != NULL &&
getenv("DS4_MOE_RECORD_SELECTED_HOTLIST_FRESH") == NULL;
}
static int ds4_gpu_moe_selected_hotlist_load_existing(const char *path) {
if (!path || !path[0] || !ds4_gpu_moe_selected_hotlist_merge_requested()) {
return 1;
}
FILE *fp = fopen(path, "rb");
if (!fp) {
if (errno == ENOENT) return 1;
fprintf(stderr, "ds4: failed to open selected hotlist merge file %s\n", path);
return 0;
}
char line[256];
uint64_t lineno = 0;
uint64_t loaded_entries = 0;
uint64_t loaded_hits = 0;
uint64_t header_records = UINT64_MAX;
uint64_t header_selections = UINT64_MAX;
while (fgets(line, sizeof(line), fp)) {
lineno++;
char *p = line;
while (*p && isspace((unsigned char)*p)) p++;
if (*p == '\0') continue;
if (*p == '#') {
unsigned long long value = 0;
if (sscanf(p, "# layer_records %llu", &value) == 1) {
header_records = (uint64_t)value;
} else if (sscanf(p, "# selections %llu", &value) == 1) {
header_selections = (uint64_t)value;
}
continue;
}
errno = 0;
char *end = NULL;
unsigned long layer = strtoul(p, &end, 10);
if (end == p || errno != 0) goto bad_line;
p = end;
while (*p && isspace((unsigned char)*p)) p++;
errno = 0;
unsigned long expert = strtoul(p, &end, 10);
if (end == p || errno != 0) goto bad_line;
p = end;
while (*p && isspace((unsigned char)*p)) p++;
errno = 0;
unsigned long long hits = strtoull(p, &end, 10);
if (end == p || errno != 0) goto bad_line;
if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER &&
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
hits != 0) {
uint64_t *dst = &g_moe_selected_hotlist_counts[layer][expert];
if (*dst > UINT64_MAX - (uint64_t)hits) {
*dst = UINT64_MAX;
} else {
*dst += (uint64_t)hits;
}
loaded_entries++;
if (loaded_hits > UINT64_MAX - (uint64_t)hits) {
loaded_hits = UINT64_MAX;
} else {
loaded_hits += (uint64_t)hits;
}
}
continue;
bad_line:
fprintf(stderr,
"ds4: invalid selected hotlist merge line %" PRIu64 " in %s\n",
lineno,
path);
fclose(fp);
return 0;
}
if (ferror(fp)) {
fprintf(stderr, "ds4: failed to read selected hotlist merge file %s\n", path);
fclose(fp);
return 0;
}
fclose(fp);
g_moe_selected_hotlist_records =
header_records != UINT64_MAX ? header_records : loaded_hits / 6u;
g_moe_selected_hotlist_selections =
header_selections != UINT64_MAX ? header_selections : loaded_hits;
fprintf(stderr,
"ds4: merged selected-id hotlist %s "
"(%" PRIu64 " entries, %" PRIu64 " hits)\n",
path,
loaded_entries,
loaded_hits);
return 1;
}
static void ds4_gpu_moe_selected_hotlist_close(void) {
const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST");
if (!g_moe_selected_hotlist_initialized || !path || !path[0]) return;
const size_t cap =
(size_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER *
DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
ds4_gpu_moe_selected_hotlist_entry *entries =
malloc(cap * sizeof(entries[0]));
if (!entries) {
fprintf(stderr, "ds4: failed to allocate selected hotlist entries\n");
return;
}
size_t n = 0;
for (uint32_t layer = 0;
layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
layer++) {
for (uint32_t expert = 0;
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
expert++) {
const uint64_t hits =
g_moe_selected_hotlist_counts[layer][expert];
if (hits == 0) continue;
entries[n++] = (ds4_gpu_moe_selected_hotlist_entry) {
.layer = layer,
.expert = expert,
.hits = hits,
};
}
}
qsort(entries, n, sizeof(entries[0]), ds4_gpu_moe_selected_hotlist_cmp);
FILE *fp = fopen(path, "wb");
if (!fp) {
fprintf(stderr, "ds4: failed to open selected hotlist file %s\n", path);
free(entries);
return;
}
fprintf(fp,
"# ds4 selected-id hotlist v1\n"
"# layer_records %" PRIu64 "\n"
"# selections %" PRIu64 "\n"
"# columns: layer expert hits weight\n",
g_moe_selected_hotlist_records,
g_moe_selected_hotlist_selections);
for (size_t i = 0; i < n; i++) {
fprintf(fp,
"%u %u %" PRIu64 " 0\n",
entries[i].layer,
entries[i].expert,
entries[i].hits);
}
free(entries);
if (fclose(fp) != 0) {
fprintf(stderr, "ds4: failed to close selected hotlist file %s\n", path);
} else {
fprintf(stderr,
"ds4: wrote selected-id hotlist to %s "
"(%" PRIu64 " layer records, %" PRIu64 " selections)\n",
path,
g_moe_selected_hotlist_records,
g_moe_selected_hotlist_selections);
}
}
static int ds4_gpu_moe_selected_hotlist_record(
uint32_t layer,
const int32_t *selected_ids,
uint32_t n_selected,
uint32_t n_total_expert) {
const char *path = getenv("DS4_MOE_RECORD_SELECTED_HOTLIST");
if (!path || !path[0]) return 1;
if (!selected_ids ||
n_selected == 0 ||
n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
return 0;
}
if (!g_moe_selected_hotlist_initialized) {
g_moe_selected_hotlist_initialized = 1;
if (!ds4_gpu_moe_selected_hotlist_load_existing(path)) return 0;
atexit(ds4_gpu_moe_selected_hotlist_close);
}
if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 1;
g_moe_selected_hotlist_records++;
for (uint32_t i = 0; i < n_selected; i++) {
if (selected_ids[i] < 0) continue;
const uint32_t expert = (uint32_t)selected_ids[i];
if (expert >= n_total_expert ||
expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
continue;
}
g_moe_selected_hotlist_counts[layer][expert]++;
g_moe_selected_hotlist_selections++;
}
return 1;
}
static void ds4_gpu_moe_selected_trace_record_close(void) {
if (g_moe_selected_trace_record_fp) {
const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS");
fclose(g_moe_selected_trace_record_fp);
g_moe_selected_trace_record_fp = NULL;
fprintf(stderr,
"ds4: recorded %" PRIu64 " routed-MoE selected-id entries to %s\n",
g_moe_selected_trace_record_count,
path && path[0] ? path : "(unknown)");
}
}
static int ds4_gpu_moe_selected_trace_record(
const int32_t selected_ids[6],
uint32_t n_selected) {
const char *path = getenv("DS4_MOE_RECORD_SELECTED_IDS");
if (!path || !path[0]) return 1;
if (n_selected != 6) {
fprintf(stderr, "ds4: selected-id recording expects exactly 6 selected experts\n");
return 0;
}
if (!g_moe_selected_trace_record_initialized) {
g_moe_selected_trace_record_initialized = 1;
g_moe_selected_trace_record_fp = fopen(path, "wb");
if (!g_moe_selected_trace_record_fp) {
fprintf(stderr, "ds4: failed to open selected-id record file %s\n", path);
return 0;
}
setvbuf(g_moe_selected_trace_record_fp, NULL, _IOFBF, 1u << 20);
atexit(ds4_gpu_moe_selected_trace_record_close);
}
if (fwrite(selected_ids, sizeof(selected_ids[0]), n_selected, g_moe_selected_trace_record_fp) != n_selected) {
fprintf(stderr, "ds4: failed to write selected-id record file %s\n", path);
return 0;
}
if (fflush(g_moe_selected_trace_record_fp) != 0) {
fprintf(stderr, "ds4: failed to flush selected-id record file %s\n", path);
return 0;
}
g_moe_selected_trace_record_count++;
return 1;
}
static int ds4_gpu_moe_selected_trace_replay(
int32_t selected_ids[6],
uint32_t n_selected) {
const char *path = getenv("DS4_MOE_REPLAY_SELECTED_IDS");
if (!path || !path[0]) return 0;
if (n_selected != 6) {
fprintf(stderr, "ds4: selected-id replay expects exactly 6 selected experts\n");
return -1;
}
if (!g_moe_selected_trace_replay_initialized) {
g_moe_selected_trace_replay_initialized = 1;
FILE *fp = fopen(path, "rb");
if (!fp) {
fprintf(stderr, "ds4: failed to open selected-id replay file %s\n", path);
return -1;
}
if (fseeko(fp, 0, SEEK_END) != 0) {
fprintf(stderr, "ds4: failed to seek selected-id replay file %s\n", path);
fclose(fp);
return -1;
}
const off_t end = ftello(fp);
if (end < 0) {
fprintf(stderr, "ds4: failed to size selected-id replay file %s\n", path);
fclose(fp);
return -1;
}
if (fseeko(fp, 0, SEEK_SET) != 0) {
fprintf(stderr, "ds4: failed to rewind selected-id replay file %s\n", path);
fclose(fp);
return -1;
}
const uint64_t bytes = (uint64_t)end;
const uint64_t entry_bytes = (uint64_t)n_selected * sizeof(selected_ids[0]);
if (bytes == 0 || (bytes % entry_bytes) != 0) {
fprintf(stderr,
"ds4: selected-id replay file %s has invalid size %" PRIu64 "\n",
path,
bytes);
fclose(fp);
return -1;
}
if (bytes > SIZE_MAX) {
fprintf(stderr, "ds4: selected-id replay file %s is too large\n", path);
fclose(fp);
return -1;
}
g_moe_selected_trace_replay_count = bytes / entry_bytes;
g_moe_selected_trace_replay_ids = malloc((size_t)bytes);
if (!g_moe_selected_trace_replay_ids) {
fprintf(stderr, "ds4: failed to allocate selected-id replay buffer\n");
fclose(fp);
return -1;
}
if (fread(g_moe_selected_trace_replay_ids, 1, (size_t)bytes, fp) != (size_t)bytes) {
fprintf(stderr, "ds4: failed to read selected-id replay file %s\n", path);
fclose(fp);
free(g_moe_selected_trace_replay_ids);
g_moe_selected_trace_replay_ids = NULL;
return -1;
}
fclose(fp);
fprintf(stderr,
"ds4: loaded %" PRIu64 " routed-MoE selected-id entries from %s\n",
g_moe_selected_trace_replay_count,
path);
}
if (g_moe_selected_trace_replay_pos >= g_moe_selected_trace_replay_count) {
fprintf(stderr,
"ds4: selected-id replay exhausted after %" PRIu64 " entries\n",
g_moe_selected_trace_replay_pos);
return -1;
}
memcpy(selected_ids,
g_moe_selected_trace_replay_ids + g_moe_selected_trace_replay_pos * n_selected,
(size_t)n_selected * sizeof(selected_ids[0]));
g_moe_selected_trace_replay_pos++;
return 1;
}
static int ds4_gpu_progress_enabled(void) {
return ds4_log_is_tty(stderr);
}
static void ds4_gpu_progress_begin(const char *what) {
if (!ds4_gpu_progress_enabled()) return;
fprintf(stderr, "ds4: %s...", what);
fflush(stderr);
}
static void ds4_gpu_progress_done(void) {
if (!ds4_gpu_progress_enabled()) return;
fputs(" done\n", stderr);
fflush(stderr);
}
static void ds4_gpu_progress_failed(void) {
if (!ds4_gpu_progress_enabled()) return;
fputs(" failed\n", stderr);
fflush(stderr);
}
static void ds4_gpu_model_views_clear(void) {
for (uint32_t i = 0; i < g_model_view_count; 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;
}
g_model_view_count = 0;
}
static void ds4_gpu_model_views_remove_map(const void *model_map) {
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[kept++] = g_model_views[i];
}
for (uint32_t i = kept; i < g_model_view_count; i++)
g_model_views[i] = (ds4_gpu_model_view){0};
g_model_view_count = kept;
}
static void ds4_gpu_model_residency_clear(void) {
#if TARGET_OS_OSX
if (@available(macOS 15.0, *)) {
if (g_model_residency_set) {
if (g_model_residency_added_to_queue &&
g_queue &&
[g_queue respondsToSelector:@selector(removeResidencySet:)]) {
[g_queue removeResidencySet:g_model_residency_set];
}
[g_model_residency_set endResidency];
[g_model_residency_set removeAllAllocations];
g_model_residency_set = nil;
}
}
#endif
g_model_residency_count = 0;
g_model_residency_added_to_queue = 0;
}
/* TP sharding keeps only this rank's expert ranges warm,
* so whole-view residency requests (which would page in the full file)
* must be skipped; pages fault in lazily through the same view buffers,
* exactly like ssd-streaming mode. */
static int g_model_residency_skipped;
void ds4_gpu_model_residency_skip(int skip) {
g_model_residency_skipped = skip;
}
static int ds4_gpu_model_residency_request_views(void) {
if (g_model_view_count == 0 ||
g_ssd_streaming_mode ||
g_model_residency_skipped ||
getenv("DS4_METAL_NO_RESIDENCY") != NULL) {
return 1;
}
#if TARGET_OS_OSX
if (@available(macOS 15.0, *)) {
/*
* Register all model views as one residency set before inference. This
* is a GPU residency/budgeting hint, not a request to fault the whole
* 80+ GB file into memory. Its purpose is to make the driver see the
* complete set of large shared allocations during setup instead of
* discovering them lazily from the first measured graph command, where
* VM validation and residency accounting would look like model compute.
*/
MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init];
desc.label = @"ds4_model";
desc.initialCapacity = g_model_view_count;
NSError *error = nil;
g_model_residency_set = [g_device newResidencySetWithDescriptor:desc error:&error];
if (!g_model_residency_set) {
fprintf(stderr, "ds4: Metal model residency set creation failed: %s\n",
[[error localizedDescription] UTF8String]);
return 0;
}
for (uint32_t i = 0; i < g_model_view_count; i++) {
[g_model_residency_set addAllocation:g_model_views[i].buffer];
}
[g_model_residency_set commit];
[g_model_residency_set requestResidency];
if (getenv("DS4_METAL_DISABLE_QUEUE_RESIDENCY_SET") == NULL &&
g_queue &&
[g_queue respondsToSelector:@selector(addResidencySet:)]) {
[g_queue addResidencySet:g_model_residency_set];
g_model_residency_added_to_queue = 1;
}
g_model_residency_count = g_model_view_count;
}
#endif
return 1;
}
static int ds4_gpu_add_model_view_range(
const void *model_map,
uint64_t model_size,
uint64_t map_offset,
uint64_t map_size,
uint64_t max_tensor_bytes,
bool use_default_view_cap,
uint64_t *mapped_model_size_out) {
const uint64_t page = (uint64_t)getpagesize();
const uintptr_t model_addr = (uintptr_t)model_map;
if ((model_addr & (uintptr_t)(page - 1)) != 0) {
fprintf(stderr, "ds4: Metal model mmap base is not page aligned\n");
return 0;
}
if (map_offset > model_size || map_size > model_size - map_offset) {
fprintf(stderr, "ds4: Metal model mapped range is outside the GGUF mapping\n");
return 0;
}
const uint64_t page_model_offset = map_offset & ~(page - 1);
const uint64_t leading = map_offset - page_model_offset;
if (map_size > UINT64_MAX - leading ||
leading + map_size > UINT64_MAX - (page - 1))
{
fprintf(stderr, "ds4: Metal model mapped range overflows page alignment\n");
return 0;
}
const uint64_t mapped_model_size = round_up_u64(leading + map_size, page);
uint64_t max_buffer = (uint64_t)[g_device maxBufferLength];
max_buffer &= ~(page - 1);
/*
* Wrap only the tensor-data part of the GGUF file. Metadata is parsed by the
* CPU and is never dereferenced by kernels, so exposing it to Metal only
* grows the residency set and the VM range the driver must validate.
*
* Metal buffers have a device-specific maximum length, and this model is
* larger than that maximum on the target machines. Creating one no-copy
* buffer per tensor would avoid the length limit, but it would also move a
* lot of VM-object creation and residency bookkeeping into graph setup. The
* stable shape here is a tiny number of page-aligned views created once.
*
* Adjacent views intentionally overlap by more than the largest tensor, plus
* one page for alignment. That invariant guarantees every tensor lies wholly
* inside at least one view, so hot paths pass one buffer and one inner byte
* offset. We never split a weight tensor across command encoders.
*/
if (max_tensor_bytes > map_size) {
fprintf(stderr, "ds4: Metal model max tensor span is larger than a mapped tensor span\n");
return 0;
}
if (max_tensor_bytes > UINT64_MAX - (page - 1)) {
fprintf(stderr, "ds4: Metal model max tensor span overflows page alignment\n");
return 0;
}
const uint64_t max_tensor_rounded = round_up_u64(max_tensor_bytes, page);
if (max_tensor_rounded > UINT64_MAX - page) {
fprintf(stderr, "ds4: Metal model view overlap overflows page slack\n");
return 0;
}
const uint64_t overlap = max_tensor_rounded + page;
if (max_buffer == 0 || max_buffer <= overlap) {
fprintf(stderr,
"ds4: Metal maxBufferLength is too small for DS4 model views "
"(max tensor %.2f GiB, max buffer %.2f GiB)\n",
ds4_gpu_gib(max_tensor_bytes),
ds4_gpu_gib(max_buffer));
return 0;
}
uint64_t view_limit = max_buffer;
const char *view_limit_env = getenv("DS4_METAL_MODEL_VIEW_MAX_GIB");
if (view_limit_env && view_limit_env[0]) {
char *end = NULL;
unsigned long long gib = strtoull(view_limit_env, &end, 10);
if (end != view_limit_env && gib > 0) {
uint64_t env_limit = gib * 1024ull * 1024ull * 1024ull;
env_limit &= ~(page - 1);
if (env_limit > 0) view_limit = env_limit;
}
} else if (use_default_view_cap && mapped_model_size > max_buffer) {
/*
* Very large no-copy buffers can make Metal's VM validation dominate
* startup or the first graph command on multi-hundred-GiB slices. Keep
* ordinary contiguous model mappings unchanged, but let distributed
* span maps use smaller overlapping views when a range already has to
* be split.
*/
const uint64_t default_limit = 128ull * 1024ull * 1024ull * 1024ull;
if (view_limit > default_limit) view_limit = default_limit;
}
if (view_limit > max_buffer) view_limit = max_buffer;
view_limit &= ~(page - 1);
if (view_limit == 0 || view_limit <= overlap) {
fprintf(stderr,
"ds4: Metal model view cap is too small for DS4 model views "
"(cap %.2f GiB, max tensor %.2f GiB)\n",
ds4_gpu_gib(view_limit),
ds4_gpu_gib(max_tensor_bytes));
return 0;
}
const uint64_t step = view_limit - overlap;
uint64_t off = 0;
while (off < mapped_model_size) {
if (g_model_view_count == DS4_METAL_MAX_MODEL_VIEWS) {
fprintf(stderr, "ds4: Metal model needs more mapped views than expected\n");
return 0;
}
uint64_t view_bytes = mapped_model_size - off;
if (view_bytes > view_limit) view_bytes = view_limit;
id<MTLBuffer> buffer = [g_device newBufferWithBytesNoCopy:(void *)(model_addr + page_model_offset + off)
length:(NSUInteger)view_bytes
options:ds4_gpu_model_resource_options()
deallocator:nil];
if (!buffer) {
fprintf(stderr,
"ds4: Metal could not wrap mmaped model view at %.2f GiB, size %.2f GiB\n",
(double)(page_model_offset + off) / (1024.0 * 1024.0 * 1024.0),
(double)view_bytes / (1024.0 * 1024.0 * 1024.0));
return 0;
}
buffer.label = [NSString stringWithFormat:@"ds4_model_view_%u", g_model_view_count];
g_model_views[g_model_view_count].buffer = buffer;
g_model_views[g_model_view_count].model_map = model_map;
g_model_views[g_model_view_count].model_size = model_size;
g_model_views[g_model_view_count].model_offset = page_model_offset + off;
g_model_views[g_model_view_count].bytes = view_bytes;
g_model_view_count++;
g_model_wrap_count++;
g_model_wrap_bytes += view_bytes;
if (view_bytes > g_model_wrap_max_bytes) g_model_wrap_max_bytes = view_bytes;
if (off + view_bytes >= mapped_model_size) break;
off += step;
}
if (mapped_model_size_out) *mapped_model_size_out += mapped_model_size;
return 1;
}
static int ds4_gpu_finish_model_views(
double t0,
uint64_t mapped_model_size,
uint64_t display_offset) {
const double t_mapped = ds4_gpu_now_ms();
const int request_residency =
!g_ssd_streaming_mode &&
getenv("DS4_METAL_NO_RESIDENCY") == NULL;
if (request_residency) ds4_gpu_progress_begin("requesting Metal residency (may take tens of seconds)");
if (!ds4_gpu_model_residency_request_views()) {
if (request_residency) ds4_gpu_progress_failed();
return 0;
}
if (request_residency) ds4_gpu_progress_done();
const double t_resident = ds4_gpu_now_ms();
int warmed = 1;
const double t_warm0 = ds4_gpu_now_ms();
const int warm_model_views = !g_ssd_streaming_mode &&
getenv("DS4_METAL_NO_RESIDENCY") == NULL &&
getenv("DS4_METAL_NO_MODEL_WARMUP") == NULL;
if (warm_model_views) {
/*
* The first GPU command touching no-copy mmap storage can pay command
* queue setup, page-table validation, and shared-allocation residency
* costs. Sample each model view here so timed graph execution starts
* after that one-time work. The stride is intentionally coarse: this is
* a validation touch over the VM ranges, not a full model prefetch. A
* dense prefetch would create exactly the kind of memory pressure and
* startup stalls this path is designed to avoid.
*/
if (g_model_residency_skipped) {
/* TP sharding: a single command buffer binding every
* view demands residency of them all and OOMs; the engine's
* CPU-side sharded warm pre-faults the owned bytes instead. */
warmed = 1;
} else {
ds4_gpu_progress_begin("warming Metal model views");
warmed = ds4_gpu_warm_model_views();
if (warmed) ds4_gpu_progress_done();
else ds4_gpu_progress_failed();
}
}
ds4_gpu_queue_keepalive_start();
const double t_warm = ds4_gpu_now_ms();
if (ds4_gpu_model_map_log_enabled()) {
fprintf(stderr,
"ds4: Metal model views created in %.3f ms, residency requested in %.3f ms, warmup %.3f ms (mapped %.2f MiB from offset %.2f MiB)\n",
t_mapped - t0,
t_resident - t_mapped,
t_warm - t_warm0,
mapped_model_size / 1024.0 / 1024.0,
display_offset / 1024.0 / 1024.0);
}
if (!warmed) return 0;
return 1;
}
static int ds4_gpu_map_model_views(
const void *model_map,
uint64_t model_size,
uint64_t map_offset,
uint64_t map_size,
uint64_t max_tensor_bytes) {
const double t0 = ds4_gpu_now_ms();
uint64_t mapped_model_size = 0;
if (!ds4_gpu_add_model_view_range(model_map,
model_size,
map_offset,
map_size,
max_tensor_bytes,
false,
&mapped_model_size)) {
return 0;
}
return ds4_gpu_finish_model_views(t0, mapped_model_size, map_offset);
}
static id<MTLBuffer> ds4_gpu_new_transient_buffer(NSUInteger bytes, const char *label) {
if (bytes == 0) bytes = 1;
id<MTLBuffer> buffer = [g_device newBufferWithLength:bytes
options:MTLResourceStorageModeShared];
if (!buffer) {
fprintf(stderr, "ds4: failed to allocate Metal transient buffer %s (%llu bytes)\n",
label ? label : "(unnamed)", (unsigned long long)bytes);
return nil;
}
if (label) buffer.label = [NSString stringWithUTF8String:label];
/*
* CPU-filled buffers must survive until their command buffer completes.
* A local ObjC strong variable is not enough when the encoder function
* returns before the caller commits the command buffer.
*/
[g_transient_buffers addObject:buffer];
return buffer;
}
static int ds4_gpu_zero_prefix_prefill_mask_cache_enabled(void) {
if (getenv("DS4_METAL_DISABLE_ZERO_PREFIX_PREFILL_MASK_CACHE") != NULL ||
getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL) {
return 0;
}
return ds4_gpu_device_name_contains("M3");
}
static ds4_gpu_zero_prefix_prefill_mask_cache_entry *
ds4_gpu_get_zero_prefix_prefill_mask_cache(
uint32_t kind,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t n_keys,
uint32_t window,
uint32_t ratio,
uint32_t nqptg,
uint32_t ncpsg,
bool has_kvpad,
bool bc_mask,
NSUInteger mask_bytes,
NSUInteger blk_bytes,
bool *created) {
if (created) *created = false;
if (!created || !ds4_gpu_zero_prefix_prefill_mask_cache_enabled() ||
mask_bytes == 0 || blk_bytes == 0) {
return NULL;
}
uint32_t slot = UINT32_MAX;
switch (kind) {
case DS4_GPU_PREFILL_MASK_CACHE_RAW: slot = 0; break;
case DS4_GPU_PREFILL_MASK_CACHE_RATIO4: slot = 1; break;
case DS4_GPU_PREFILL_MASK_CACHE_RATIO128: slot = 2; break;
default: return NULL;
}
ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry =
&g_zero_prefix_prefill_mask_cache[slot];
if (entry->valid && entry->mask && entry->blk &&
entry->mask_bytes == mask_bytes && entry->blk_bytes == blk_bytes &&
entry->kind == kind && entry->n_tokens == n_tokens &&
entry->n_comp == n_comp && entry->n_keys == n_keys &&
entry->window == window && entry->ratio == ratio &&
entry->nqptg == nqptg && entry->ncpsg == ncpsg &&
entry->has_kvpad == has_kvpad && entry->bc_mask == bc_mask) {
return entry;
}
id<MTLBuffer> mask = [g_device newBufferWithLength:mask_bytes
options:MTLResourceStorageModeShared];
id<MTLBuffer> blk = [g_device newBufferWithLength:blk_bytes
options:MTLResourceStorageModePrivate];
if (!blk) {
blk = [g_device newBufferWithLength:blk_bytes
options:MTLResourceStorageModeShared];
}
if (!mask || !blk) return NULL;
mask.label = [NSString stringWithFormat:@"ds4_prefill_mask_cache_%u", kind];
blk.label = [NSString stringWithFormat:@"ds4_prefill_blk_cache_%u", kind];
/* A replaced entry may still be referenced by an uncommitted or in-flight
* command buffer. Keep its resources alive with the other batch-scoped
* buffers until command completion instead of mutating or releasing them. */
if (entry->mask) [g_transient_buffers addObject:entry->mask];
if (entry->blk) [g_transient_buffers addObject:entry->blk];
entry->mask = mask;
entry->blk = blk;
entry->mask_bytes = mask_bytes;
entry->blk_bytes = blk_bytes;
entry->kind = kind;
entry->n_tokens = n_tokens;
entry->n_comp = n_comp;
entry->n_keys = n_keys;
entry->window = window;
entry->ratio = ratio;
entry->nqptg = nqptg;
entry->ncpsg = ncpsg;
entry->has_kvpad = has_kvpad;
entry->bc_mask = bc_mask;
/* The caller publishes the entry only after synchronously filling mask. */
entry->valid = false;
entry->blk_ready = false;
*created = true;
return entry;
}
static void ds4_gpu_invalidate_zero_prefix_prefill_block_maps(void) {
for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) {
g_zero_prefix_prefill_mask_cache[i].blk_ready = false;
}
}
static void ds4_gpu_clear_zero_prefix_prefill_mask_cache(void) {
for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) {
ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry =
&g_zero_prefix_prefill_mask_cache[i];
entry->mask = nil;
entry->blk = nil;
entry->mask_bytes = 0;
entry->blk_bytes = 0;
entry->kind = 0;
entry->n_tokens = 0;
entry->n_comp = 0;
entry->n_keys = 0;
entry->window = 0;
entry->ratio = 0;
entry->nqptg = 0;
entry->ncpsg = 0;
entry->has_kvpad = false;
entry->bc_mask = false;
entry->valid = false;
entry->blk_ready = false;
}
}
void ds4_gpu_release_zero_prefix_prefill_mask_cache(void) {
/* Layer-major prefill waits each layer before advancing, so its final
* release point has no outstanding cache users. Keep the guard here for
* diagnostic callers that may have an open or asynchronously flushed CB. */
if (!g_initialized || g_batch_cb ||
(g_pending_cbs && [g_pending_cbs count] != 0)) {
return;
}
ds4_gpu_clear_zero_prefix_prefill_mask_cache();
}
static id<MTLComputePipelineState> ds4_gpu_get_mul_mm_pipeline(
const char *function_name,
bool bc_inp,
bool bc_out) {
NSString *key = [NSString stringWithFormat:@"%s_bci=%d_bco=%d",
function_name, bc_inp ? 1 : 0, bc_out ? 1 : 0];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) return cached;
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700];
[constants setConstantValue:&bc_out type:MTLDataTypeBool atIndex:701];
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal %s function not found: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
return pipeline;
}
static id<MTLComputePipelineState> ds4_gpu_get_mul_mm_id_pipeline(
const char *function_name,
bool bc_inp) {
NSString *key = [NSString stringWithFormat:@"%s_bci=%d",
function_name, bc_inp ? 1 : 0];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) return cached;
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&bc_inp type:MTLDataTypeBool atIndex:700];
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal %s function not found: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
return pipeline;
}
static id<MTLComputePipelineState> ds4_gpu_get_pipeline(
const char *function_name) {
NSString *key = [NSString stringWithFormat:@"%s", function_name];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
/* Failed lookups are remembered as NSNull so the error prints once
* per name instead of on every dispatch. */
if (cached == (id<MTLComputePipelineState>)[NSNull null]) return nil;
return cached;
}
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name];
if (!fn) {
fprintf(stderr, "ds4: Metal %s function not found\n", function_name);
[g_pipeline_cache setObject:(id<MTLComputePipelineState>)[NSNull null] forKey:key];
return nil;
}
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
function_name, [[error localizedDescription] UTF8String]);
[g_pipeline_cache setObject:(id<MTLComputePipelineState>)[NSNull null] forKey:key];
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
return pipeline;
}
static int ds4_gpu_disable_hot_pipeline_statics(void) {
static int initialized;
static int disabled;
if (!initialized) {
disabled = getenv("DS4_METAL_DISABLE_HOT_PIPELINE_STATICS") != NULL;
initialized = 1;
}
return disabled;
}
static id<MTLComputePipelineState> ds4_gpu_hot_pipeline(
id<MTLComputePipelineState> pipeline,
const char *fallback_name) {
if (!ds4_gpu_disable_hot_pipeline_statics()) return pipeline;
return ds4_gpu_get_pipeline(fallback_name);
}
static int ds4_gpu_use_compressor_pair_nr4(void) {
static int initialized;
static int enabled;
if (!initialized) {
enabled = getenv("DS4_METAL_COMPRESSOR_PAIR_NR4") != NULL;
initialized = 1;
}
return enabled;
}
static int ds4_gpu_device_name_contains(const char *needle);
static int ds4_gpu_env_value_eq(const char *v, size_t n, const char *literal) {
size_t m = strlen(literal);
if (n != m) return 0;
for (size_t i = 0; i < n; i++) {
if (tolower((unsigned char)v[i]) != tolower((unsigned char)literal[i])) return 0;
}
return 1;
}
static int ds4_gpu_env_bool(const char *name) {
const char *v = getenv(name);
if (!v) return -1;
while (isspace((unsigned char)*v)) v++;
size_t n = strlen(v);
while (n > 0 && isspace((unsigned char)v[n - 1])) n--;
if (n == 0) return 1;
if (ds4_gpu_env_value_eq(v, n, "1") ||
ds4_gpu_env_value_eq(v, n, "true") ||
ds4_gpu_env_value_eq(v, n, "yes") ||
ds4_gpu_env_value_eq(v, n, "on")) {
return 1;
}
if (ds4_gpu_env_value_eq(v, n, "0") ||
ds4_gpu_env_value_eq(v, n, "false") ||
ds4_gpu_env_value_eq(v, n, "no") ||
ds4_gpu_env_value_eq(v, n, "off")) {
return 0;
}
if (!g_mpp_invalid_env_reported) {
fprintf(stderr,
"ds4: invalid Metal boolean environment value %s=%.*s; treating presence as enabled\n",
name, (int)n, v);
g_mpp_invalid_env_reported = 1;
}
return 1;
}
static uint64_t ds4_gpu_env_u64(const char *name,
uint64_t fallback,
uint64_t min_value,
uint64_t max_value) {
const char *v = getenv(name);
if (!v) return fallback;
while (isspace((unsigned char)*v)) v++;
if (!*v) return fallback;
errno = 0;
char *end = NULL;
unsigned long long parsed = strtoull(v, &end, 10);
if (end == v || errno == ERANGE) return fallback;
while (isspace((unsigned char)*end)) end++;
if (*end) return fallback;
if (parsed < min_value) return fallback;
uint64_t value = (uint64_t)parsed;
if (value > max_value) value = max_value;
return value;
}
static uint32_t ds4_gpu_glm_full_attention_max_cache_len(void) {
/*
* kernel_glm_attention_full stores one score per visible token in
* threadgroup memory, plus 256 reduction slots. Keep the default under
* the 32 KiB envelope used by current Apple GPUs while allowing short
* dense decode to step past the 4096-token prefill boundary.
*/
return 7680u;
}
static uint32_t ds4_gpu_glm_flash_attention_max_cache_len(void) {
/*
* Staged FlashAttention uses fixed-size threadgroup scratch and separate
* KV staging buffers, so it is not bound by the legacy full-attention
* kernel's one-score-per-token threadgroup-memory envelope.
*/
return 8192u;
}
static int ds4_gpu_mpp_available(void) {
return g_metal4_tensor_api_enabled && !g_quality_mode;
}
/*
* Retained Metal4 defaults live here instead of behind user-visible options.
* The public runtime has one automatic accelerated path plus the global
* DS4_METAL_DISABLE_METAL4 comparison switch. Benchmark-only alternatives that
* lost during M5 work are removed or kept out of the dispatch path so future
* changes do not accidentally turn old experiments into new modes.
*/
static int ds4_gpu_use_mpp_attn_out_low_matmul(void) {
return ds4_gpu_mpp_available();
}
enum {
DS4_METAL_ATTN_OUT_MPP_TILE_N = 64,
};
static void ds4_gpu_warn_mpp_fallback(void) {
static int warned;
if (!warned) {
fprintf(stderr, "ds4: accelerated Metal prefill matmul unavailable; falling back to legacy kernel\n");
warned = 1;
}
}
static int ds4_gpu_device_name_contains(const char *needle) {
return g_metal_device_name[0] != '\0' && strstr(g_metal_device_name, needle) != NULL;
}
int ds4_gpu_device_is_pre_m5_apple_silicon(void) {
return strncmp(g_metal_device_name, "Apple M", 7) == 0 &&
g_metal_device_name[7] >= '1' &&
g_metal_device_name[7] <= '4' &&
(g_metal_device_name[8] == '\0' ||
g_metal_device_name[8] == ' ');
}
int ds4_gpu_device_is_m5_apple_silicon(void) {
return strncmp(g_metal_device_name, "Apple M5", 8) == 0 &&
(g_metal_device_name[8] == '\0' ||
g_metal_device_name[8] == ' ');
}
static bool ds4_gpu_ported_m5_decode_feature_enabled(
const char *pre_m5_disable_env,
const char *m5_disable_env) {
if (m5_disable_env && getenv(m5_disable_env) != NULL) return false;
const bool pre_m5 = ds4_gpu_device_is_pre_m5_apple_silicon();
if (pre_m5 &&
(getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_PORTS") != NULL ||
(pre_m5_disable_env && getenv(pre_m5_disable_env) != NULL))) {
return false;
}
return pre_m5 || ds4_gpu_device_is_m5_apple_silicon();
}
static int ds4_gpu_compile_tensor_probe(void) {
#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000
if (!g_device) return 0;
if (@available(macOS 26.0, *)) {
const char *src =
"#include <metal_stdlib>\n"
"#include <metal_tensor>\n"
"#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>\n"
"using namespace metal;\n"
"using namespace mpp::tensor_ops;\n"
"kernel void ds4_tensor_probe(\n"
" tensor<device half, dextents<int32_t, 2>> A [[buffer(0)]],\n"
" tensor<device half, dextents<int32_t, 2>> B [[buffer(1)]],\n"
" device float *C [[buffer(2)]],\n"
" uint2 tgid [[threadgroup_position_in_grid]]) {\n"
" auto tA = A.slice(0, (int)tgid.y);\n"
" auto tB = B.slice((int)tgid.x, 0);\n"
" matmul2d<matmul2d_descriptor(16, 16, dynamic_extent), execution_simdgroups<4>> mm;\n"
" auto cT = mm.get_destination_cooperative_tensor<decltype(tA), decltype(tB), float>();\n"
" auto sA = tA.slice(0, 0);\n"
" auto sB = tB.slice(0, 0);\n"
" mm.run(sB, sA, cT);\n"
" auto tC = tensor<device float, dextents<int32_t, 2>, tensor_inline>(C, dextents<int32_t, 2>(16, 16));\n"
" cT.store(tC);\n"
"}\n";
NSError *error = nil;
NSString *source = [NSString stringWithUTF8String:src];
id<MTLLibrary> probe_library = [g_device newLibraryWithSource:source options:[MTLCompileOptions new] error:&error];
if (!probe_library) {
fprintf(stderr, "ds4: Metal 4 tensor API probe compile failed: %s\n",
error ? [[error localizedDescription] UTF8String] : "(unknown)");
return 0;
}
id<MTLFunction> fn = [probe_library newFunctionWithName:@"ds4_tensor_probe"];
if (!fn) {
fprintf(stderr, "ds4: Metal 4 tensor API probe function missing\n");
return 0;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal 4 tensor API probe pipeline failed: %s\n",
error ? [[error localizedDescription] UTF8String] : "(unknown)");
return 0;
}
return 1;
}
#endif
return 0;
}
static void ds4_gpu_detect_metal4_features(void) {
g_metal4_runtime_available = 0;
g_metal4_family_supported = 0;
g_metal4_queue_supported = 0;
g_metal4_m5_neural_accelerators_hint = 0;
g_metal4_tensor_api_enabled = 0;
g_metal4_tensor_api_compile_supported = 0;
g_metal_device_name[0] = '\0';
if (!g_device) return;
const char *name = [[g_device name] UTF8String];
if (name) {
snprintf(g_metal_device_name, sizeof(g_metal_device_name), "%s", name);
}
const int metal4_disabled = ds4_gpu_env_bool("DS4_METAL_DISABLE_METAL4") > 0;
#if defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000
if (@available(macOS 26.0, *)) {
g_metal4_runtime_available = 1;
g_metal4_family_supported =
!metal4_disabled && [g_device supportsFamily:MTLGPUFamilyMetal4] ? 1 : 0;
g_metal4_queue_supported = [g_device respondsToSelector:@selector(newMTL4CommandQueue)] ? 1 : 0;
/*
* Apple does not currently expose a separate "Neural Accelerator" bit
* through Metal. On public M5 systems the hardware signal is the device
* generation plus Metal 4 support, so keep this as a conservative hint.
*/
if (g_metal4_family_supported && ds4_gpu_device_name_contains("M5")) {
g_metal4_m5_neural_accelerators_hint = 1;
}
if (g_metal4_family_supported) {
const int default_enable =
ds4_gpu_device_name_contains("M5") ||
ds4_gpu_device_name_contains("M6") ||
ds4_gpu_device_name_contains("A19") ||
ds4_gpu_device_name_contains("A20");
/*
* Metal 4 TensorOps are portable in source, but on pre-M5 hardware
* they can map to ordinary shader fallbacks. Keep the automatic
* fast path restricted to hardware generations where the Neural
* Accelerator/TensorOps path is expected to pay off; older Metal
* machines continue to use the established kernels unless a future
* device is explicitly added here.
*/
if (default_enable) {
g_metal4_tensor_api_compile_supported = ds4_gpu_compile_tensor_probe();
g_metal4_tensor_api_enabled = g_metal4_tensor_api_compile_supported;
if (!g_metal4_tensor_api_enabled) {
fprintf(stderr, "ds4: Metal 4 tensor API probe failed; using legacy Metal kernels\n");
}
} else {
fprintf(stderr, "ds4: Metal 4 tensor API disabled for pre-M5/pre-A19 devices\n");
}
}
}
#endif
}
static int ds4_gpu_warm_model_views(void) {
if (g_model_view_count == 0) return 1;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_touch_u8_stride");
if (!pipeline) return 0;
uint64_t stride = 1024ull * 1024ull;
const char *stride_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_MB");
if (stride_env && stride_env[0]) {
char *end = NULL;
unsigned long long mb = strtoull(stride_env, &end, 10);
if (end != stride_env && mb > 0 && mb <= 1024) {
stride = mb * 1024ull * 1024ull;
}
}
const char *stride_kb_env = getenv("DS4_METAL_MODEL_WARMUP_STRIDE_KB");
if (stride_kb_env && stride_kb_env[0]) {
char *end = NULL;
unsigned long long kb = strtoull(stride_kb_env, &end, 10);
if (end != stride_kb_env && kb > 0 && kb <= 1024ull * 1024ull) {
stride = kb * 1024ull;
const uint64_t page = (uint64_t)getpagesize();
if (stride < page) stride = page;
}
}
uint64_t total_touches = 0;
for (uint32_t i = 0; 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;
const NSUInteger out_bytes = (NSUInteger)total_touches;
id<MTLBuffer> out = [g_device newBufferWithLength:out_bytes
options:MTLResourceStorageModeShared];
if (!out) {
fprintf(stderr, "ds4: Metal model warmup scratch allocation failed\n");
return 0;
}
out.label = @"ds4_model_warmup";
id<MTLCommandBuffer> cb = ds4_gpu_new_command_buffer();
if (!cb) {
fprintf(stderr, "ds4: Metal model warmup command buffer allocation failed\n");
return 0;
}
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++) {
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];
[enc setBuffer:out offset:0 atIndex:1];
[enc setBytes:&stride length:sizeof(stride) atIndex:2];
[enc setBytes:&bytes length:sizeof(bytes) atIndex:3];
[enc setBytes:&dst_offset length:sizeof(dst_offset) atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
dst_offset += n;
}
ds4_gpu_end_compute_encoder(cb, enc);
[cb commit];
[cb waitUntilCompleted];
if (cb.status == MTLCommandBufferStatusError) {
fprintf(stderr, "ds4: Metal model warmup failed: %s\n",
[[cb.error localizedDescription] UTF8String]);
return 0;
}
return 1;
}
static const char *ds4_gpu_mul_mm_id_map0_name(uint32_t ne20) {
switch (ne20) {
case 1: return "kernel_mul_mm_id_map0_ne20_1";
case 2: return "kernel_mul_mm_id_map0_ne20_2";
case 4: return "kernel_mul_mm_id_map0_ne20_4";
case 5: return "kernel_mul_mm_id_map0_ne20_5";
case 6: return "kernel_mul_mm_id_map0_ne20_6";
case 8: return "kernel_mul_mm_id_map0_ne20_8";
case 10: return "kernel_mul_mm_id_map0_ne20_10";
case 16: return "kernel_mul_mm_id_map0_ne20_16";
case 22: return "kernel_mul_mm_id_map0_ne20_22";
default: return NULL;
}
}
static void ds4_gpu_decode_pipeline_fast_cache_reset(void) {
for (uint32_t i = 0; i < DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS; i++) {
ds4_gpu_decode_pipeline_fast_cache_entry *entry =
&g_decode_pipeline_fast_cache[i];
/* ARC must see every strong release; do not memset this table. */
entry->pipeline = nil;
entry->hash = 0;
entry->nsg = 0;
entry->name_len = 0;
entry->used = false;
entry->name[0] = '\0';
}
g_decode_pipeline_fast_lookup_active = false;
g_decode_pipeline_fast_cache_entries = 0;
}
static bool ds4_gpu_decode_pipeline_fast_key(
const char *function_name,
int16_t nsg,
int16_t nxpsg,
uint16_t *name_len_out,
uint64_t *hash_out) {
if (!function_name || !name_len_out || !hash_out) return false;
uint64_t hash = UINT64_C(14695981039346656037);
size_t name_len = 0;
while (function_name[name_len] != '\0') {
if (name_len + 1u >= DS4_METAL_DECODE_PIPELINE_FAST_NAME_BYTES) {
return false;
}
hash ^= (uint8_t)function_name[name_len++];
hash *= UINT64_C(1099511628211);
}
const uint16_t nsg_bits = (uint16_t)nsg;
hash ^= (uint8_t)nsg_bits;
hash *= UINT64_C(1099511628211);
hash ^= (uint8_t)(nsg_bits >> 8u);
hash *= UINT64_C(1099511628211);
const uint16_t nxpsg_bits = (uint16_t)nxpsg;
hash ^= (uint8_t)nxpsg_bits;
hash *= UINT64_C(1099511628211);
hash ^= (uint8_t)(nxpsg_bits >> 8u);
hash *= UINT64_C(1099511628211);
*name_len_out = (uint16_t)name_len;
*hash_out = hash;
return true;
}
static id<MTLComputePipelineState> ds4_gpu_decode_pipeline_fast_cache_lookup(
const char *function_name,
int16_t nsg,
int16_t nxpsg,
uint16_t name_len,
uint64_t hash) {
const uint32_t mask = DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS - 1u;
uint32_t slot = (uint32_t)hash & mask;
for (uint32_t probe = 0;
probe < DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS;
probe++, slot = (slot + 1u) & mask) {
ds4_gpu_decode_pipeline_fast_cache_entry *entry =
&g_decode_pipeline_fast_cache[slot];
if (!entry->used) break;
if (entry->hash == hash &&
entry->nsg == nsg &&
entry->nxpsg == nxpsg &&
entry->name_len == name_len &&
memcmp(entry->name, function_name, name_len) == 0) {
return entry->pipeline;
}
}
return nil;
}
static void ds4_gpu_decode_pipeline_fast_cache_insert(
const char *function_name,
int16_t nsg,
int16_t nxpsg,
uint16_t name_len,
uint64_t hash,
id<MTLComputePipelineState> pipeline) {
if (!pipeline ||
g_decode_pipeline_fast_cache_entries >=
DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS) {
return;
}
const uint32_t mask = DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS - 1u;
uint32_t slot = (uint32_t)hash & mask;
for (uint32_t probe = 0;
probe < DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS;
probe++, slot = (slot + 1u) & mask) {
ds4_gpu_decode_pipeline_fast_cache_entry *entry =
&g_decode_pipeline_fast_cache[slot];
if (entry->used) {
if (entry->hash == hash &&
entry->nsg == nsg &&
entry->nxpsg == nxpsg &&
entry->name_len == name_len &&
memcmp(entry->name, function_name, name_len) == 0) {
entry->pipeline = pipeline;
return;
}
continue;
}
memcpy(entry->name, function_name, name_len);
entry->name[name_len] = '\0';
entry->hash = hash;
entry->nsg = nsg;
entry->nxpsg = nxpsg;
entry->name_len = name_len;
entry->pipeline = pipeline;
entry->used = true;
g_decode_pipeline_fast_cache_entries++;
return;
}
}
int ds4_gpu_set_decode_pipeline_fast_lookup(int enabled) {
const int previous = g_decode_pipeline_fast_lookup_active ? 1 : 0;
g_decode_pipeline_fast_lookup_active = enabled != 0;
return previous;
}
static id<MTLComputePipelineState> ds4_gpu_get_mul_mv_pipeline(
const char *function_name,
int16_t nsg) {
uint16_t fast_name_len = 0;
uint64_t fast_hash = 0;
const bool fast_key_valid =
g_decode_pipeline_fast_lookup_active &&
ds4_gpu_decode_pipeline_fast_key(
function_name, nsg,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
&fast_name_len, &fast_hash);
if (fast_key_valid) {
id<MTLComputePipelineState> fast_cached =
ds4_gpu_decode_pipeline_fast_cache_lookup(
function_name, nsg,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
fast_name_len, fast_hash);
if (fast_cached) return fast_cached;
}
NSString *key = [NSString stringWithFormat:@"%s_nsg=%d", function_name, (int)nsg];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
if (fast_key_valid) {
ds4_gpu_decode_pipeline_fast_cache_insert(
function_name, nsg,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
fast_name_len, fast_hash, cached);
}
return cached;
}
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600];
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal %s function not found: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
if (fast_key_valid) {
ds4_gpu_decode_pipeline_fast_cache_insert(
function_name, nsg,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
fast_name_len, fast_hash, pipeline);
}
return pipeline;
}
/* The ordinary mul-mv cache key covers function name and nsg only. Keep the
* descriptor-hinted PSO separate so a cache hit cannot erase this compiler
* contract or substitute it for the fallback pipeline. */
static id<MTLComputePipelineState> ds4_gpu_new_mul_mv_tg_multiple_pipeline(
const char *function_name,
int16_t nsg) {
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600];
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal %s tg-multiple function not found: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
MTLComputePipelineDescriptor *descriptor = [[MTLComputePipelineDescriptor alloc] init];
descriptor.label = name;
descriptor.computeFunction = fn;
descriptor.threadGroupSizeIsMultipleOfThreadExecutionWidth = YES;
error = nil;
id<MTLComputePipelineState> pipeline =
[g_device newComputePipelineStateWithDescriptor:descriptor
options:MTLPipelineOptionNone
reflection:nil
error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal %s tg-multiple pipeline failed: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
return pipeline;
}
int ds4_gpu_test_decode_pipeline_fast_lookup(void) {
if (!g_initialized || !g_pipeline_cache || !g_library || !g_device) {
return 0;
}
static const char base_name[] = "kernel_mul_mv_id_mxfp4_f32";
char copied_name[sizeof(base_name)];
memcpy(copied_name, base_name, sizeof(base_name));
ds4_gpu_decode_pipeline_fast_cache_reset();
id<MTLComputePipelineState> baseline =
ds4_gpu_get_mul_mv_pipeline(base_name, 2);
if (!baseline ||
g_decode_pipeline_fast_cache_entries != 0u) {
ds4_gpu_decode_pipeline_fast_cache_reset();
return 0;
}
(void)ds4_gpu_set_decode_pipeline_fast_lookup(1);
id<MTLComputePipelineState> populated =
ds4_gpu_get_mul_mv_pipeline(base_name, 2);
id<MTLComputePipelineState> hit =
ds4_gpu_get_mul_mv_pipeline(copied_name, 2);
uint16_t copied_name_len = 0;
uint64_t copied_hash = 0;
const bool copied_key_valid =
ds4_gpu_decode_pipeline_fast_key(
copied_name, 2,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
&copied_name_len, &copied_hash);
id<MTLComputePipelineState> direct_hit =
copied_key_valid ?
ds4_gpu_decode_pipeline_fast_cache_lookup(
copied_name, 2,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
copied_name_len, copied_hash) : nil;
id<MTLComputePipelineState> distinct_name =
ds4_gpu_get_mul_mv_pipeline(
"kernel_mul_mv_id_mxfp4_pair_swiglu_f32", 2);
id<MTLComputePipelineState> distinct_nsg =
ds4_gpu_get_mul_mv_pipeline(base_name, 4);
const bool populated_and_hit =
populated == baseline &&
hit == baseline &&
direct_hit == baseline &&
distinct_name != nil &&
distinct_nsg != nil &&
g_decode_pipeline_fast_cache_entries == 3u;
(void)ds4_gpu_set_decode_pipeline_fast_lookup(0);
const uint32_t entries_before_disabled_lookup =
g_decode_pipeline_fast_cache_entries;
id<MTLComputePipelineState> disabled =
ds4_gpu_get_mul_mv_pipeline(base_name, 2);
const bool inactive_unchanged =
disabled == baseline &&
g_decode_pipeline_fast_cache_entries == entries_before_disabled_lookup;
ds4_gpu_decode_pipeline_fast_cache_reset();
(void)ds4_gpu_set_decode_pipeline_fast_lookup(1);
bool filled = true;
for (uint32_t i = 0;
i < DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS && filled;
i++) {
char synthetic_name[DS4_METAL_DECODE_PIPELINE_FAST_NAME_BYTES];
const int written = snprintf(
synthetic_name, sizeof(synthetic_name),
"ds4_test_decode_pipeline_fast_%u", i);
uint16_t synthetic_name_len = 0;
uint64_t synthetic_hash = 0;
filled = written > 0 &&
(size_t)written < sizeof(synthetic_name) &&
ds4_gpu_decode_pipeline_fast_key(
synthetic_name, 2,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
&synthetic_name_len, &synthetic_hash);
if (filled) {
ds4_gpu_decode_pipeline_fast_cache_insert(
synthetic_name, 2,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
synthetic_name_len,
synthetic_hash, baseline);
}
}
id<MTLComputePipelineState> full_fallback =
ds4_gpu_get_mul_mv_pipeline(base_name, 2);
const bool full_table_unchanged =
filled && full_fallback == baseline &&
g_decode_pipeline_fast_cache_entries ==
DS4_METAL_DECODE_PIPELINE_FAST_CACHE_SLOTS;
ds4_gpu_decode_pipeline_fast_cache_reset();
return populated_and_hit && inactive_unchanged && full_table_unchanged;
}
/*
* Exercises the second SIMD-group constant key shape used by
* ds4_gpu_get_mul_mv_ext_pipeline (name, nsg, nxpsg). Confirms the extended key
* populates and hits, distinguishes a real nxpsg from the single-constant
* sentinel, and coexists with the single-constant mv entries in the shared
* 64-slot table. Uses baseline as a stand-in pipeline object so the guard does
* not depend on compiling a specific extended-kernel function.
*/
int ds4_gpu_test_decode_pipeline_fast_lookup_ext(void) {
if (!g_initialized || !g_pipeline_cache || !g_library || !g_device) {
return 0;
}
static const char base_name[] = "kernel_mul_mv_id_mxfp4_f32";
char copied_name[sizeof(base_name)];
memcpy(copied_name, base_name, sizeof(base_name));
id<MTLComputePipelineState> baseline =
ds4_gpu_get_mul_mv_pipeline(base_name, 2);
if (!baseline) return 0;
ds4_gpu_decode_pipeline_fast_cache_reset();
(void)ds4_gpu_set_decode_pipeline_fast_lookup(1);
const int16_t ext_nsg = 2;
const int16_t ext_nxpsg = 4;
uint16_t ext_name_len = 0;
uint64_t ext_hash = 0;
const bool ext_key_valid =
ds4_gpu_decode_pipeline_fast_key(
copied_name, ext_nsg, ext_nxpsg, &ext_name_len, &ext_hash);
id<MTLComputePipelineState> ext_populated = nil;
if (ext_key_valid) {
ds4_gpu_decode_pipeline_fast_cache_insert(
copied_name, ext_nsg, ext_nxpsg, ext_name_len, ext_hash, baseline);
ext_populated = ds4_gpu_decode_pipeline_fast_cache_lookup(
copied_name, ext_nsg, ext_nxpsg, ext_name_len, ext_hash);
}
/*
* The single-constant sentinel must not match an extended key even when the
* name and nsg agree, proving the two key shapes never alias.
*/
uint16_t none_name_len = 0;
uint64_t none_hash = 0;
const bool none_key_valid =
ds4_gpu_decode_pipeline_fast_key(
copied_name, ext_nsg,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
&none_name_len, &none_hash);
id<MTLComputePipelineState> none_hit =
none_key_valid ?
ds4_gpu_decode_pipeline_fast_cache_lookup(
copied_name, ext_nsg,
DS4_METAL_DECODE_PIPELINE_FAST_NXPSG_NONE,
none_name_len, none_hash) : nil;
const int16_t alt_nxpsg = 8;
uint16_t alt_name_len = 0;
uint64_t alt_hash = 0;
const bool alt_key_valid =
ds4_gpu_decode_pipeline_fast_key(
copied_name, ext_nsg, alt_nxpsg, &alt_name_len, &alt_hash);
id<MTLComputePipelineState> alt_hit =
alt_key_valid ?
ds4_gpu_decode_pipeline_fast_cache_lookup(
copied_name, ext_nsg, alt_nxpsg, alt_name_len, alt_hash) : nil;
const bool ext_populated_and_hit =
ext_key_valid &&
ext_populated == baseline &&
g_decode_pipeline_fast_cache_entries == 1u &&
none_hit == nil &&
alt_hit == nil;
(void)ds4_gpu_set_decode_pipeline_fast_lookup(0);
ds4_gpu_decode_pipeline_fast_cache_reset();
return ext_populated_and_hit;
}
static id<MTLComputePipelineState> ds4_gpu_get_mul_mv_ext_pipeline(
const char *function_name,
int16_t nsg,
int16_t nxpsg) {
uint16_t fast_name_len = 0;
uint64_t fast_hash = 0;
const bool fast_key_valid =
g_decode_pipeline_fast_lookup_active &&
ds4_gpu_decode_pipeline_fast_key(
function_name, nsg, nxpsg, &fast_name_len, &fast_hash);
if (fast_key_valid) {
id<MTLComputePipelineState> fast_cached =
ds4_gpu_decode_pipeline_fast_cache_lookup(
function_name, nsg, nxpsg, fast_name_len, fast_hash);
if (fast_cached) return fast_cached;
}
NSString *key = [NSString stringWithFormat:@"%s_nsg=%d_nxpsg=%d",
function_name, (int)nsg, (int)nxpsg];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
if (fast_key_valid) {
ds4_gpu_decode_pipeline_fast_cache_insert(
function_name, nsg, nxpsg, fast_name_len, fast_hash, cached);
}
return cached;
}
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&nsg type:MTLDataTypeShort atIndex:600];
[constants setConstantValue:&nxpsg type:MTLDataTypeShort atIndex:601];
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal %s function not found: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
if (fast_key_valid) {
ds4_gpu_decode_pipeline_fast_cache_insert(
function_name, nsg, nxpsg, fast_name_len, fast_hash, pipeline);
}
return pipeline;
}
static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_pad_pipeline(
bool has_mask,
int32_t ncpsg) {
/*
* Decode calls this once per layer with identical arguments, so memoize
* the last hit and skip the NSString key + dictionary lookup on the hot
* path. The generic cache below remains the fallback for new variants.
* The rollback switch restores the dictionary path for same-binary A/B.
*/
static struct {
bool m;
int32_t nc;
id<MTLComputePipelineState> pipeline;
} memo;
const bool memo_disabled =
getenv("DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_PAD_BLK_MEMO") != NULL;
if (!memo_disabled && memo.pipeline && memo.m == has_mask && memo.nc == ncpsg) {
return memo.pipeline;
}
NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_pad_mask=%d_ncpsg=%d",
has_mask ? 1 : 0, (int)ncpsg];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
if (!memo_disabled) {
memo = (typeof(memo)){ has_mask, ncpsg, cached };
}
return cached;
}
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:100];
[constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:125];
NSError *error = nil;
id<MTLFunction> fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_pad"
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad function not found: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_pad pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
if (!memo_disabled) {
memo = (typeof(memo)){ has_mask, ncpsg, pipeline };
}
return pipeline;
}
static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_blk_pipeline(
int32_t nqptg,
int32_t ncpsg) {
/*
* Decode calls this once per layer with identical arguments, so memoize
* the last hit and skip the NSString key + dictionary lookup on the hot
* path. The generic cache below remains the fallback for new variants.
* The rollback switch restores the dictionary path for same-binary A/B.
*/
static struct {
int32_t nq;
int32_t nc;
id<MTLComputePipelineState> pipeline;
} memo;
const bool memo_disabled =
getenv("DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_PAD_BLK_MEMO") != NULL;
if (!memo_disabled && memo.pipeline && memo.nq == nqptg && memo.nc == ncpsg) {
return memo.pipeline;
}
NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_blk_nqptg=%d_ncpsg=%d",
(int)nqptg, (int)ncpsg];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
if (!memo_disabled) {
memo = (typeof(memo)){ nqptg, ncpsg, cached };
}
return cached;
}
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&nqptg type:MTLDataTypeInt atIndex:224];
[constants setConstantValue:&ncpsg type:MTLDataTypeInt atIndex:225];
NSError *error = nil;
id<MTLFunction> fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_blk"
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk function not found: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_blk pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
if (!memo_disabled) {
memo = (typeof(memo)){ nqptg, ncpsg, pipeline };
}
return pipeline;
}
static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_pipeline(
const char *function_name,
bool has_mask,
bool has_sinks,
bool has_bias,
bool has_scap,
bool has_kvpad,
bool bc_mask,
int32_t ns10,
int32_t ns20,
int32_t nsg) {
/*
* Prefill and batched decode call this once per layer with identical
* arguments, so memoize the last hit and skip the NSString key + dictionary
* lookup on the hot path. The generic cache below remains the fallback for
* new variants. The rollback switch restores the dictionary path for
* same-binary A/B.
*/
static struct {
const char *fn;
bool m, s, b, c, k, bc;
int32_t n10, n20, sg;
id<MTLComputePipelineState> pipeline;
} memo;
const bool memo_disabled =
getenv("DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_BATCHED_MEMO") != NULL;
if (!memo_disabled && memo.pipeline && memo.fn != NULL &&
strcmp(memo.fn, function_name) == 0 &&
memo.m == has_mask && memo.s == has_sinks && memo.b == has_bias &&
memo.c == has_scap && memo.k == has_kvpad && memo.bc == bc_mask &&
memo.n10 == ns10 && memo.n20 == ns20 && memo.sg == nsg) {
return memo.pipeline;
}
NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d",
function_name,
has_mask ? 1 : 0,
has_sinks ? 1 : 0,
has_bias ? 1 : 0,
has_scap ? 1 : 0,
has_kvpad ? 1 : 0,
bc_mask ? 1 : 0,
(int)ns10,
(int)ns20,
(int)nsg];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
if (!memo_disabled) {
memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias,
has_scap, has_kvpad, bc_mask, ns10, ns20,
nsg, cached };
}
return cached;
}
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:300];
[constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:301];
[constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:302];
[constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:303];
[constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:304];
[constants setConstantValue:&bc_mask type:MTLDataTypeBool atIndex:310];
[constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:320];
[constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:321];
[constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:322];
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal %s function not found: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
if (!memo_disabled) {
memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias,
has_scap, has_kvpad, bc_mask, ns10, ns20,
nsg, pipeline };
}
return pipeline;
}
static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_vec_pipeline(
const char *function_name,
bool has_mask,
bool has_sinks,
bool has_bias,
bool has_scap,
bool has_kvpad,
bool shared_kvpad,
int32_t ns10,
int32_t ns20,
int32_t nsg,
int32_t nwg) {
/*
* Decode calls this once per layer with identical arguments, so memoize
* the last hit and skip the NSString key + dictionary lookup on the hot
* path. The generic cache below remains the fallback for new variants.
*/
static struct {
const char *fn;
bool m, s, b, c, k, sp;
int32_t n10, n20, sg, wg;
id<MTLComputePipelineState> pipeline;
} memo;
if (memo.pipeline && memo.fn != NULL && strcmp(memo.fn, function_name) == 0 &&
memo.m == has_mask && memo.s == has_sinks && memo.b == has_bias &&
memo.c == has_scap && memo.k == has_kvpad && memo.sp == shared_kvpad &&
memo.n10 == ns10 && memo.n20 == ns20 && memo.sg == nsg && memo.wg == nwg) {
return memo.pipeline;
}
NSString *key = [NSString stringWithFormat:@"%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_sharedpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
function_name,
has_mask ? 1 : 0,
has_sinks ? 1 : 0,
has_bias ? 1 : 0,
has_scap ? 1 : 0,
has_kvpad ? 1 : 0,
shared_kvpad ? 1 : 0,
(int)ns10,
(int)ns20,
(int)nsg,
(int)nwg];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias,
has_scap, has_kvpad, shared_kvpad, ns10, ns20,
nsg, nwg, cached };
return cached;
}
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&has_mask type:MTLDataTypeBool atIndex:400];
[constants setConstantValue:&has_sinks type:MTLDataTypeBool atIndex:401];
[constants setConstantValue:&has_bias type:MTLDataTypeBool atIndex:402];
[constants setConstantValue:&has_scap type:MTLDataTypeBool atIndex:403];
[constants setConstantValue:&has_kvpad type:MTLDataTypeBool atIndex:404];
[constants setConstantValue:&shared_kvpad type:MTLDataTypeBool atIndex:405];
[constants setConstantValue:&ns10 type:MTLDataTypeInt atIndex:420];
[constants setConstantValue:&ns20 type:MTLDataTypeInt atIndex:421];
[constants setConstantValue:&nsg type:MTLDataTypeInt atIndex:422];
[constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:423];
NSError *error = nil;
NSString *name = [NSString stringWithUTF8String:function_name];
id<MTLFunction> fn = [g_library newFunctionWithName:name
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal %s function not found: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal %s pipeline failed: %s\n",
function_name, [[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
memo = (typeof(memo)){ function_name, has_mask, has_sinks, has_bias,
has_scap, has_kvpad, shared_kvpad, ns10, ns20,
nsg, nwg, pipeline };
return pipeline;
}
/* Pipeline for the RoPE-fused decode reduce. Same function constants as the
* plain reduce so the split-K geometry is identical. */
static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_reduce_rope_pipeline(
int32_t dv,
int32_t nwg) {
static int32_t memo_dv, memo_nwg;
static id<MTLComputePipelineState> memo_pipeline;
if (memo_pipeline && memo_dv == dv && memo_nwg == nwg) {
return memo_pipeline;
}
NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_vec_reduce_rope_dv=%d_nwg=%d",
(int)dv, (int)nwg];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
memo_dv = dv; memo_nwg = nwg; memo_pipeline = cached;
return cached;
}
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&dv type:MTLDataTypeInt atIndex:500];
[constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:501];
NSError *error = nil;
id<MTLFunction> fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_vec_reduce_rope"
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce_rope not found: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce_rope pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
memo_dv = dv; memo_nwg = nwg; memo_pipeline = pipeline;
return pipeline;
}
int ds4_gpu_decode_attn_rope_fuse_available(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (g_rope_tail_inplace_pair_affine_pipeline == nil) return 0;
if (getenv("DS4_METAL_DISABLE_INPLACE_ROPE_PAIR") != NULL) return 0;
if (getenv("DS4_METAL_DISABLE_AFFINE_ROPE_PAIR") != NULL) return 0;
if (!ds4_gpu_device_name_contains("M3") && !ds4_gpu_device_name_contains("M5")) return 0;
return 1;
}
static id<MTLComputePipelineState> ds4_gpu_get_flash_attn_reduce_pipeline(
int32_t dv,
int32_t nwg) {
/* Same per-layer memo pattern as the vec getter above. */
static int32_t memo_dv, memo_nwg;
static id<MTLComputePipelineState> memo_pipeline;
if (memo_pipeline && memo_dv == dv && memo_nwg == nwg) {
return memo_pipeline;
}
NSString *key = [NSString stringWithFormat:@"kernel_flash_attn_ext_vec_reduce_dv=%d_nwg=%d",
(int)dv, (int)nwg];
id<MTLComputePipelineState> cached = [g_pipeline_cache objectForKey:key];
if (cached) {
memo_dv = dv; memo_nwg = nwg; memo_pipeline = cached;
return cached;
}
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&dv type:MTLDataTypeInt atIndex:500];
[constants setConstantValue:&nwg type:MTLDataTypeInt atIndex:501];
NSError *error = nil;
id<MTLFunction> fn = [g_library newFunctionWithName:@"kernel_flash_attn_ext_vec_reduce"
constantValues:constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce function not found: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
error = nil;
id<MTLComputePipelineState> pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Metal kernel_flash_attn_ext_vec_reduce pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
[g_pipeline_cache setObject:pipeline forKey:key];
memo_dv = dv; memo_nwg = nwg; memo_pipeline = pipeline;
return pipeline;
}
static uint32_t ds4_gpu_flash_attn_vec_nsg(uint32_t n_keys, uint32_t nwg, uint32_t ncpsg) {
uint32_t nsg = 1;
while (2u * nwg * nsg * ncpsg < n_keys && nsg < 4u) {
nsg *= 2u;
}
return nsg;
}
static int ds4_gpu_trace_allocs(void) {
static int initialized;
static int enabled;
if (!initialized) {
enabled = getenv("DS4_METAL_TRACE_ALLOCS") != NULL;
initialized = 1;
}
return enabled;
}
static double ds4_gpu_mib(uint64_t bytes) {
return (double)bytes / (1024.0 * 1024.0);
}
static double ds4_gpu_gib(uint64_t bytes) {
return (double)bytes / (1024.0 * 1024.0 * 1024.0);
}
static ds4_gpu_stream_expert_timing_snapshot
ds4_gpu_stream_expert_timing_current(void) {
return (ds4_gpu_stream_expert_timing_snapshot) {
.selected_calls = g_stream_expert_timing_selected_calls,
.selected_read_ms = g_stream_expert_timing_selected_read_ms,
.selected_sync_ms = g_stream_expert_timing_selected_sync_ms,
.selected_copy_ms = g_stream_expert_timing_selected_copy_ms,
.selected_bind_ms = g_stream_expert_timing_selected_bind_ms,
.split_layers = g_stream_expert_timing_split_layers,
.split_resident_experts = g_stream_expert_timing_split_resident_experts,
.split_missing_experts = g_stream_expert_timing_split_missing_experts,
.split_resident_ms = g_stream_expert_timing_split_resident_ms,
.split_missing_ms = g_stream_expert_timing_split_missing_ms,
.split_missing_load_ms =
g_stream_expert_timing_split_missing_load_ms,
.split_missing_slot_ms =
g_stream_expert_timing_split_missing_slot_ms,
.split_missing_prune_ms =
g_stream_expert_timing_split_missing_prune_ms,
.split_missing_addr_ms =
g_stream_expert_timing_split_missing_addr_ms,
.split_missing_wait_ms =
g_stream_expert_timing_split_missing_wait_ms,
.load_calls = g_stream_expert_timing_load_calls,
.load_prepare_ms = g_stream_expert_timing_load_prepare_ms,
.load_pread_ms = g_stream_expert_timing_load_pread_ms,
.load_install_ms = g_stream_expert_timing_load_install_ms,
.prepare_batch_reuse_calls =
g_stream_expert_timing_prepare_batch_reuse_calls,
.prepare_batch_reuse_ms =
g_stream_expert_timing_prepare_batch_reuse_ms,
.prepare_buffer_calls =
g_stream_expert_timing_prepare_buffer_calls,
.prepare_buffer_ms =
g_stream_expert_timing_prepare_buffer_ms,
.prepare_task_experts =
g_stream_expert_timing_prepare_task_experts,
.prepare_task_ms =
g_stream_expert_timing_prepare_task_ms,
.reuse_scan_calls =
g_stream_expert_timing_reuse_scan_calls,
.reuse_scan_entries =
g_stream_expert_timing_reuse_scan_entries,
.reuse_scan_ms =
g_stream_expert_timing_reuse_scan_ms,
.reuse_clear_ms =
g_stream_expert_timing_reuse_clear_ms,
.readahead_calls =
g_stream_expert_timing_readahead_calls,
.readahead_bytes =
g_stream_expert_timing_readahead_bytes,
.readahead_ms =
g_stream_expert_timing_readahead_ms,
.cache_all_resident_layers =
g_stream_expert_timing_cache_all_resident_layers,
.cache_all_missing_layers =
g_stream_expert_timing_cache_all_missing_layers,
.cache_mixed_layers = g_stream_expert_timing_cache_mixed_layers,
.cache_resident_experts =
g_stream_expert_timing_cache_resident_experts,
.cache_missing_experts =
g_stream_expert_timing_cache_missing_experts,
};
}
static uint64_t ds4_gpu_stream_expert_timing_delta_u64(
uint64_t current,
uint64_t previous) {
return current >= previous ? current - previous : current;
}
static double ds4_gpu_stream_expert_timing_delta_f64(
double current,
double previous) {
return current >= previous ? current - previous : current;
}
static ds4_gpu_stream_expert_timing_snapshot
ds4_gpu_stream_expert_timing_delta(
ds4_gpu_stream_expert_timing_snapshot current,
ds4_gpu_stream_expert_timing_snapshot previous) {
return (ds4_gpu_stream_expert_timing_snapshot) {
.selected_calls =
ds4_gpu_stream_expert_timing_delta_u64(current.selected_calls,
previous.selected_calls),
.selected_read_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.selected_read_ms,
previous.selected_read_ms),
.selected_sync_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.selected_sync_ms,
previous.selected_sync_ms),
.selected_copy_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.selected_copy_ms,
previous.selected_copy_ms),
.selected_bind_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.selected_bind_ms,
previous.selected_bind_ms),
.split_layers =
ds4_gpu_stream_expert_timing_delta_u64(current.split_layers,
previous.split_layers),
.split_resident_experts =
ds4_gpu_stream_expert_timing_delta_u64(current.split_resident_experts,
previous.split_resident_experts),
.split_missing_experts =
ds4_gpu_stream_expert_timing_delta_u64(current.split_missing_experts,
previous.split_missing_experts),
.split_resident_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.split_resident_ms,
previous.split_resident_ms),
.split_missing_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.split_missing_ms,
previous.split_missing_ms),
.split_missing_load_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.split_missing_load_ms,
previous.split_missing_load_ms),
.split_missing_slot_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.split_missing_slot_ms,
previous.split_missing_slot_ms),
.split_missing_prune_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.split_missing_prune_ms,
previous.split_missing_prune_ms),
.split_missing_addr_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.split_missing_addr_ms,
previous.split_missing_addr_ms),
.split_missing_wait_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.split_missing_wait_ms,
previous.split_missing_wait_ms),
.load_calls =
ds4_gpu_stream_expert_timing_delta_u64(current.load_calls,
previous.load_calls),
.load_prepare_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.load_prepare_ms,
previous.load_prepare_ms),
.load_pread_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.load_pread_ms,
previous.load_pread_ms),
.load_install_ms =
ds4_gpu_stream_expert_timing_delta_f64(current.load_install_ms,
previous.load_install_ms),
.prepare_batch_reuse_calls =
ds4_gpu_stream_expert_timing_delta_u64(
current.prepare_batch_reuse_calls,
previous.prepare_batch_reuse_calls),
.prepare_batch_reuse_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.prepare_batch_reuse_ms,
previous.prepare_batch_reuse_ms),
.prepare_buffer_calls =
ds4_gpu_stream_expert_timing_delta_u64(
current.prepare_buffer_calls,
previous.prepare_buffer_calls),
.prepare_buffer_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.prepare_buffer_ms,
previous.prepare_buffer_ms),
.prepare_task_experts =
ds4_gpu_stream_expert_timing_delta_u64(
current.prepare_task_experts,
previous.prepare_task_experts),
.prepare_task_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.prepare_task_ms,
previous.prepare_task_ms),
.reuse_scan_calls =
ds4_gpu_stream_expert_timing_delta_u64(
current.reuse_scan_calls,
previous.reuse_scan_calls),
.reuse_scan_entries =
ds4_gpu_stream_expert_timing_delta_u64(
current.reuse_scan_entries,
previous.reuse_scan_entries),
.reuse_scan_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.reuse_scan_ms,
previous.reuse_scan_ms),
.reuse_clear_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.reuse_clear_ms,
previous.reuse_clear_ms),
.readahead_calls =
ds4_gpu_stream_expert_timing_delta_u64(
current.readahead_calls,
previous.readahead_calls),
.readahead_bytes =
ds4_gpu_stream_expert_timing_delta_u64(
current.readahead_bytes,
previous.readahead_bytes),
.readahead_ms =
ds4_gpu_stream_expert_timing_delta_f64(
current.readahead_ms,
previous.readahead_ms),
.cache_all_resident_layers =
ds4_gpu_stream_expert_timing_delta_u64(
current.cache_all_resident_layers,
previous.cache_all_resident_layers),
.cache_all_missing_layers =
ds4_gpu_stream_expert_timing_delta_u64(
current.cache_all_missing_layers,
previous.cache_all_missing_layers),
.cache_mixed_layers =
ds4_gpu_stream_expert_timing_delta_u64(current.cache_mixed_layers,
previous.cache_mixed_layers),
.cache_resident_experts =
ds4_gpu_stream_expert_timing_delta_u64(
current.cache_resident_experts,
previous.cache_resident_experts),
.cache_missing_experts =
ds4_gpu_stream_expert_timing_delta_u64(
current.cache_missing_experts,
previous.cache_missing_experts),
};
}
static int ds4_gpu_stream_expert_timing_has_data(
ds4_gpu_stream_expert_timing_snapshot s) {
return s.selected_calls != 0 ||
s.split_layers != 0 ||
s.load_calls != 0 ||
s.cache_all_resident_layers != 0 ||
s.cache_all_missing_layers != 0 ||
s.cache_mixed_layers != 0;
}
static void ds4_gpu_stream_expert_timing_print(
const char *scope,
ds4_gpu_stream_expert_timing_snapshot s) {
if (!ds4_gpu_stream_expert_timing_has_data(s)) return;
const double selected_calls = (double)s.selected_calls;
const double split_layers = (double)s.split_layers;
const double selected_read_avg =
selected_calls != 0.0 ? s.selected_read_ms / selected_calls : 0.0;
const double selected_sync_avg =
selected_calls != 0.0 ? s.selected_sync_ms / selected_calls : 0.0;
const double selected_copy_avg =
selected_calls != 0.0 ? s.selected_copy_ms / selected_calls : 0.0;
const double selected_bind_avg =
selected_calls != 0.0 ? s.selected_bind_ms / selected_calls : 0.0;
const double split_resident_avg =
split_layers != 0.0 ? s.split_resident_ms / split_layers : 0.0;
const double split_missing_avg =
split_layers != 0.0 ? s.split_missing_ms / split_layers : 0.0;
const double split_missing_load_avg =
split_layers != 0.0 ?
s.split_missing_load_ms / split_layers : 0.0;
const double split_missing_slot_avg =
split_layers != 0.0 ?
s.split_missing_slot_ms / split_layers : 0.0;
const double split_missing_prune_avg =
split_layers != 0.0 ?
s.split_missing_prune_ms / split_layers : 0.0;
const double split_missing_addr_avg =
split_layers != 0.0 ?
s.split_missing_addr_ms / split_layers : 0.0;
const double split_missing_wait_avg =
split_layers != 0.0 ?
s.split_missing_wait_ms / split_layers : 0.0;
const double load_calls = (double)s.load_calls;
const double load_prepare_avg =
load_calls != 0.0 ? s.load_prepare_ms / load_calls : 0.0;
const double load_pread_avg =
load_calls != 0.0 ? s.load_pread_ms / load_calls : 0.0;
const double load_install_avg =
load_calls != 0.0 ? s.load_install_ms / load_calls : 0.0;
const double prepare_batch_reuse_avg =
s.prepare_batch_reuse_calls != 0 ?
s.prepare_batch_reuse_ms /
(double)s.prepare_batch_reuse_calls : 0.0;
const double prepare_buffer_avg =
s.prepare_buffer_calls != 0 ?
s.prepare_buffer_ms / (double)s.prepare_buffer_calls : 0.0;
const double prepare_task_avg =
s.prepare_task_experts != 0 ?
s.prepare_task_ms / (double)s.prepare_task_experts : 0.0;
const double reuse_scan_avg =
s.reuse_scan_calls != 0 ?
s.reuse_scan_ms / (double)s.reuse_scan_calls : 0.0;
const double reuse_scan_entries_avg =
s.reuse_scan_calls != 0 ?
(double)s.reuse_scan_entries / (double)s.reuse_scan_calls : 0.0;
const double readahead_avg =
s.readahead_calls != 0 ?
s.readahead_ms / (double)s.readahead_calls : 0.0;
const double split_resident_experts_avg =
split_layers != 0.0 ?
(double)s.split_resident_experts / split_layers : 0.0;
const double split_missing_experts_avg =
split_layers != 0.0 ?
(double)s.split_missing_experts / split_layers : 0.0;
const uint64_t cache_layers =
s.cache_all_resident_layers +
s.cache_all_missing_layers +
s.cache_mixed_layers;
const double cache_layer_count = (double)cache_layers;
const double cache_resident_experts_avg =
cache_layer_count != 0.0 ?
(double)s.cache_resident_experts / cache_layer_count : 0.0;
const double cache_missing_experts_avg =
cache_layer_count != 0.0 ?
(double)s.cache_missing_experts / cache_layer_count : 0.0;
fprintf(stderr,
"ds4: streaming expert timing %s selected_calls=%llu read_avg=%.3f ms sync_avg=%.3f ms copy_avg=%.3f ms bind_avg=%.3f ms read_total=%.3f ms sync_total=%.3f ms copy_total=%.3f ms bind_total=%.3f ms split_layers=%llu resident_experts_avg=%.2f missing_experts_avg=%.2f resident_submit_avg=%.3f ms missing_bind_avg=%.3f ms resident_submit_total=%.3f ms missing_bind_total=%.3f ms missing_load_avg=%.3f ms missing_slot_avg=%.3f ms missing_prune_avg=%.3f ms missing_addr_avg=%.3f ms missing_wait_avg=%.3f ms missing_wait_total=%.3f ms load_calls=%llu load_prepare_avg=%.3f ms load_pread_avg=%.3f ms load_install_avg=%.3f ms prepare_batch_reuse_calls=%llu prepare_batch_reuse_avg=%.3f ms prepare_batch_reuse_total=%.3f ms prepare_buffer_calls=%llu prepare_buffer_avg=%.3f ms prepare_buffer_total=%.3f ms prepare_task_experts=%llu prepare_task_avg=%.3f ms prepare_task_total=%.3f ms reuse_scan_calls=%llu reuse_scan_entries_avg=%.1f reuse_scan_avg=%.3f ms reuse_scan_total=%.3f ms reuse_clear_total=%.3f ms readahead_calls=%llu readahead_avg=%.3f ms readahead_total=%.3f ms readahead_gib=%.2f cache_all_resident=%llu cache_all_missing=%llu cache_mixed=%llu cache_resident_avg=%.2f cache_missing_avg=%.2f\n",
scope ? scope : "total",
(unsigned long long)s.selected_calls,
selected_read_avg,
selected_sync_avg,
selected_copy_avg,
selected_bind_avg,
s.selected_read_ms,
s.selected_sync_ms,
s.selected_copy_ms,
s.selected_bind_ms,
(unsigned long long)s.split_layers,
split_resident_experts_avg,
split_missing_experts_avg,
split_resident_avg,
split_missing_avg,
s.split_resident_ms,
s.split_missing_ms,
split_missing_load_avg,
split_missing_slot_avg,
split_missing_prune_avg,
split_missing_addr_avg,
split_missing_wait_avg,
s.split_missing_wait_ms,
(unsigned long long)s.load_calls,
load_prepare_avg,
load_pread_avg,
load_install_avg,
(unsigned long long)s.prepare_batch_reuse_calls,
prepare_batch_reuse_avg,
s.prepare_batch_reuse_ms,
(unsigned long long)s.prepare_buffer_calls,
prepare_buffer_avg,
s.prepare_buffer_ms,
(unsigned long long)s.prepare_task_experts,
prepare_task_avg,
s.prepare_task_ms,
(unsigned long long)s.reuse_scan_calls,
reuse_scan_entries_avg,
reuse_scan_avg,
s.reuse_scan_ms,
s.reuse_clear_ms,
(unsigned long long)s.readahead_calls,
readahead_avg,
s.readahead_ms,
ds4_gpu_gib(s.readahead_bytes),
(unsigned long long)s.cache_all_resident_layers,
(unsigned long long)s.cache_all_missing_layers,
(unsigned long long)s.cache_mixed_layers,
cache_resident_experts_avg,
cache_missing_experts_avg);
}
static void ds4_gpu_print_task_memory_report(void) {
task_vm_info_data_t info;
mach_msg_type_number_t count = TASK_VM_INFO_COUNT;
const kern_return_t kr = task_info(mach_task_self(),
TASK_VM_INFO,
(task_info_t)&info,
&count);
if (kr != KERN_SUCCESS) return;
fprintf(stderr,
"ds4: macOS task memory footprint %.2f GiB, resident %.2f GiB, virtual %.2f GiB\n",
ds4_gpu_gib((uint64_t)info.phys_footprint),
ds4_gpu_gib((uint64_t)info.resident_size),
ds4_gpu_gib((uint64_t)info.virtual_size));
}
void ds4_gpu_print_memory_report(const char *label) {
uint64_t cached_prefill_mask_bytes = 0;
uint64_t cached_prefill_blk_bytes = 0;
for (uint32_t i = 0; i < DS4_GPU_PREFILL_MASK_CACHE_SLOTS; i++) {
const ds4_gpu_zero_prefix_prefill_mask_cache_entry *entry =
&g_zero_prefix_prefill_mask_cache[i];
if (entry->mask) cached_prefill_mask_bytes += entry->mask_bytes;
if (entry->blk) cached_prefill_blk_bytes += entry->blk_bytes;
}
const uint64_t scratch =
(uint64_t)g_flash_attn_mask_bytes +
(uint64_t)g_flash_attn_zero_mask_bytes +
cached_prefill_mask_bytes +
(uint64_t)g_flash_attn_pad_bytes +
(uint64_t)g_flash_attn_tmp_bytes +
(uint64_t)g_flash_attn_blk_bytes +
cached_prefill_blk_bytes +
(uint64_t)g_flash_attn_ring_bytes +
(uint64_t)g_flash_attn_kv_bytes +
(uint64_t)g_glm_flash_attn_mask_bytes +
(uint64_t)g_compressor_pool_kv_bytes +
(uint64_t)g_compressor_pool_score_bytes +
(uint64_t)g_compressor_pool_score_cont_bytes +
(uint64_t)g_compressor_pool_softmax_bytes +
(uint64_t)g_compressor_pool_product_bytes +
(uint64_t)g_compressor_store_ape_bytes +
(uint64_t)g_compressor_store_score_bytes +
(uint64_t)g_embed_rows_bytes +
(uint64_t)g_router_selection_bytes +
(uint64_t)g_router_weight_sum_bytes +
(uint64_t)g_indexer_head_scores_bytes +
(uint64_t)g_indexer_topk_bytes +
(uint64_t)g_indexed_topk_bytes +
(uint64_t)g_f16_round_scratch_bytes +
(uint64_t)g_q8_prefill_scratch_bytes + (uint64_t)g_qwen4_conv_halo_bytes +
(uint64_t)g_raw_store_round_bytes +
(uint64_t)g_moe_gate_scratch_bytes +
(uint64_t)g_moe_down_scratch_bytes +
(uint64_t)g_moe_id_map_bytes + (uint64_t)g_moe_packed_rhs_bytes +
(uint64_t)g_moe_q4_gate_slots_bytes +
(uint64_t)g_moe_q4_up_slots_bytes +
(uint64_t)g_moe_q4_down_slots_bytes;
pthread_mutex_lock(&g_tensor_mu);
const uint64_t tensor_live_snap = g_tensor_alloc_live_bytes;
const uint64_t tensor_peak_snap = g_tensor_alloc_peak_bytes;
pthread_mutex_unlock(&g_tensor_mu);
uint64_t tracked_live = tensor_live_snap;
if (tracked_live > UINT64_MAX - g_stream_expert_cache_bytes) {
tracked_live = UINT64_MAX;
} else {
tracked_live += g_stream_expert_cache_bytes;
}
const bool color = ds4_log_is_tty(stderr);
const char *green = color ? "\x1b[32m" : "";
const char *bright_green = color ? "\x1b[1;32m" : "";
const char *reset = color ? "\x1b[0m" : "";
fprintf(stderr,
"%sds4: Metal memory%s%s: runtime %.2f GiB + streaming experts %.2f GiB = %s%.2f GiB tracked live%s\n",
green,
label && label[0] ? " " : "",
label && label[0] ? label : "",
ds4_gpu_gib(tensor_live_snap),
ds4_gpu_gib(g_stream_expert_cache_bytes),
bright_green,
ds4_gpu_gib(tracked_live),
reset);
if (color) fputs(green, stderr);
fprintf(stderr,
"ds4: runtime tensors live %.2f MiB peak %.2f MiB\n",
ds4_gpu_mib(tensor_live_snap),
ds4_gpu_mib(tensor_peak_snap));
ds4_gpu_print_task_memory_report();
fprintf(stderr,
"ds4: mmap model wrapper creations %llu buffers %.2f GiB cumulative, %.2f GiB max (not copied)\n",
(unsigned long long)g_model_wrap_count,
ds4_gpu_gib(g_model_wrap_bytes),
ds4_gpu_gib(g_model_wrap_max_bytes));
if (g_model_buffer_cache && [g_model_buffer_cache count] != 0) {
const uint64_t limit = ds4_gpu_exact_view_cache_limit_bytes();
if (limit == 0) {
fprintf(stderr,
"ds4: exact model view cache %lu buffers %.2f GiB unlimited, %llu evictions (not copied)\n",
(unsigned long)[g_model_buffer_cache count],
ds4_gpu_gib(g_model_buffer_cache_bytes),
(unsigned long long)g_model_buffer_cache_evictions);
} else {
fprintf(stderr,
"ds4: exact model view cache %lu buffers %.2f GiB / %.2f GiB, %llu evictions (not copied)\n",
(unsigned long)[g_model_buffer_cache count],
ds4_gpu_gib(g_model_buffer_cache_bytes),
ds4_gpu_gib(limit),
(unsigned long long)g_model_buffer_cache_evictions);
}
}
if (g_stream_expert_cache_hits != 0 ||
g_stream_expert_cache_misses != 0 ||
g_stream_expert_cache_bytes != 0) {
const uint64_t budget = ds4_gpu_stream_expert_cache_configured_budget();
uint64_t target_bytes = 0;
if (budget != 0 && g_stream_expert_cache_expert_bytes != 0) {
target_bytes =
budget > UINT64_MAX / g_stream_expert_cache_expert_bytes ?
UINT64_MAX :
budget * g_stream_expert_cache_expert_bytes;
}
const uint64_t lookups = g_stream_expert_cache_hits + g_stream_expert_cache_misses;
const double hit_rate = lookups ?
(double)g_stream_expert_cache_hits / (double)lookups : 0.0;
if (g_stream_expert_cache_evict_advise_bytes != 0 ||
g_stream_expert_cache_willneed_advise_bytes != 0 ||
g_stream_expert_cache_pread_bytes != 0) {
fprintf(stderr,
"ds4: streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu evict_dontneed=%.2f GiB miss_willneed=%.2f GiB miss_pread=%.2f GiB pread_ms=%.3f\n",
(unsigned long long)budget,
g_stream_expert_cache_entry_count,
ds4_gpu_mib(g_stream_expert_cache_expert_bytes),
ds4_gpu_gib(target_bytes),
ds4_gpu_gib(g_stream_expert_cache_bytes),
(unsigned long long)g_stream_expert_cache_hits,
(unsigned long long)g_stream_expert_cache_misses,
hit_rate,
(unsigned long long)g_stream_expert_cache_wraps,
(unsigned long long)g_stream_expert_cache_evictions,
(unsigned long long)g_stream_expert_cache_buffer_allocs,
(unsigned long long)g_stream_expert_cache_buffer_reuses,
ds4_gpu_gib(g_stream_expert_cache_evict_advise_bytes),
ds4_gpu_gib(g_stream_expert_cache_willneed_advise_bytes),
ds4_gpu_gib(g_stream_expert_cache_pread_bytes),
g_stream_expert_cache_pread_ms);
} else {
fprintf(stderr,
"ds4: streaming expert cache budget=%llu experts entries=%u expert=%.2f MiB target=%.2f GiB live=%.2f GiB, hits=%llu misses=%llu hit_rate=%.3f wraps=%llu evictions=%llu buffer_allocs=%llu buffer_reuses=%llu\n",
(unsigned long long)budget,
g_stream_expert_cache_entry_count,
ds4_gpu_mib(g_stream_expert_cache_expert_bytes),
ds4_gpu_gib(target_bytes),
ds4_gpu_gib(g_stream_expert_cache_bytes),
(unsigned long long)g_stream_expert_cache_hits,
(unsigned long long)g_stream_expert_cache_misses,
hit_rate,
(unsigned long long)g_stream_expert_cache_wraps,
(unsigned long long)g_stream_expert_cache_evictions,
(unsigned long long)g_stream_expert_cache_buffer_allocs,
(unsigned long long)g_stream_expert_cache_buffer_reuses);
}
if (g_stream_expert_cache_mlock_bytes != 0 ||
g_stream_expert_cache_mlock_failures != 0) {
fprintf(stderr,
"ds4: streaming expert buffer mlock locked=%.2f GiB failed=%.2f GiB failures=%llu time=%.3f ms\n",
ds4_gpu_gib(g_stream_expert_cache_mlock_bytes),
ds4_gpu_gib(g_stream_expert_cache_mlock_fail_bytes),
(unsigned long long)g_stream_expert_cache_mlock_failures,
g_stream_expert_cache_mlock_ms);
}
if (ds4_gpu_stream_expert_timing_summary_enabled()) {
fprintf(stderr,
"ds4: streaming cache aging tokens=%llu first_layer=%u decay_token=%llu\n",
(unsigned long long)g_stream_expert_cache_decode_tokens,
g_stream_expert_cache_first_decode_layer,
(unsigned long long)g_stream_expert_cache_hotness_decay_token);
const ds4_gpu_stream_expert_timing_snapshot total =
ds4_gpu_stream_expert_timing_current();
if (ds4_gpu_stream_expert_timing_has_data(total)) {
const ds4_gpu_stream_expert_timing_snapshot delta =
ds4_gpu_stream_expert_timing_delta(
total,
g_stream_expert_timing_last_report);
ds4_gpu_stream_expert_timing_print("total", total);
ds4_gpu_stream_expert_timing_print("delta", delta);
g_stream_expert_timing_last_report = total;
}
}
if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS") != NULL) {
fprintf(stderr, "ds4: streaming expert cache per-layer stats:\n");
for (uint32_t layer = 0;
layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
layer++) {
const uint64_t hits = g_stream_expert_cache_layer_hits[layer];
const uint64_t misses = g_stream_expert_cache_layer_misses[layer];
const uint64_t lookups = hits + misses;
const uint64_t evictions =
g_stream_expert_cache_layer_evictions[layer];
const uint64_t pread_bytes =
g_stream_expert_cache_layer_pread_bytes[layer];
const double pread_ms =
g_stream_expert_cache_layer_pread_ms[layer];
const uint32_t cached =
g_stream_expert_cache_layer_count[layer];
const uint32_t layer_slots =
ds4_gpu_stream_expert_cache_configured_count() != 0 ?
DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT : 0;
if (lookups == 0 && evictions == 0 && pread_bytes == 0 &&
cached == 0) {
continue;
}
const double layer_hit_rate = lookups ?
(double)hits / (double)lookups : 0.0;
fprintf(stderr,
"ds4: layer=%u layer_slots=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n",
layer,
layer_slots,
cached,
(unsigned long long)hits,
(unsigned long long)misses,
layer_hit_rate,
(unsigned long long)evictions,
ds4_gpu_gib(pread_bytes),
pread_ms);
}
if (getenv("DS4_METAL_STREAMING_EXPERT_LAYER_STATS_DELTA") != NULL) {
fprintf(stderr, "ds4: streaming expert cache per-layer delta:\n");
for (uint32_t layer = 0;
layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
layer++) {
const uint64_t hits =
g_stream_expert_cache_layer_hits[layer];
const uint64_t misses =
g_stream_expert_cache_layer_misses[layer];
const uint64_t evictions =
g_stream_expert_cache_layer_evictions[layer];
const uint64_t pread_bytes =
g_stream_expert_cache_layer_pread_bytes[layer];
const double pread_ms =
g_stream_expert_cache_layer_pread_ms[layer];
const uint64_t delta_hits =
hits - g_stream_expert_cache_layer_last_hits[layer];
const uint64_t delta_misses =
misses - g_stream_expert_cache_layer_last_misses[layer];
const uint64_t delta_evictions =
evictions -
g_stream_expert_cache_layer_last_evictions[layer];
const uint64_t delta_pread_bytes =
pread_bytes -
g_stream_expert_cache_layer_last_pread_bytes[layer];
const double delta_pread_ms =
pread_ms -
g_stream_expert_cache_layer_last_pread_ms[layer];
const uint64_t lookups = delta_hits + delta_misses;
if (lookups == 0 && delta_evictions == 0 &&
delta_pread_bytes == 0) {
continue;
}
const double hit_rate = lookups ?
(double)delta_hits / (double)lookups : 0.0;
fprintf(stderr,
"ds4: layer=%u cached=%u hits=%llu misses=%llu hit_rate=%.3f evictions=%llu miss_pread=%.2f GiB pread_ms=%.3f\n",
layer,
g_stream_expert_cache_layer_count[layer],
(unsigned long long)delta_hits,
(unsigned long long)delta_misses,
hit_rate,
(unsigned long long)delta_evictions,
ds4_gpu_gib(delta_pread_bytes),
delta_pread_ms);
g_stream_expert_cache_layer_last_hits[layer] = hits;
g_stream_expert_cache_layer_last_misses[layer] = misses;
g_stream_expert_cache_layer_last_evictions[layer] =
evictions;
g_stream_expert_cache_layer_last_pread_bytes[layer] =
pread_bytes;
g_stream_expert_cache_layer_last_pread_ms[layer] =
pread_ms;
}
}
}
}
fprintf(stderr,
"ds4: model residency requests %llu%s\n",
(unsigned long long)g_model_residency_count,
g_ssd_streaming_mode ? " (ssd-streaming)" :
(getenv("DS4_METAL_NO_RESIDENCY") != NULL ? " (disabled)" : ""));
fprintf(stderr,
"ds4: device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n",
g_metal_device_name[0] ? g_metal_device_name : "(unknown)",
g_metal4_runtime_available ? "yes" : "no",
g_metal4_family_supported ? "yes" : "no",
g_metal4_queue_supported ? "yes" : "no",
g_metal4_tensor_api_enabled ? "enabled" :
(g_metal4_tensor_api_compile_supported ? "available" : "disabled"),
g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected");
fprintf(stderr,
"ds4: accelerated Metal path %s%s\n",
ds4_gpu_mpp_available() ? "enabled" : "disabled",
g_quality_mode ? " by --quality" :
(!g_metal4_tensor_api_enabled ? " (tensor API unavailable)" : ""));
fprintf(stderr,
"ds4: device %s, Metal 4 runtime %s, family %s, MTL4 queue %s, tensor API %s, M5 neural accelerators %s\n",
g_metal_device_name[0] ? g_metal_device_name : "(unknown)",
g_metal4_runtime_available ? "yes" : "no",
g_metal4_family_supported ? "yes" : "no",
g_metal4_queue_supported ? "yes" : "no",
g_metal4_tensor_api_enabled ? "enabled" :
(g_metal4_tensor_api_compile_supported ? "available" : "disabled"),
g_metal4_m5_neural_accelerators_hint ? "likely" : "not detected");
fprintf(stderr,
"ds4: scratch %.2f MiB (flash mask %.2f, pad %.2f, tmp %.2f, blk %.2f, ring %.2f, kv %.2f, compressor %.2f, router %.2f, indexer %.2f, moe %.2f, f16 %.2f, raw-store %.2f)\n",
ds4_gpu_mib(scratch),
ds4_gpu_mib((uint64_t)g_flash_attn_mask_bytes +
(uint64_t)g_glm_flash_attn_mask_bytes +
(uint64_t)g_flash_attn_zero_mask_bytes +
cached_prefill_mask_bytes),
ds4_gpu_mib((uint64_t)g_flash_attn_pad_bytes),
ds4_gpu_mib((uint64_t)g_flash_attn_tmp_bytes),
ds4_gpu_mib((uint64_t)g_flash_attn_blk_bytes +
cached_prefill_blk_bytes),
ds4_gpu_mib((uint64_t)g_flash_attn_ring_bytes),
ds4_gpu_mib((uint64_t)g_flash_attn_kv_bytes),
ds4_gpu_mib((uint64_t)g_compressor_pool_kv_bytes +
(uint64_t)g_compressor_pool_score_bytes +
(uint64_t)g_compressor_pool_score_cont_bytes +
(uint64_t)g_compressor_pool_softmax_bytes +
(uint64_t)g_compressor_pool_product_bytes +
(uint64_t)g_compressor_store_ape_bytes +
(uint64_t)g_compressor_store_score_bytes +
(uint64_t)g_embed_rows_bytes),
ds4_gpu_mib((uint64_t)g_router_selection_bytes +
(uint64_t)g_router_weight_sum_bytes),
ds4_gpu_mib((uint64_t)g_indexer_head_scores_bytes +
(uint64_t)g_indexer_topk_bytes +
(uint64_t)g_indexed_topk_bytes),
ds4_gpu_mib((uint64_t)g_moe_gate_scratch_bytes +
(uint64_t)g_moe_down_scratch_bytes +
(uint64_t)g_moe_id_map_bytes + (uint64_t)g_moe_packed_rhs_bytes +
(uint64_t)g_moe_q4_gate_slots_bytes +
(uint64_t)g_moe_q4_up_slots_bytes +
(uint64_t)g_moe_q4_down_slots_bytes),
ds4_gpu_mib((uint64_t)g_f16_round_scratch_bytes),
ds4_gpu_mib((uint64_t)g_raw_store_round_bytes));
if (color) fputs(reset, stderr);
}
void ds4_gpu_set_quality(bool quality) {
g_quality_mode = quality ? 1 : 0;
}
void ds4_gpu_set_glm_model(bool enabled) {
g_glm_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);
if (g_ssd_streaming_mode) {
fprintf(stderr,
"ds4: Metal SSD streaming mode enabled; full model residency and warmup are skipped\n");
}
}
void ds4_gpu_set_glm_streaming_prefill_full_layer(bool enabled) {
g_glm_streaming_prefill_full_layer_runtime = enabled ? 1 : 0;
}
void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) {
if (experts > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {
experts = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES;
}
g_stream_expert_cache_budget_override = experts;
ds4_gpu_stream_expert_cache_clear_all(1);
}
void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) {
/*
* Pre-seed the cache's single slab size class with the model's uniform
* per-expert bytes (first routed layer). With a mixed-precision GGUF this
* pins the class to the majority layers so the boosted ones are rejected
* deterministically from startup, instead of depending on which layer
* happens to touch the cache first.
*/
g_stream_expert_cache_expert_bytes = bytes;
}
uint64_t ds4_gpu_recommended_working_set_size(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!g_device) return 0;
return (uint64_t)[g_device recommendedMaxWorkingSetSize];
}
static int ds4_gpu_model_map_log_enabled(void) {
if (!g_ssd_streaming_mode) return 1;
const char *trace = getenv("DS4_METAL_STREAMING_MAP_TRACE");
return trace && trace[0] && strcmp(trace, "0") != 0;
}
static id<MTLBuffer> ds4_gpu_wrap_model_range(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t len,
uint64_t *inner_offset);
static id<MTLBuffer> ds4_gpu_wrap_model_exact_range(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t len,
uint64_t *inner_offset);
static const char *ds4_gpu_source =
"#include <metal_stdlib>\n"
"#ifdef DS4_METAL_HAS_TENSOR\n"
"#include <metal_tensor>\n"
"#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>\n"
"#endif\n"
"using namespace metal;\n"
"#ifdef DS4_METAL_HAS_TENSOR\n"
"using namespace mpp::tensor_ops;\n"
"#endif\n"
"\n"
"#define MAX(x, y) ((x) > (y) ? (x) : (y))\n"
"#define MIN(x, y) ((x) < (y) ? (x) : (y))\n"
"#define SWAP(x, y) { auto tmp = (x); (x) = (y); (y) = tmp; }\n"
"#define QK8_0 32\n"
"#ifndef QK_K\n"
"#define QK_K 256\n"
"#endif\n"
"#define N_SIMDWIDTH 32\n"
"#define N_R0_Q8_0 2\n"
"#define N_SG_Q8_0 4\n"
"#define FC_MUL_MV 600\n"
"#define FC_MUL_MM 700\n"
"#define FC_BIN 1300\n"
"#define FOR_UNROLL(x) _Pragma(\"clang loop unroll(full)\") for (x)\n"
"#define M_PI_F 3.14159265358979323846f\n"
"\n"
"// Reads one byte per stride to warm model-backed pages without copying the\n"
"// model. This is outside inference and exists only to reduce first-use stalls.\n"
"kernel void kernel_touch_u8_stride(\n"
" device const uchar *src [[buffer(0)]],\n"
" device uchar *dst [[buffer(1)]],\n"
" constant ulong &stride [[buffer(2)]],\n"
" constant ulong &bytes [[buffer(3)]],\n"
" constant ulong &dst_offset [[buffer(4)]],\n"
" uint gid [[thread_position_in_grid]]) {\n"
" ulong off = (ulong)gid * stride;\n"
" if (off >= bytes) return;\n"
" dst[dst_offset + (ulong)gid] = src[off];\n"
"}\n"
"\n"
"// Gate-time prefetch: one dispatch over up to 8 ranges given as GPU\n"
"// addresses (bindless views over the model), 128 B per thread, ranges\n"
"// assigned by cumulative line counts. One launch instead of one per\n"
"// range: the fixed cost per dispatch is what a rank pays when its peer's\n"
"// partial is already there.\n"
"kernel void kernel_touch_u8_stride_table(\n"
" constant ulong *addr [[buffer(0)]],\n"
" constant uint *lines [[buffer(1)]],\n"
" constant uint &count [[buffer(2)]],\n"
" device uchar *dst [[buffer(3)]],\n"
" uint gid [[thread_position_in_grid]]) {\n"
" uint line = gid;\n"
" uint i = 0;\n"
" while (i < count && line >= lines[i]) { line -= lines[i]; i++; }\n"
" if (i >= count) return;\n"
" device const uchar *src = (device const uchar *)addr[i];\n"
" dst[gid] = src[(ulong)line * 128ul];\n"
"}\n"
"\n"
"enum ds4_sort_order {\n"
" DS4_SORT_ORDER_ASC,\n"
" DS4_SORT_ORDER_DESC,\n"
"};\n"
"\n"
"struct block_q8_0 {\n"
" half d;\n"
" int8_t qs[QK8_0];\n"
"};\n"
"\n"
"struct block_q8_K {\n"
" float d;\n"
" int8_t qs[QK_K];\n"
" int16_t bsums[QK_K / 16];\n"
"};\n"
"\n"
"\n";
static NSString *ds4_gpu_full_source(void) {
NSString *base = [NSString stringWithUTF8String:ds4_gpu_source];
NSFileManager *fm = [NSFileManager defaultManager];
/*
* Kernels are kept as separate files for review, then concatenated into one
* Metal library. Environment overrides are still honored so a diagnostic
* run can swap one source file without changing the executable.
*/
NSArray<NSArray<NSString *> *> *required_sources = @[
@[@"DS4_METAL_FLASH_ATTN_SOURCE", @"metal/flash_attn.metal"],
@[@"DS4_METAL_DENSE_SOURCE", @"metal/dense.metal"],
@[@"DS4_METAL_GLM53_BF16_SOURCE", @"metal/glm53_bf16.metal"],
@[@"DS4_METAL_GLM53_VISION_SOURCE", @"metal/glm53_vision.metal"],
@[@"DS4_METAL_DEEPSEEK4_VISION_SOURCE", @"metal/deepseek4_vision.metal"],
@[@"DS4_METAL_GLM53_KDA_SOURCE", @"metal/glm53_kda.metal"],
@[@"DS4_METAL_MOE_SOURCE", @"metal/moe.metal"],
@[@"DS4_METAL_DSV4_HC_SOURCE", @"metal/dsv4_hc.metal"],
@[@"DS4_METAL_UNARY_SOURCE", @"metal/unary.metal"],
@[@"DS4_METAL_DSV4_KV_SOURCE", @"metal/dsv4_kv.metal"],
@[@"DS4_METAL_DSV41_SOURCE", @"metal/dsv41.metal"],
@[@"DS4_METAL_DSV4_ROPE_SOURCE", @"metal/dsv4_rope.metal"],
@[@"DS4_METAL_DSV4_MISC_SOURCE", @"metal/dsv4_misc.metal"],
@[@"DS4_METAL_ARGSORT_SOURCE", @"metal/argsort.metal"],
@[@"DS4_METAL_CPY_SOURCE", @"metal/cpy.metal"],
@[@"DS4_METAL_CONCAT_SOURCE", @"metal/concat.metal"],
@[@"DS4_METAL_GET_ROWS_SOURCE", @"metal/get_rows.metal"],
@[@"DS4_METAL_SUM_ROWS_SOURCE", @"metal/sum_rows.metal"],
@[@"DS4_METAL_SOFTMAX_SOURCE", @"metal/softmax.metal"],
@[@"DS4_METAL_REPEAT_SOURCE", @"metal/repeat.metal"],
@[@"DS4_METAL_GLU_SOURCE", @"metal/glu.metal"],
@[@"DS4_METAL_NORM_SOURCE", @"metal/norm.metal"],
@[@"DS4_METAL_BIN_SOURCE", @"metal/bin.metal"],
@[@"DS4_METAL_SET_ROWS_SOURCE", @"metal/set_rows.metal"],
@[@"DS4_METAL_QWEN4_SOURCE", @"metal/qwen4.metal"],
@[@"DS4_METAL_QWEN4_VISION_SOURCE", @"metal/qwen4_vision.metal"],
];
NSMutableString *source = [NSMutableString stringWithString:base];
for (NSArray<NSString *> *spec in required_sources) {
const char *override_path = getenv([spec[0] UTF8String]);
NSMutableArray<NSString *> *paths = [NSMutableArray array];
if (override_path && override_path[0]) {
[paths addObject:[NSString stringWithUTF8String:override_path]];
}
[paths addObject:spec[1]];
[paths addObject:[@"./" stringByAppendingString:spec[1]]];
NSString *loaded = nil;
NSString *loaded_path = nil;
for (NSString *path in paths) {
if (![fm fileExistsAtPath:path]) continue;
NSError *error = nil;
loaded = [NSString stringWithContentsOfFile:path
encoding:NSUTF8StringEncoding
error:&error];
if (!loaded) {
fprintf(stderr, "ds4: failed to read Metal source %s: %s\n",
[path UTF8String], [[error localizedDescription] UTF8String]);
return nil;
}
loaded_path = path;
break;
}
if (!loaded) {
fprintf(stderr,
"ds4: Metal source %s not found (set %s to override)\n",
[spec[1] UTF8String], [spec[0] UTF8String]);
return nil;
}
[source appendFormat:@"\n// appended %@\n%@\n", loaded_path, loaded];
}
return source;
}
typedef struct {
int32_t ne00t;
int32_t ne00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne10;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
} ds4_gpu_get_rows_args;
typedef struct {
int32_t n_embd;
int32_t n_vocab;
int32_t n_tokens;
uint64_t src_row_bytes;
uint64_t dst_row_bytes;
uint64_t token_stride;
} ds4_gpu_get_rows_q8_0_args;
typedef struct {
int32_t ne00;
int32_t ne01;
int32_t ne02;
int32_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne0;
int32_t ne1;
int32_t ne2;
int32_t ne3;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
} ds4_gpu_repeat_args;
typedef struct {
int32_t nk0;
int32_t ne01;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne11;
int32_t ne12;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
} ds4_gpu_set_rows_args;
typedef struct {
int32_t ne00;
int32_t ne01;
int32_t ne02;
int32_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne10;
int32_t ne11;
int32_t ne12;
int32_t ne13;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
int32_t ne0;
int32_t ne1;
int32_t ne2;
int32_t ne3;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
int32_t dim;
} ds4_gpu_concat_args;
typedef struct {
int64_t nk0;
int64_t ne00;
int64_t ne01;
int64_t ne02;
int64_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int64_t ne0;
int64_t ne1;
int64_t ne2;
int64_t ne3;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
} ds4_gpu_cpy_args;
static ds4_gpu_cpy_args ds4_gpu_make_cpy_1d_args(
uint32_t n,
uint64_t src_elem,
uint64_t dst_elem) {
return (ds4_gpu_cpy_args) {
.nk0 = (int64_t)n,
.ne00 = (int64_t)n,
.ne01 = 1,
.ne02 = 1,
.ne03 = 1,
.nb00 = src_elem,
.nb01 = (uint64_t)n * src_elem,
.nb02 = (uint64_t)n * src_elem,
.nb03 = (uint64_t)n * src_elem,
.ne0 = (int64_t)n,
.ne1 = 1,
.ne2 = 1,
.ne3 = 1,
.nb0 = dst_elem,
.nb1 = (uint64_t)n * dst_elem,
.nb2 = (uint64_t)n * dst_elem,
.nb3 = (uint64_t)n * dst_elem,
};
}
static NSUInteger ds4_gpu_cpy_threads(uint32_t n, id<MTLComputePipelineState> pipeline) {
NSUInteger nth = 32u;
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
while (nth < (NSUInteger)n && nth < max_threads) nth *= 2u;
if (nth > max_threads) nth = max_threads;
if (nth > (NSUInteger)n) nth = (NSUInteger)n;
return nth ? nth : 1u;
}
static float ds4_gpu_negative_infinity(void) {
union { uint32_t u; float f; } v = { 0xff800000u };
return v.f;
}
static float ds4_gpu_positive_infinity(void) {
union { uint32_t u; float f; } v = { 0x7f800000u };
return v.f;
}
static int ds4_gpu_encode_cpy_f32_f32_3d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t cols,
uint32_t rows,
uint32_t planes,
uint64_t src_row_stride,
uint64_t src_plane_stride,
uint64_t dst_row_stride,
uint64_t dst_plane_stride);
static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t cols,
uint32_t rows,
uint32_t planes,
uint64_t src_col_stride,
uint64_t src_row_stride,
uint64_t src_plane_stride,
uint64_t dst_row_stride,
uint64_t dst_plane_stride);
static int ds4_gpu_encode_cpy_f32_f16_1d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t n);
static int ds4_gpu_encode_cpy_f32_f16_2d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t cols,
uint32_t rows,
uint64_t src_row_stride,
uint64_t dst_row_stride);
static int ds4_gpu_encode_cpy_f16_f32_1d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t n);
static int ds4_gpu_encode_fill_f32_rows(
id<MTLCommandBuffer> cb,
id<MTLBuffer> buf,
NSUInteger offset,
uint32_t width,
uint32_t rows,
float value);
static int ds4_gpu_encode_add_f32_1d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> a,
NSUInteger a_off,
id<MTLBuffer> b,
NSUInteger b_off,
id<MTLBuffer> out,
NSUInteger out_off,
uint32_t n);
typedef struct {
int32_t ne00;
uint64_t nb01;
int32_t ne10;
uint64_t nb11;
int32_t ne0;
uint64_t nb1;
int32_t i00;
int32_t i10;
float alpha;
float limit;
} ds4_gpu_glu_args;
typedef struct {
uint32_t n;
} ds4_gpu_add_flat_args;
typedef struct {
int32_t ne00;
int32_t ne01;
int32_t ne02;
int32_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne10;
int32_t ne11;
int32_t ne12;
int32_t ne13;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
int32_t ne0;
int32_t ne1;
int32_t ne2;
int32_t ne3;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
uint64_t offs;
uint64_t o1[8];
} ds4_gpu_bin_args;
typedef struct {
int32_t ne00;
int32_t ne01;
int32_t ne02;
int32_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne0;
int32_t ne1;
int32_t ne2;
int32_t ne3;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
float slope;
float scale;
float bias;
float val;
float min;
float max;
} ds4_gpu_unary_args;
static ds4_gpu_bin_args ds4_gpu_make_bin_rows_args(uint32_t n, uint32_t rows, uint32_t rhs_n) {
const uint64_t row_bytes = (uint64_t)n * sizeof(float);
const uint64_t rhs_row_bytes = (uint64_t)rhs_n * sizeof(float);
return (ds4_gpu_bin_args) {
.ne00 = (int32_t)n,
.ne01 = (int32_t)rows,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = row_bytes,
.nb02 = row_bytes,
.nb03 = row_bytes,
.ne10 = (int32_t)rhs_n,
.ne11 = 1,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = rhs_row_bytes,
.nb12 = rhs_row_bytes,
.nb13 = rhs_row_bytes,
.ne0 = (int32_t)n,
.ne1 = (int32_t)rows,
.ne2 = 1,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = row_bytes,
.nb2 = row_bytes,
.nb3 = row_bytes,
.offs = 0,
.o1 = { 0 },
};
}
static ds4_gpu_unary_args ds4_gpu_make_unary_rows_args(
uint32_t n,
uint32_t rows,
int c4,
float scale,
float bias) {
const uint64_t row_bytes = (uint64_t)n * sizeof(float);
const uint32_t n_kernel = c4 ? n / 4u : n;
return (ds4_gpu_unary_args) {
.ne00 = (int32_t)n_kernel,
.ne01 = (int32_t)rows,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = row_bytes,
.nb02 = row_bytes,
.nb03 = row_bytes,
.ne0 = (int32_t)n_kernel,
.ne1 = (int32_t)rows,
.ne2 = 1,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = row_bytes,
.nb2 = row_bytes,
.nb3 = row_bytes,
.slope = 0.0f,
.scale = scale,
.bias = bias,
.val = 0.0f,
.min = 0.0f,
.max = 0.0f,
};
}
static ds4_gpu_bin_args ds4_gpu_make_bin_same_rows_args(uint32_t n, uint32_t rows) {
const uint64_t row_bytes = (uint64_t)n * sizeof(float);
return (ds4_gpu_bin_args) {
.ne00 = (int32_t)n,
.ne01 = (int32_t)rows,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = row_bytes,
.nb02 = (uint64_t)rows * row_bytes,
.nb03 = (uint64_t)rows * row_bytes,
.ne10 = (int32_t)n,
.ne11 = (int32_t)rows,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = row_bytes,
.nb12 = (uint64_t)rows * row_bytes,
.nb13 = (uint64_t)rows * row_bytes,
.ne0 = (int32_t)n,
.ne1 = (int32_t)rows,
.ne2 = 1,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = row_bytes,
.nb2 = (uint64_t)rows * row_bytes,
.nb3 = (uint64_t)rows * row_bytes,
.offs = 0,
.o1 = { 0 },
};
}
static int ds4_gpu_encode_bin_f32_rows(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_bin_args *args,
id<MTLBuffer> a,
NSUInteger a_off,
id<MTLBuffer> b,
NSUInteger b_off,
id<MTLBuffer> out,
NSUInteger out_off);
static int ds4_gpu_encode_sum_rows_f32(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t width,
uint32_t rows);
typedef struct {
int32_t ne00;
int32_t ne01;
int32_t ne02;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne10;
int32_t ne11;
int32_t ne12;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
int32_t ne0;
int32_t ne1;
int32_t nr0;
int16_t r2;
int16_t r3;
} ds4_gpu_q8_0_matvec_args;
typedef struct {
int32_t ne00;
int32_t ne02;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne12;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
int32_t ne0;
int32_t ne1;
int16_t r2;
int16_t r3;
} ds4_gpu_mul_mm_args;
typedef struct {
int32_t ne00;
int32_t ne01;
int32_t ne02;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne10;
int32_t ne11;
int32_t ne12;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
int32_t ne0;
int32_t ne1;
int16_t r2;
int16_t r3;
} ds4_gpu_mul_mv_ext_args;
typedef ds4_gpu_q8_0_matvec_args ds4_gpu_f16_matvec_args;
static ds4_gpu_q8_0_matvec_args ds4_gpu_make_q8_0_mv_args(uint64_t in_dim, uint64_t out_dim) {
const uint64_t row_bytes = (in_dim / 32u) * 34u;
return (ds4_gpu_q8_0_matvec_args) {
.ne00 = (int32_t)in_dim,
.ne01 = (int32_t)out_dim,
.ne02 = 1,
.nb00 = 34,
.nb01 = row_bytes,
.nb02 = row_bytes * out_dim,
.nb03 = row_bytes * out_dim,
.ne10 = (int32_t)in_dim,
.ne11 = 1,
.ne12 = 1,
.nb10 = sizeof(float),
.nb11 = in_dim * sizeof(float),
.nb12 = in_dim * sizeof(float),
.nb13 = in_dim * sizeof(float),
.ne0 = (int32_t)out_dim,
.ne1 = 1,
.nr0 = 2,
.r2 = 1,
.r3 = 1,
};
}
static ds4_gpu_f16_matvec_args ds4_gpu_make_f16_mv_args(uint64_t in_dim, uint64_t out_dim) {
const uint64_t row_bytes = in_dim * sizeof(uint16_t);
return (ds4_gpu_f16_matvec_args) {
.ne00 = (int32_t)in_dim,
.ne01 = (int32_t)out_dim,
.ne02 = 1,
.nb00 = sizeof(uint16_t),
.nb01 = row_bytes,
.nb02 = row_bytes * out_dim,
.nb03 = row_bytes * out_dim,
.ne10 = (int32_t)in_dim,
.ne11 = 1,
.ne12 = 1,
.nb10 = sizeof(float),
.nb11 = in_dim * sizeof(float),
.nb12 = in_dim * sizeof(float),
.nb13 = in_dim * sizeof(float),
.ne0 = (int32_t)out_dim,
.ne1 = 1,
.nr0 = 2,
.r2 = 1,
.r3 = 1,
};
}
static ds4_gpu_q8_0_matvec_args ds4_gpu_make_f32_mv_args(
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_vec) {
const uint64_t row_bytes = in_dim * sizeof(float);
return (ds4_gpu_q8_0_matvec_args) {
.ne00 = (int32_t)in_dim,
.ne01 = (int32_t)out_dim,
.ne02 = 1,
.nb00 = sizeof(float),
.nb01 = row_bytes,
.nb02 = row_bytes * out_dim,
.nb03 = row_bytes * out_dim,
.ne10 = (int32_t)in_dim,
.ne11 = (int32_t)n_vec,
.ne12 = 1,
.nb10 = sizeof(float),
.nb11 = in_dim * sizeof(float),
.nb12 = in_dim * n_vec * sizeof(float),
.nb13 = in_dim * n_vec * sizeof(float),
.ne0 = (int32_t)out_dim,
.ne1 = (int32_t)n_vec,
.nr0 = 2,
.r2 = 1,
.r3 = 1,
};
}
typedef struct {
const char *function_name;
int16_t nsg;
int32_t nr0;
NSUInteger smem;
} ds4_gpu_mv_dispatch;
static int ds4_gpu_tp_world_is_two(void);
static ds4_gpu_mv_dispatch ds4_gpu_make_q8_0_mv_dispatch(void) {
const uint64_t default_nsg = ds4_gpu_tp_world_is_two() ? 2u : 4u;
const int16_t nsg =
(int16_t)ds4_gpu_env_u64("DS4_METAL_Q8_MV_NSG", default_nsg, 1u, 8u);
return (ds4_gpu_mv_dispatch) {
.function_name = "kernel_mul_mv_q8_0_f32",
.nsg = nsg,
.nr0 = 2,
.smem = 32u * 2u * sizeof(float),
};
}
/* Single-token F16/F32 matvecs with few output rows (the Qwen hyper-connection
* low-rank down projections and routers) launch one row per SIMD group on M5,
* where two-row tiles leave most of the 40 cores idle. The per-row K walk,
* simdgroup count and reduction tree are unchanged. DS4_METAL_PLAIN_MV_NR0
* forces 1 or 2 rows on any device. */
static bool ds4_gpu_plain_mv_single_row(uint64_t out_dim) {
const uint64_t override = ds4_gpu_env_u64("DS4_METAL_PLAIN_MV_NR0", 0u, 0u, 2u);
if (override) return override == 1u;
return out_dim <= 1024u && ds4_gpu_device_is_m5_apple_silicon();
}
static ds4_gpu_mv_dispatch ds4_gpu_make_plain_mv_dispatch(
uint64_t in_dim,
int f32_weights) {
if (in_dim < 32) {
return (ds4_gpu_mv_dispatch) {
.function_name = f32_weights ? "kernel_mul_mv_f32_f32_short" : "kernel_mul_mv_f16_f32_short",
.nsg = 1,
.nr0 = 32,
.smem = 0,
};
}
const int16_t nsg = (int16_t)((in_dim + 127u) / 128u > 8u ? 8u : (in_dim + 127u) / 128u);
const int use_4 = (in_dim % 4u) == 0;
return (ds4_gpu_mv_dispatch) {
.function_name = f32_weights
? (use_4 ? "kernel_mul_mv_f32_f32_4" : "kernel_mul_mv_f32_f32")
: (use_4 ? "kernel_mul_mv_f16_f32_4" : "kernel_mul_mv_f16_f32"),
.nsg = nsg,
.nr0 = 2,
.smem = 32u * 2u * sizeof(float),
};
}
static ds4_gpu_mul_mm_args ds4_gpu_make_mm_args(
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t row_bytes) {
return (ds4_gpu_mul_mm_args) {
.ne00 = (int32_t)in_dim,
.ne02 = 1,
.nb01 = row_bytes,
.nb02 = row_bytes * out_dim,
.nb03 = row_bytes * out_dim,
.ne12 = 1,
.nb10 = sizeof(float),
.nb11 = in_dim * sizeof(float),
.nb12 = in_dim * n_tok * sizeof(float),
.nb13 = in_dim * n_tok * sizeof(float),
.ne0 = (int32_t)out_dim,
.ne1 = (int32_t)n_tok,
.r2 = 1,
.r3 = 1,
};
}
static ds4_gpu_mul_mv_ext_args ds4_gpu_make_mv_ext_args(
uint64_t in_dim,
uint64_t out_dim,
uint64_t n_tok,
uint64_t elem_bytes,
uint64_t row_bytes) {
return (ds4_gpu_mul_mv_ext_args) {
.ne00 = (int32_t)in_dim,
.ne01 = (int32_t)out_dim,
.ne02 = 1,
.nb00 = elem_bytes,
.nb01 = row_bytes,
.nb02 = row_bytes * out_dim,
.nb03 = row_bytes * out_dim,
.ne10 = (int32_t)in_dim,
.ne11 = (int32_t)n_tok,
.ne12 = 1,
.nb10 = sizeof(float),
.nb11 = in_dim * sizeof(float),
.nb12 = in_dim * n_tok * sizeof(float),
.nb13 = in_dim * n_tok * sizeof(float),
.ne0 = (int32_t)out_dim,
.ne1 = (int32_t)n_tok,
.r2 = 1,
.r3 = 1,
};
}
/* Simdgroups per threadgroup for the few-row (2..16 tokens) Q8/F16 matvecs
* that carry the MTP verify rows. Rows per threadgroup only: each row keeps
* its lane walk and shuffle tree, so outputs are byte-identical at any
* value (tests/test_qwen4_kernels.c pins 1/2/4/8). */
/* Row-exact geometry for the 3-row speculative verify: the few-row matvec
* and routed-expert dispatches keep the lane maps and kernel choices of the
* T <= 2 decode paths, so every committed row stays byte-identical. */
static bool g_qwen4_verify_rows_exact;
void ds4_gpu_qwen4_set_verify_rows_exact(bool on) {
g_qwen4_verify_rows_exact = on;
}
static int16_t ds4_gpu_mv_ext_nsg(void) {
return (int16_t)ds4_gpu_env_u64("DS4_METAL_MV_EXT_NSG", 2u, 1u, 8u);
}
static int16_t ds4_gpu_mv_ext_nxpsg(uint64_t in_dim, uint64_t n_tok) {
if ((in_dim % 256u) == 0 && (n_tok < 3 || (n_tok == 3 && g_qwen4_verify_rows_exact))) return 16;
if ((in_dim % 128u) == 0) return 8;
return 4;
}
static int16_t ds4_gpu_mv_ext_r1ptg(uint64_t n_tok) {
switch (n_tok) {
case 2: return 2;
case 3:
case 6: return 3;
case 4:
case 7:
case 8: return 4;
case 5: return 5;
default: return n_tok > 8 ? 4 : 0;
}
}
static const char *ds4_gpu_mv_ext_name(int q8, int16_t r1ptg) {
if (q8) {
switch (r1ptg) {
case 2: return "kernel_mul_mv_ext_q8_0_f32_r1_2";
case 3: return "kernel_mul_mv_ext_q8_0_f32_r1_3";
case 4: return "kernel_mul_mv_ext_q8_0_f32_r1_4";
case 5: return "kernel_mul_mv_ext_q8_0_f32_r1_5";
default: return NULL;
}
}
switch (r1ptg) {
case 2: return "kernel_mul_mv_ext_f16_f32_r1_2";
case 3: return "kernel_mul_mv_ext_f16_f32_r1_3";
case 4: return "kernel_mul_mv_ext_f16_f32_r1_4";
case 5: return "kernel_mul_mv_ext_f16_f32_r1_5";
default: return NULL;
}
}
static const char *ds4_gpu_mv_ext_f32_name(int16_t r1ptg) {
switch (r1ptg) {
case 2: return "kernel_mul_mv_ext_f32_f32_r1_2";
case 3: return "kernel_mul_mv_ext_f32_f32_r1_3";
case 4: return "kernel_mul_mv_ext_f32_f32_r1_4";
case 5: return "kernel_mul_mv_ext_f32_f32_r1_5";
default: return NULL;
}
}
static const char *ds4_gpu_mv_ext_q8_pair_swiglu_name(int16_t r1ptg) {
switch (r1ptg) {
case 2: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_2";
case 3: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_3";
case 4: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_4";
case 5: return "kernel_mul_mv_ext_q8_0_pair_swiglu_f32_r1_5";
default: return NULL;
}
}
typedef struct {
int32_t ne00;
int32_t ne00_t;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
float eps;
int32_t nef1[3];
int32_t nef2[3];
int32_t nef3[3];
uint64_t nbf1[3];
uint64_t nbf2[3];
uint64_t nbf3[3];
} ds4_gpu_rms_norm_args;
typedef struct {
int32_t q_n;
int32_t q_n4;
int32_t kv_n;
int32_t kv_n4;
uint64_t q_row_stride;
uint64_t kv_row_stride;
float eps;
uint32_t tasks;
} ds4_gpu_qkv_rms_norm_args;
static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_args(uint32_t n, uint32_t rows, float eps) {
const uint64_t row_bytes = (uint64_t)n * sizeof(float);
return (ds4_gpu_rms_norm_args) {
.ne00 = (int32_t)n,
.ne00_t = (int32_t)(n / 4u),
.nb1 = row_bytes,
.nb2 = row_bytes * rows,
.nb3 = row_bytes * rows,
.eps = eps,
.nef1 = { (int32_t)rows, 1, 1 },
.nef2 = { 1, 1, 1 },
.nef3 = { 1, 1, 1 },
.nbf1 = { row_bytes, row_bytes, row_bytes },
.nbf2 = { row_bytes * rows, row_bytes, row_bytes },
.nbf3 = { row_bytes * rows, row_bytes, row_bytes },
};
}
static ds4_gpu_rms_norm_args ds4_gpu_make_rms_norm_3d_args(
uint32_t n0,
uint32_t n1,
uint32_t n2,
float eps) {
const uint64_t row_bytes = (uint64_t)n0 * sizeof(float);
const uint64_t plane_bytes = row_bytes * n1;
return (ds4_gpu_rms_norm_args) {
.ne00 = (int32_t)n0,
.ne00_t = (int32_t)(n0 / 4u),
.nb1 = row_bytes,
.nb2 = plane_bytes,
.nb3 = plane_bytes * n2,
.eps = eps,
.nef1 = { (int32_t)n1, 1, 1 },
.nef2 = { (int32_t)n2, 1, 1 },
.nef3 = { 1, 1, 1 },
.nbf1 = { row_bytes, row_bytes, row_bytes },
.nbf2 = { plane_bytes, row_bytes, row_bytes },
.nbf3 = { plane_bytes * n2, row_bytes, row_bytes },
};
}
static NSUInteger ds4_gpu_rms_norm_threads(uint32_t n) {
NSUInteger ne00_t = n / 4u;
NSUInteger nth = 32u;
while (nth < ne00_t && nth < 1024u) nth *= 2u;
if (nth > ne00_t) nth = ne00_t;
return nth ? nth : 1u;
}
static NSUInteger ds4_gpu_rms_norm_pipeline_threads(
uint32_t n,
id<MTLComputePipelineState> pipeline) {
NSUInteger ne00_t = n / 4u;
NSUInteger max_threads = pipeline ? [pipeline maxTotalThreadsPerThreadgroup] : 1024u;
NSUInteger nth = 32u;
while (nth < ne00_t && nth < max_threads) nth *= 2u;
if (nth > max_threads) nth = max_threads;
if (nth > ne00_t) nth = ne00_t;
return nth ? nth : 1u;
}
typedef struct {
int32_t n_hc;
int32_t sinkhorn_iters;
int64_t n_rows;
int64_t mix_hc;
uint64_t nb01;
uint64_t nb1;
float eps;
} ds4_gpu_hc_split_args;
typedef struct {
int64_t n_embd;
int64_t n_hc;
int64_t n_tokens;
uint64_t nb_x0;
uint64_t nb_x1;
uint64_t nb_x2;
uint64_t nb_w0;
uint64_t nb_w1;
uint64_t nb0;
uint64_t nb1;
uint64_t round_bf16;
} ds4_gpu_hc_weighted_sum_args;
typedef struct {
float post_scale;
float eps;
} ds4_gpu_output_hc_weights4_args;
typedef struct {
int64_t n_embd;
int32_t n_hc;
int32_t sinkhorn_iters;
int64_t n_rows;
int64_t mix_hc;
uint64_t nb_mix1;
uint64_t nb_split1;
uint64_t nb_x0;
uint64_t nb_x1;
uint64_t nb_x2;
uint64_t nb0;
uint64_t nb1;
float eps;
} ds4_gpu_hc_split_weighted_sum_args;
typedef struct {
int64_t n_embd;
int32_t n_hc;
int32_t sinkhorn_iters;
int64_t n_rows;
int64_t mix_hc;
uint64_t nb_mix1;
uint64_t nb_split1;
uint64_t nb_x0;
uint64_t nb_x1;
uint64_t nb_x2;
uint64_t nb0;
uint64_t nb1;
uint64_t nb_norm1;
float eps;
float norm_eps;
} ds4_gpu_hc_split_weighted_sum_norm_args;
typedef struct {
int64_t n_embd;
int64_t n_hc;
int64_t n_tokens;
uint64_t nb_block0;
uint64_t nb_block1;
uint64_t nb_add0;
uint64_t nb_add1;
uint64_t nb_res0;
uint64_t nb_res1;
uint64_t nb_res2;
uint64_t nb_post0;
uint64_t nb_post1;
uint64_t nb_comb0;
uint64_t nb_comb1;
uint64_t nb_comb2;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
int32_t has_add;
int32_t round_bf16;
int32_t round_block_in;
} ds4_gpu_hc_expand_args;
typedef struct {
int32_t nei0;
int32_t nei1;
uint64_t nbi1;
int32_t ne00;
int32_t ne01;
int32_t ne02;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
int32_t ne10;
int32_t ne11;
int32_t ne12;
int32_t ne13;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
int32_t ne0;
int32_t ne1;
uint64_t nb1;
int32_t nr0;
/* Tensor-parallel expert ownership; see ds4_metal_args_mul_mv_id in
* metal/moe.metal. Zero (from struct literals) means no split. */
int32_t tp_rank;
int32_t tp_world;
int32_t tp_addend;
int32_t tp_expert_base;
} ds4_gpu_mul_mv_id_args;
typedef struct {
uint32_t n_total_expert;
uint32_t n_expert;
} ds4_gpu_stream_expert_validate_args;
typedef struct {
uint32_t active_mask;
uint32_t accumulate;
} ds4_gpu_stream_expert_split_args;
typedef struct {
int32_t ne02;
int32_t ne10;
int32_t ne11;
uint64_t nb11;
uint64_t nb12;
int32_t ne21;
int32_t ne20;
uint64_t nb21;
} ds4_gpu_mul_mm_id_map_args;
typedef struct {
int32_t ne00;
int32_t ne02;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne11;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
int32_t ne20;
int32_t ne21;
int32_t ne0;
int32_t ne1;
int16_t r2;
int16_t r3;
int32_t tp_rank;
int32_t tp_world;
int32_t tp_expert_base;
} ds4_gpu_mul_mm_id_args;
static int ds4_gpu_encode_mul_mv_id(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0);
static int ds4_gpu_encode_attn_out_low_q8_direct(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0);
static int ds4_gpu_encode_attn_out_low_mpp(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off);
static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args(
uint32_t src0_cols,
uint32_t src0_experts,
uint32_t src1_expert_rows,
uint32_t selected_experts,
uint32_t n_tokens);
static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args(
uint32_t src0_cols,
uint32_t src0_rows,
uint32_t src0_experts,
uint64_t src0_row_bytes,
uint64_t src0_expert_bytes,
uint32_t src1_expert_rows,
uint32_t selected_experts,
uint32_t n_tokens);
static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size(
uint32_t src0_cols,
uint32_t src0_rows,
uint32_t src0_experts,
uint64_t src0_row_bytes,
uint64_t src0_expert_bytes,
uint32_t src1_expert_rows,
uint32_t selected_experts,
uint32_t n_tokens,
uint32_t src1_elem_size);
static int ds4_gpu_encode_mul_mm_id(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> map_pipeline,
id<MTLComputePipelineState> mm_pipeline,
const ds4_gpu_mul_mm_id_map_args *map_args,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off);
static int ds4_gpu_encode_mul_mm_id_map(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> map_pipeline,
const ds4_gpu_mul_mm_id_map_args *map_args,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> ids,
NSUInteger ids_off);
static int ds4_gpu_encode_mul_mm_id_mapped(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> mm_pipeline,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off);
static int ds4_gpu_encode_mul_mm_id_mapped_tile(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> mm_pipeline,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
NSUInteger threadgroup_bytes);
static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> mm_pipeline,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0_addrs,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
NSUInteger threadgroup_bytes,
ds4_gpu_stream_expert_cache_entry * const *resources,
uint32_t resource_count,
uint32_t resource_kind,
id<MTLBuffer> overflow_resource);
typedef struct {
int32_t ne11;
int32_t ne_12_2;
int32_t ne_12_3;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
uint64_t nb21;
uint64_t nb22;
uint64_t nb23;
int32_t ne31;
int32_t ne32;
int32_t ne33;
uint64_t nb31;
uint64_t nb32;
uint64_t nb33;
} ds4_gpu_flash_attn_pad_args;
typedef struct {
uint32_t raw_cap;
uint32_t raw_start;
uint32_t n_raw;
uint32_t n_comp;
uint32_t pad_rows;
uint32_t shared_pad;
} ds4_gpu_flash_kv_stage_f16_args;
typedef struct {
int32_t ne01;
int32_t ne30;
int32_t ne31;
int32_t ne32;
int32_t ne33;
uint64_t nb31;
uint64_t nb32;
uint64_t nb33;
} ds4_gpu_flash_attn_blk_args;
typedef struct {
int32_t ne01;
int32_t ne02;
int32_t ne03;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne11;
int32_t ne_12_2;
int32_t ne_12_3;
int32_t ns10;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
int32_t ns20;
uint64_t nb21;
uint64_t nb22;
uint64_t nb23;
int32_t ne31;
int32_t ne32;
int32_t ne33;
uint64_t nb31;
uint64_t nb32;
uint64_t nb33;
int32_t ne1;
int32_t ne2;
int32_t ne3;
float scale;
float max_bias;
float m0;
float m1;
int32_t n_head_log2;
float logit_softcap;
} ds4_gpu_flash_attn_vec_args;
typedef struct {
int32_t nrows;
} ds4_gpu_flash_attn_reduce_args;
typedef struct {
int64_t ne00;
int64_t ne01;
int64_t ne02;
int64_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
int32_t n_dims;
int32_t mode;
int32_t n_ctx_orig;
int32_t inverse;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
bool src2;
} ds4_gpu_rope_tail_batch_args;
typedef struct {
uint64_t row_bytes;
uint64_t token_bytes;
int32_t head_dim;
int32_t n_dims;
int32_t n_ctx_orig;
int32_t inverse;
uint32_t pos0;
uint32_t pos_step;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
} ds4_gpu_rope_affine_pair_args;
/* Matches ds4_metal_args_dsv4_comp_finalize in dsv4_rope.metal. */
typedef struct {
ds4_gpu_rope_affine_pair_args rope;
float rms_eps;
uint32_t pad0;
} ds4_gpu_comp_finalize_args;
/* Set by ds4.c immediately before the decode attention call when the inverse
* RoPE tail is deferred into the reduce kernel. Cleared by the encoder. */
static ds4_gpu_rope_affine_pair_args g_decode_attn_rope_args;
static int g_decode_attn_rope_fuse;
/* Set only by the encoder that actually applied the deferred rotation, so ds4.c
* can fall back to the standalone RoPE on any attention path that does not
* consume it (for example ratio-0 layers that take the raw-heads encoder). */
static int g_decode_attn_rope_fuse_used;
void ds4_gpu_set_decode_attn_rope_fuse(
uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t n_ctx_orig,
bool inverse, float freq_base, float freq_scale, float ext_factor,
float attn_factor, float beta_fast, float beta_slow) {
const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float);
g_decode_attn_rope_args = (ds4_gpu_rope_affine_pair_args) {
.row_bytes = row_bytes,
.token_bytes = row_bytes,
.head_dim = (int32_t)head_dim,
.n_dims = (int32_t)n_rot,
.n_ctx_orig = (int32_t)n_ctx_orig,
.inverse = inverse ? 1 : 0,
.pos0 = pos0,
.pos_step = 1,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
};
g_decode_attn_rope_fuse = 1;
g_decode_attn_rope_fuse_used = 0;
}
int ds4_gpu_decode_attn_rope_fuse_used(void) {
return g_decode_attn_rope_fuse_used;
}
_Static_assert(sizeof(ds4_gpu_rope_affine_pair_args) == 64,
"Metal affine RoPE argument ABI changed");
static ds4_gpu_rope_tail_batch_args ds4_gpu_make_rope_tail_args(
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float);
const uint64_t tok_bytes = (uint64_t)n_head * row_bytes;
return (ds4_gpu_rope_tail_batch_args) {
.ne00 = head_dim,
.ne01 = n_head,
.ne02 = n_tok,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = row_bytes,
.nb02 = tok_bytes,
.nb03 = (uint64_t)n_tok * tok_bytes,
.nb0 = sizeof(float),
.nb1 = row_bytes,
.nb2 = tok_bytes,
.nb3 = (uint64_t)n_tok * tok_bytes,
.n_dims = (int32_t)n_rot,
.mode = 0,
.n_ctx_orig = (int32_t)n_ctx_orig,
.inverse = inverse ? 1 : 0,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
.src2 = false,
};
}
static int ds4_gpu_encode_rope_tail_inplace(
id<MTLCommandBuffer> cb,
id<MTLBuffer> xbuf,
NSUInteger xoff,
const ds4_gpu_rope_tail_batch_args *args,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t pos0,
uint32_t pos_step) {
const uint32_t tail_threads = args->n_dims > 0 ? (uint32_t)args->n_dims : 0u;
const bool lane_compatible =
tail_threads <= head_dim && ((head_dim - tail_threads) & 31u) == 0u;
const bool use_inplace_pair =
g_rope_tail_inplace_pair_pipeline != nil &&
getenv("DS4_METAL_DISABLE_INPLACE_ROPE_PAIR") == NULL &&
args->mode == 0 && !args->src2 &&
lane_compatible &&
(ds4_gpu_device_name_contains("M3") ||
(ds4_gpu_device_name_contains("M5") && n_tok == 1u));
const bool use_shared_coeff =
use_inplace_pair &&
g_rope_tail_inplace_pair_shared4_pipeline != nil &&
getenv("DS4_METAL_DISABLE_SHARED_ROPE_COEFF") == NULL &&
/* The 256-thread grouped schedule helps long prefill, but reduces the
* per-head parallelism that short batches and decode rely on. */
tail_threads == 64u && n_head >= 4u && n_tok >= 32u;
/* Keep long prefill on the proven shared4 kernel. Reconstructing affine
* positions inside its coefficient cohort perturbs YaRN fast-math codegen;
* the compact affine specialization is exact for the decode pair schedule. */
const bool use_affine_position =
use_inplace_pair && !use_shared_coeff &&
g_rope_tail_inplace_pair_affine_pipeline != nil &&
getenv("DS4_METAL_DISABLE_AFFINE_ROPE_PAIR") == NULL &&
n_tok == 1u &&
(ds4_gpu_device_name_contains("M3") ||
ds4_gpu_device_name_contains("M5"));
int32_t pos_stack[256];
int32_t *pos = NULL;
id<MTLBuffer> posbuf = nil;
const NSUInteger pos_bytes = (NSUInteger)n_tok * sizeof(int32_t);
if (!use_affine_position) {
pos = pos_stack;
if (n_tok > (uint32_t)(sizeof(pos_stack) / sizeof(pos_stack[0]))) {
pos = malloc((size_t)n_tok * sizeof(*pos));
if (!pos) {
fprintf(stderr, "ds4: failed to allocate Metal RoPE position buffer\n");
return 0;
}
}
for (uint32_t t = 0; t < n_tok; t++) {
pos[t] = (int32_t)(pos0 + t * pos_step);
}
if (pos_bytes > 4096u) {
/*
* Metal inline setBytes data is meant for small constants. Long
* prefill RoPE calls need thousands of positions; passing that much
* inline can make the Apple driver abort the process.
*/
posbuf = ds4_gpu_new_transient_buffer(
pos_bytes, "ds4_rope_positions");
if (!posbuf) {
if (pos != pos_stack) free(pos);
return 0;
}
memcpy([posbuf contents], pos, pos_bytes);
}
}
ds4_gpu_rope_affine_pair_args affine_args;
if (use_affine_position) {
const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float);
affine_args = (ds4_gpu_rope_affine_pair_args) {
.row_bytes = row_bytes,
.token_bytes = (uint64_t)n_head * row_bytes,
.head_dim = (int32_t)head_dim,
.n_dims = args->n_dims,
.n_ctx_orig = args->n_ctx_orig,
.inverse = args->inverse,
.pos0 = pos0,
.pos_step = pos_step,
.freq_base = args->freq_base,
.freq_scale = args->freq_scale,
.ext_factor = args->ext_factor,
.attn_factor = args->attn_factor,
.beta_fast = args->beta_fast,
.beta_slow = args->beta_slow,
};
}
const NSUInteger reference_nth =
(NSUInteger)(head_dim < 256u ? head_dim : 256u);
const NSUInteger pair_nth =
(NSUInteger)(tail_threads < 256u ? tail_threads : 256u);
const NSUInteger nth = use_shared_coeff ? 256u :
(use_inplace_pair ? pair_nth : reference_nth);
const NSUInteger head_groups = use_shared_coeff ?
(NSUInteger)((n_head + 3u) / 4u) : (NSUInteger)n_head;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
id<MTLComputePipelineState> pipeline = use_affine_position ?
g_rope_tail_inplace_pair_affine_pipeline :
(use_shared_coeff ?
g_rope_tail_inplace_pair_shared4_pipeline :
(use_inplace_pair ?
g_rope_tail_inplace_pair_pipeline : g_rope_tail_batch_pipeline));
[enc setComputePipelineState:pipeline];
if (use_affine_position) {
[enc setBytes:&affine_args length:sizeof(affine_args) atIndex:0];
} else {
[enc setBytes:args length:sizeof(*args) atIndex:0];
}
[enc setBuffer:xbuf offset:xoff atIndex:1];
if (!use_affine_position) {
if (posbuf) {
[enc setBuffer:posbuf offset:0 atIndex:2];
} else {
[enc setBytes:pos length:pos_bytes atIndex:2];
}
[enc setBuffer:xbuf offset:xoff atIndex:3];
}
[enc setBuffer:xbuf offset:xoff atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(head_groups, n_tok, 1)
threadsPerThreadgroup:MTLSizeMake(nth ? nth : 1u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (pos && pos != pos_stack) free(pos);
return 1;
}
typedef struct {
int64_t ne00;
int64_t ne01;
int64_t ne02;
int64_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
int32_t n_rot;
} ds4_gpu_dsv4_fp8_kv_quantize_args;
typedef struct {
int32_t head_dim;
int32_t n_rot;
int32_t raw_row;
} ds4_gpu_dsv4_kv_fp8_store_args;
typedef struct {
uint32_t n_rows;
uint32_t head_dim;
uint64_t row_stride;
} ds4_gpu_dsv4_indexer_qat_args;
typedef struct {
uint32_t width;
} ds4_gpu_dsv4_ratio4_shift_args;
typedef struct {
uint32_t head_dim;
uint32_t n_comp;
uint32_t replay;
uint32_t n_threads;
} ds4_gpu_dsv4_compressor_pack_ratio4_args;
typedef struct {
int64_t n_rows;
uint32_t head_dim;
uint32_t n_comp;
uint32_t replay;
uint32_t pad;
} ds4_gpu_dsv4_softmax_pool_ratio4_direct_args;
typedef struct {
uint32_t width;
uint32_t ratio;
uint32_t pos;
uint32_t ape_type;
} ds4_gpu_dsv4_compressor_store_one_args;
typedef struct {
int64_t ne00;
int64_t ne01;
int64_t ne02;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb10;
uint64_t nb11;
uint64_t nb12;
int64_t ne0;
int64_t ne1;
uint64_t nb0;
uint64_t nb1;
} ds4_gpu_dsv4_softmax_pool_args;
typedef struct {
uint32_t width;
uint32_t ratio;
uint32_t pos0;
uint32_t n_tokens;
} ds4_gpu_dsv4_compressor_score_ape_args;
typedef struct {
int32_t ne00;
int32_t ne01;
int32_t ne02;
int32_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne0;
int32_t ne1;
int32_t ne2;
int32_t ne3;
int32_t top_k;
uint32_t causal_start;
uint32_t causal_ratio;
} ds4_gpu_kargs_argsort;
typedef struct {
int64_t ne00;
int64_t ne01;
int64_t ne02;
int64_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne0;
int32_t ne1;
int32_t ne2;
int32_t ne3;
int32_t top_k;
int32_t len;
uint32_t causal_start;
uint32_t causal_ratio;
uint32_t block_width;
uint32_t block_top_k;
} ds4_gpu_kargs_argsort_merge;
typedef struct {
uint32_t n;
uint32_t rows;
uint32_t top_k;
uint32_t pass;
uint32_t parity;
uint32_t n_tg;
} ds4_gpu_kargs_topk_select;
typedef struct {
int64_t ne00;
int64_t ne01;
int64_t ne02;
int64_t ne03;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int64_t ne0;
int64_t ne1;
int64_t ne2;
int64_t ne3;
uint64_t nb0;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
} ds4_gpu_kargs_sum_rows;
typedef struct {
int32_t ne00;
int32_t ne01;
int32_t ne02;
uint64_t nb01;
uint64_t nb02;
uint64_t nb03;
int32_t ne11;
int32_t ne12;
int32_t ne13;
uint64_t nb11;
uint64_t nb12;
uint64_t nb13;
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
float scale;
float max_bias;
float m0;
float m1;
int32_t n_head_log2;
} ds4_gpu_softmax_args;
typedef struct {
int64_t ne00;
int64_t ne01;
uint64_t nb00;
uint64_t nb01;
int64_t ne0;
int64_t ne1;
uint64_t nb0;
uint64_t nb1;
} ds4_gpu_dsv4_topk_mask_args;
typedef struct {
int64_t ne00;
int64_t ne01;
int64_t ne02;
uint64_t nb00;
uint64_t nb01;
uint64_t nb02;
int64_t ne10;
int64_t ne11;
uint64_t nb10;
uint64_t nb11;
int64_t ne0;
int64_t ne1;
uint64_t nb0;
uint64_t nb1;
float scale;
} ds4_gpu_dsv4_indexer_weighted_sum_args;
typedef struct {
uint32_t has_bias;
uint32_t hash_mode;
uint32_t use_token_buffer;
uint32_t token;
uint32_t hash_rows;
} ds4_gpu_dsv4_router_select_one_args;
typedef struct {
uint32_t hash_rows;
uint32_t vocab_size;
uint32_t n_tokens;
uint32_t has_bias;
uint32_t hash_mode;
} ds4_gpu_dsv4_router_select_visual_args;
typedef struct {
uint32_t n_expert;
uint32_t n_expert_used;
float expert_weight_scale;
uint32_t pad0;
} ds4_gpu_glm_router_select_one_args;
typedef struct {
uint32_t n_tokens;
uint32_t kv_raw_dim;
uint32_t kv_lora_dim;
float eps;
} ds4_gpu_glm_kv_lora_rms_norm_args;
typedef struct {
uint32_t n_tokens;
uint32_t kv_lora_dim;
uint32_t qk_nope;
uint32_t n_head;
uint32_t row_bytes;
uint32_t weight_type;
uint32_t pad1;
uint32_t pad2;
} ds4_gpu_glm_k_b_project_args;
typedef struct {
uint32_t pos0;
uint32_t n_tokens;
uint32_t cache_cap;
uint32_t kv_raw_dim;
uint32_t kv_lora_dim;
uint32_t qk_rope;
uint32_t cache_f16;
uint32_t pad1;
} ds4_gpu_glm_store_compact_kv_args;
typedef struct {
uint32_t pos0;
uint32_t n_tokens;
uint32_t cache_cap;
uint32_t q_n;
uint32_t q_n4;
uint32_t kv_raw_dim;
uint32_t kv_lora_dim;
uint32_t kv_lora_n4;
uint32_t qk_rope;
uint32_t cache_f16;
float eps;
uint32_t pad0;
} ds4_gpu_glm_qkv_norm_store_compact_kv_args;
typedef struct {
uint32_t pos0;
uint32_t n_tokens;
uint32_t cache_cap;
uint32_t head_dim;
uint32_t rot_dim;
uint32_t n_ctx_orig;
uint32_t cache_f16;
uint32_t pad0;
float eps;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
float pad1;
} ds4_gpu_glm_store_indexer_k_args;
typedef struct {
uint32_t pos0;
uint32_t n_tokens;
uint32_t cache_cap;
uint32_t head_dim;
uint32_t pool_size;
uint32_t cache_f16;
float eps;
uint32_t pad0;
} ds4_gpu_glm53_indexer_pool_update_args;
typedef struct {
uint32_t pos0;
uint32_t n_tokens;
uint32_t cache_cap;
uint32_t n_head;
uint32_t kv_raw_dim;
uint32_t kv_lora_dim;
uint32_t qk_nope;
uint32_t qk_rope;
uint32_t value_dim;
uint32_t n_ctx_orig;
uint32_t cache_f16;
uint32_t pad0;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
} ds4_gpu_glm_build_kv_cache_args;
typedef struct {
uint32_t pos0;
uint32_t n_tokens;
uint32_t cache_len;
uint32_t cache_cap;
uint32_t n_head;
uint32_t qk_dim;
uint32_t value_dim;
uint32_t pad0;
uint32_t cache_f16;
uint32_t pad1;
uint32_t pad2;
float scale;
} ds4_gpu_glm_attention_full_args;
typedef struct {
uint32_t n_selected;
} ds4_gpu_glm_fill_selected_range_args;
typedef struct {
uint32_t n_tokens;
uint32_t pos0;
uint32_t n_selected;
uint32_t pad_row;
} ds4_gpu_glm_fill_selected_range_batch_args;
typedef struct {
uint32_t n_tokens;
uint32_t pos0;
uint32_t selected_pools;
uint32_t index_topk;
uint32_t pool_size;
uint32_t output_width;
} ds4_gpu_glm53_expand_pool_selection_args;
typedef struct {
uint32_t n_tokens;
uint32_t n_head;
uint32_t head_dim;
uint32_t rot_dim;
uint32_t rot_offset;
uint32_t pos0;
uint32_t n_ctx_orig;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
} ds4_gpu_glm_rope_tail_args;
typedef struct {
uint32_t n_rows;
uint32_t n_head;
uint32_t head_dim;
uint32_t cache_f16;
float scale;
} ds4_gpu_glm_indexer_score_one_args;
typedef struct {
uint32_t n_rows;
uint32_t n_tokens;
uint32_t n_head;
uint32_t head_dim;
uint32_t pos0;
uint32_t cache_f16;
uint32_t row_group_size;
uint32_t pad0;
uint64_t q_token_stride;
uint64_t q_head_stride;
uint64_t weights_token_stride;
uint64_t score_token_stride;
float scale;
} ds4_gpu_glm_indexer_scores_batch_args;
typedef struct {
uint32_t n_head;
uint32_t kv_lora_dim;
uint32_t qk_nope;
uint32_t qk_dim;
uint32_t row_bytes;
uint32_t weight_type;
uint32_t pad1;
uint32_t pad2;
} ds4_gpu_glm_qk_lowrank_args;
typedef struct {
uint32_t n_tokens;
uint32_t n_head;
uint32_t kv_lora_dim;
uint32_t qk_nope;
uint32_t qk_dim;
uint32_t row_bytes;
uint32_t weight_type;
uint32_t head_base;
} ds4_gpu_glm_qk_lowrank_batch_args;
typedef struct {
uint32_t n_selected;
uint32_t cache_cap;
uint32_t cache_f16;
uint32_t n_head;
uint32_t kv_lora_dim;
uint32_t qk_nope;
uint32_t qk_rope;
uint32_t value_dim;
uint32_t n_ctx_orig;
uint32_t value_row_bytes;
float scale;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
uint32_t value_type;
} ds4_gpu_glm_attention_indexed_decode_args;
typedef struct {
uint32_t n_selected;
uint32_t cache_cap;
uint32_t cache_f16;
uint32_t n_head;
uint32_t kv_lora_dim;
uint32_t qk_nope;
uint32_t qk_rope;
uint32_t value_dim;
uint32_t n_ctx_orig;
uint32_t value_row_bytes;
uint32_t block_rows;
uint32_t n_blocks;
float scale;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
uint32_t value_type;
} ds4_gpu_glm_attention_indexed_decode_split_args;
typedef struct {
uint32_t n_tokens;
uint32_t n_selected;
uint32_t cache_cap;
uint32_t cache_f16;
uint32_t n_head;
uint32_t kv_lora_dim;
uint32_t qk_nope;
uint32_t qk_rope;
uint32_t value_dim;
uint32_t n_ctx_orig;
uint32_t value_row_bytes;
uint32_t value_type;
uint32_t pos0;
float scale;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
uint32_t head_base;
} ds4_gpu_glm_attention_indexed_batch_args;
typedef struct {
uint32_t in_dim;
uint32_t mid_dim;
uint32_t out_dim;
uint32_t n_total_expert;
uint32_t n_expert_used;
uint32_t n_tokens;
uint32_t mid_token_stride;
uint32_t down_type;
float swiglu_clamp;
int32_t tp_rank;
int32_t tp_world;
int32_t tp_expert_base;
uint64_t gate_expert_bytes;
uint64_t gate_row_bytes;
uint64_t up_expert_bytes;
uint64_t up_row_bytes;
uint64_t down_expert_bytes;
uint64_t down_row_bytes;
} ds4_gpu_glm_routed_moe_args;
typedef struct {
uint32_t n_tokens;
uint32_t n_head;
uint32_t n_raw;
uint32_t raw_cap;
uint32_t raw_start;
uint32_t n_comp;
uint32_t top_k;
uint32_t pos0;
uint32_t window;
uint32_t ratio;
uint32_t comp_kv_f16;
uint32_t n_splits;
uint64_t q_token_stride;
uint64_t q_head_stride;
uint64_t raw_row_stride;
uint64_t comp_row_stride;
uint64_t topk_token_stride;
uint64_t dst_token_stride;
uint64_t dst_head_stride;
float scale;
} ds4_gpu_dsv4_indexed_attention_args;
typedef struct {
uint32_t n_comp;
uint32_t n_tokens;
uint32_t n_head;
uint32_t head_dim;
uint32_t pos0;
uint32_t ratio;
uint64_t q_token_stride;
uint64_t q_head_stride;
uint64_t weights_token_stride;
uint64_t index_row_stride;
uint64_t score_token_stride;
float scale;
} ds4_gpu_dsv4_indexer_scores_fused_args;
typedef struct {
uint32_t width;
uint32_t rows;
uint64_t gate_row_stride;
uint64_t up_row_stride;
uint64_t mid_row_stride;
uint64_t weight_stride;
uint32_t write_clamped;
float clamp_value;
} ds4_gpu_dsv4_moe_swiglu_weight_args;
typedef struct {
uint32_t expert_base;
uint32_t expert_count;
uint32_t accumulate;
uint32_t pad0;
} ds4_gpu_moe_expert_group_args;
typedef struct {
uint64_t expert_bytes;
uint32_t group_size;
uint32_t n_slots;
} ds4_gpu_q4_gather_slots6_args;
typedef struct {
uint32_t width;
uint32_t tokens;
uint64_t src_token_stride;
uint64_t dst_token_stride;
} ds4_gpu_dsv4_moe_sum6_args;
/* Compile the single in-repo Metal source and create the pipelines that every
* session uses. Shape-dependent kernels with function constants are built
* lazily by the small ds4_gpu_get_* caches, so startup stays predictable
* while long-context prefill and decode can still pick specialized variants. */
int ds4_gpu_init(void) {
if (g_initialized) return 1;
@autoreleasepool {
ds4_gpu_decode_pipeline_fast_cache_reset();
g_device = MTLCreateSystemDefaultDevice();
if (!g_device) {
fprintf(stderr, "ds4: Metal device not available\n");
return 0;
}
ds4_gpu_timeline_probe(g_device);
ds4_gpu_print_device_summary();
ds4_gpu_detect_metal4_features();
g_queue = [g_device newCommandQueueWithMaxCommandBufferCount:256];
if (!g_queue) {
fprintf(stderr, "ds4: failed to create Metal command queue\n");
g_device = nil;
return 0;
}
g_model_buffer_cache = [NSMutableDictionary dictionary];
g_model_buffer_cache_bytes = 0;
g_model_buffer_cache_evictions = 0;
g_model_buffer_cache_over_limit = 0;
g_q4_expert_table_cache = [NSMutableDictionary dictionary];
g_q4_expert_layer_residency_cache = [NSMutableDictionary dictionary];
g_pipeline_cache = [NSMutableDictionary dictionary];
g_dsv4_completion_cache = [NSCache new];
g_dsv4_completion_cache.countLimit = 256u;
g_transient_buffers = [NSMutableArray array];
g_pending_cbs = [NSMutableArray array];
if (!g_model_buffer_cache || !g_q4_expert_table_cache ||
!g_q4_expert_layer_residency_cache ||
!g_pipeline_cache || !g_dsv4_completion_cache ||
!g_transient_buffers || !g_pending_cbs) {
fprintf(stderr, "ds4: Metal bookkeeping allocation failed\n");
g_pending_cbs = nil;
g_transient_buffers = nil;
g_dsv4_completion_cache = nil;
g_pipeline_cache = nil;
g_q4_expert_layer_residency_cache = nil;
g_q4_expert_table_cache = nil;
g_model_buffer_cache = nil;
g_queue = nil;
g_device = nil;
return 0;
}
NSError *error = nil;
NSString *source = ds4_gpu_full_source();
if (!source) {
g_queue = nil;
g_device = nil;
return 0;
}
MTLCompileOptions *options = [MTLCompileOptions new];
NSMutableDictionary *macros = [NSMutableDictionary new];
if (g_metal4_tensor_api_enabled) {
macros[@"DS4_METAL_HAS_TENSOR"] = @"1";
fprintf(stderr, "ds4: Metal 4 tensor API enabled for Tensor kernels\n");
}
const int drift_hc_stable = ds4_gpu_env_bool("DS4_METAL_HC_STABLE") != 0; // default ON
const int drift_norm_unify = ds4_gpu_env_bool("DS4_METAL_NORM_RSQRT_DISABLE") != 0; // default ON
const int drift_kv_raw_f32 = ds4_gpu_env_bool("DS4_METAL_KV_RAW_F32") > 0; // default OFF
const int drift_rope_exp2_log2 = ds4_gpu_env_bool("DS4_METAL_ROPE_EXP2_LOG2") > 0; // default OFF
const int drift_math_safe = ds4_gpu_env_bool("DS4_METAL_MATH_SAFE") > 0; // default OFF
if (drift_math_safe) {
// MTLCompileOptions.fastMathEnabled defaults to YES and Apple's
// headers explicitly say this "may violate the IEEE 754 standard".
// Different fast-math optimizations get applied across the
// matmul2d cooperative-tensor path and the legacy
// simdgroup_multiply_accumulate path on M5, amplifying the
// mismatch. MTLMathModeSafe pins the entire library to strict
// IEEE-754 semantics. Diagnostic-only: useful to localize drift
// sources but not to ship as a default.
if (@available(macOS 15.0, *)) {
options.mathMode = MTLMathModeSafe;
fprintf(stderr, "ds4: Metal shader library math mode = safe (strict IEEE-754) by DS4_METAL_MATH_SAFE\n");
} else {
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
options.fastMathEnabled = NO;
#pragma clang diagnostic pop
fprintf(stderr, "ds4: Metal shader library fast-math disabled by DS4_METAL_MATH_SAFE (pre-macOS 15)\n");
}
}
if (drift_hc_stable) macros[@"DS4_METAL_HC_STABLE"] = @"1";
if (drift_norm_unify) macros[@"DS4_METAL_NORM_RSQRT_DISABLE"] = @"1";
if (drift_kv_raw_f32) macros[@"DS4_METAL_KV_RAW_F32"] = @"1";
if (drift_rope_exp2_log2) macros[@"DS4_METAL_ROPE_EXP2_LOG2"] = @"1";
fprintf(stderr,
"ds4: drift-patch flags hc_stable=%s norm_unify=%s kv_raw_f32=%s rope_exp2_log2=%s math_safe=%s tensor_matmul=%s\n",
drift_hc_stable ? "on" : "off",
drift_norm_unify ? "on" : "off",
drift_kv_raw_f32 ? "on" : "off",
drift_rope_exp2_log2 ? "on" : "off",
drift_math_safe ? "on" : "off",
g_metal4_tensor_api_enabled ? "on" : "off");
options.preprocessorMacros = macros;
id<MTLLibrary> library = [g_device newLibraryWithSource:source options:options error:&error];
if (!library) {
fprintf(stderr, "ds4: Metal shader compilation failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_library = library;
id<MTLFunction> fn = [library newFunctionWithName:@"kernel_get_rows_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_get_rows_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_get_rows_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_get_rows_f32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_get_rows_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_get_rows_f16"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_get_rows_f16 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_get_rows_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_get_rows_f16_pipeline) {
fprintf(stderr, "ds4: Metal kernel_get_rows_f16 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_get_rows_i32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_get_rows_i32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_get_rows_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_get_rows_i32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_get_rows_i32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_get_rows_q8_0_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_get_rows_q8_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_get_rows_q8_0_pipeline) {
fprintf(stderr, "ds4: Metal kernel_get_rows_q8_0_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_get_rows_q4_0_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_get_rows_q4_0_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_get_rows_q4_0_pipeline) {
fprintf(stderr, "ds4: Metal kernel_get_rows_q4_0_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_get_rows_q4_K_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_get_rows_q4_K_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_get_rows_q4_K_pipeline) {
fprintf(stderr, "ds4: Metal kernel_get_rows_q4_K_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_repeat_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_repeat_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_repeat_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_repeat_f32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_repeat_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_set_rows_f32_i32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_set_rows_f32_i32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_set_rows_f32_i32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_set_rows_f32_i32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_concat"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_concat function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_concat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_concat_pipeline) {
fprintf(stderr, "ds4: Metal kernel_concat pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_cpy_f32_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_cpy_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_cpy_f32_f32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_cpy_f32_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_cpy_f32_f16"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_cpy_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_cpy_f32_f16_pipeline) {
fprintf(stderr, "ds4: Metal kernel_cpy_f32_f16 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_cpy_contig_f32_f16_4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_cpy_contig_f32_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_cpy_contig_f32_f16_pipeline) {
fprintf(stderr, "ds4: Metal kernel_cpy_contig_f32_f16_4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_cpy_f16_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_cpy_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_cpy_f16_f32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_cpy_f16_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_cpy_f16_f16"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_cpy_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_cpy_f16_f16_pipeline) {
fprintf(stderr, "ds4: Metal kernel_cpy_f16_f16 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f32_4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_cpy_contig_f16_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_cpy_contig_f16_f32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f32_4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_cpy_contig_f16_f16_bits_4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_cpy_contig_f16_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_cpy_contig_f16_f16_pipeline) {
fprintf(stderr, "ds4: Metal kernel_cpy_contig_f16_f16_bits_4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_dsv4_flash_kv_stage_f16"];
if (fn) {
g_flash_kv_stage_f16_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_flash_kv_stage_f16_pipeline) {
fprintf(stderr,
"ds4: optional Metal gathered KV staging pipeline unavailable: %s\n",
[[error localizedDescription] UTF8String]);
}
} else {
fprintf(stderr,
"ds4: optional Metal gathered KV staging kernel unavailable\n");
}
fn = [library newFunctionWithName:@"kernel_dsv4_fp8_kv_quantize_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_fp8_kv_quantize_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_fp8_kv_quantize_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_fp8_kv_quantize_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_indexer_hadamard_fp4_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_indexer_qat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_indexer_qat_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_hadamard_fp4_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_kv_fp8_store_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_kv_fp8_store_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_kv_fp8_store_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_kv_fp8_store_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_kv_rope_fp8_store_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_kv_rope_fp8_store_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
g_dsv4_kv_rope_fp8_store_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_kv_rope_fp8_store_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_kv_rope_fp8_store_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_ratio4_shift_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_ratio4_shift_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_ratio4_shift_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_ratio4_shift_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_swiglu_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_swiglu_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_swiglu_flat_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
g_swiglu_flat_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_swiglu_flat_pipeline) {
fprintf(stderr, "ds4: Metal kernel_swiglu_flat_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum6_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_moe_sum8_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_sum8_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_sum8_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_moe_sum8_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *bin_constants = [[MTLFunctionConstantValues alloc] init];
int16_t bin_op = 0;
int16_t bin_f = 1;
bool bin_rb = false;
bool bin_cb = false;
[bin_constants setConstantValue:&bin_op type:MTLDataTypeShort atIndex:1300];
[bin_constants setConstantValue:&bin_f type:MTLDataTypeShort atIndex:1301];
[bin_constants setConstantValue:&bin_rb type:MTLDataTypeBool atIndex:1302];
[bin_constants setConstantValue:&bin_cb type:MTLDataTypeBool atIndex:1303];
error = nil;
fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32"
constantValues:bin_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_add_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_add_pipeline) {
fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_add2_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_add2_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
g_add2_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_add2_pipeline) {
fprintf(stderr, "ds4: Metal kernel_add2_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_add3_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_add3_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
g_add3_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_add3_pipeline) {
fprintf(stderr, "ds4: Metal kernel_add3_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *bin_mul_constants = [[MTLFunctionConstantValues alloc] init];
int16_t bin_mul_plain_op = 2;
int16_t bin_mul_plain_f = 1;
bool bin_mul_plain_rb = false;
bool bin_mul_plain_cb = false;
[bin_mul_constants setConstantValue:&bin_mul_plain_op type:MTLDataTypeShort atIndex:1300];
[bin_mul_constants setConstantValue:&bin_mul_plain_f type:MTLDataTypeShort atIndex:1301];
[bin_mul_constants setConstantValue:&bin_mul_plain_rb type:MTLDataTypeBool atIndex:1302];
[bin_mul_constants setConstantValue:&bin_mul_plain_cb type:MTLDataTypeBool atIndex:1303];
error = nil;
fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32"
constantValues:bin_mul_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_mul_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_mul_pipeline) {
fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *bin_mul_scalar_constants = [[MTLFunctionConstantValues alloc] init];
int16_t bin_mul_op = 2;
int16_t bin_mul_f = 1;
bool bin_mul_rb = false;
bool bin_mul_cb = true;
[bin_mul_scalar_constants setConstantValue:&bin_mul_op type:MTLDataTypeShort atIndex:1300];
[bin_mul_scalar_constants setConstantValue:&bin_mul_f type:MTLDataTypeShort atIndex:1301];
[bin_mul_scalar_constants setConstantValue:&bin_mul_rb type:MTLDataTypeBool atIndex:1302];
[bin_mul_scalar_constants setConstantValue:&bin_mul_cb type:MTLDataTypeBool atIndex:1303];
error = nil;
fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32"
constantValues:bin_mul_scalar_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_bin_mul_scalar_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_bin_mul_scalar_pipeline) {
fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 mul-scalar pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *bin_div_row_constants = [[MTLFunctionConstantValues alloc] init];
int16_t bin_div_op = 3;
int16_t bin_div_f = 1;
bool bin_div_rb = false;
bool bin_div_cb = true;
[bin_div_row_constants setConstantValue:&bin_div_op type:MTLDataTypeShort atIndex:1300];
[bin_div_row_constants setConstantValue:&bin_div_f type:MTLDataTypeShort atIndex:1301];
[bin_div_row_constants setConstantValue:&bin_div_rb type:MTLDataTypeBool atIndex:1302];
[bin_div_row_constants setConstantValue:&bin_div_cb type:MTLDataTypeBool atIndex:1303];
error = nil;
fn = [library newFunctionWithName:@"kernel_bin_fuse_f32_f32_f32"
constantValues:bin_div_row_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_bin_div_row_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_bin_div_row_pipeline) {
fprintf(stderr, "ds4: Metal kernel_bin_fuse_f32_f32_f32 div-row pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_rms_norm_mul_f32_4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_rms_norm_pipeline) {
fprintf(stderr, "ds4: Metal kernel_rms_norm_mul_f32_4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_rms_norm_f32_4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_rms_norm_plain_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_rms_norm_plain_pipeline) {
fprintf(stderr, "ds4: Metal kernel_rms_norm_f32_4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_add_rms_norm_mul_f32_4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_add_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_add_rms_norm_pipeline) {
fprintf(stderr, "ds4: Metal kernel_add_rms_norm_mul_f32_4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_qkv_rms_norm_f32_4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_qkv_rms_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_qkv_rms_norm_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_qkv_rms_norm_f32_4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_head_rms_norm_rope_tail_f32"];
if (!fn) {
fprintf(stderr,
"ds4: Metal kernel_dsv4_head_rms_norm_rope_tail_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
g_dsv4_head_rms_norm_rope_tail_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_head_rms_norm_rope_tail_pipeline) {
fprintf(stderr,
"ds4: Metal kernel_dsv4_head_rms_norm_rope_tail_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_use_dsv4_head_rms_norm_rope_tail_pipeline =
ds4_gpu_device_is_pre_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_PRE_M5_HEAD_RMS_ROPE_PIPELINE_STATIC") == NULL;
MTLFunctionConstantValues *moe_mv_id_constants = [[MTLFunctionConstantValues alloc] init];
int16_t moe_mv_id_nsg = 2;
[moe_mv_id_constants setConstantValue:&moe_mv_id_nsg type:MTLDataTypeShort atIndex:600];
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_iq2_xxs_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_iq2_xxs_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_iq2_xxs_pair_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_pair_swiglu_pack2_overlap_f32"
constantValues:moe_mv_id_constants
error:&error];
if (fn) {
g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
}
if (!g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline) {
fprintf(stderr,
"ds4: optional Metal IQ2 pair pack2 pipeline unavailable: %s\n",
error ? [[error localizedDescription] UTF8String] :
"function not found");
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_q2_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_q2_k_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_q2_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_q2_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q2_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_iq2_xxs_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_iq2_xxs_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_iq2_xxs_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q2_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_slots6_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_slots6_q2_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q2_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_addr_iq2_xxs_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_addr_iq2_xxs_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_addr_q2_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_addr_q2_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_iq2_xxs_pair_swiglu_masked_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q2_K_sum6_masked_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q2_K_sum6_masked_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_stream_expert_cache_validate"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_stream_expert_cache_validate_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_stream_expert_cache_validate_pipeline) {
fprintf(stderr, "ds4: Metal kernel_stream_expert_cache_validate pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_q4_k_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_q4_k_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_q4_k_pair_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_q4_k_pair_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_id_q4_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_id_q4_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_id_q4_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_group_q4_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_group_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_group_q4_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group_q4_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_group6_q4_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_group6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_group6_q4_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group6_q4_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_group8_q4_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_group8_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_group8_q4_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group8_q4_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_id_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_group24_q4_k_id_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_group24_q4_k_id_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_id_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_group24_q4_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_group24_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_group24_q4_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_group24_q4_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_slots6_q4_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_slots6_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_slots6_q4_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_slots6_q4_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_id_mxfp4_pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_f32", 2);
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_pair_swiglu_f32", 2);
g_moe_mul_mv_id_mxfp4_sum6_pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_sum6_f32", 2);
/* Single-simdgroup variants of the MXFP4 decode MoE kernels. The
* per-row per-lane arithmetic is identical; only the row-to-simdgroup
* mapping changes, so outputs remain bit-exact. Smaller threadgroups
* measurably improve GPU scheduling granularity on these
* latency-bound small-grid matvecs. Pipelines are built
* unconditionally; the dispatch-time selection gate reads the
* environment so A/B harnesses can toggle it after engine creation. */
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1 =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_pair_swiglu_f32", 1);
g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1 =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_mxfp4_sum6_f32", 1);
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_pair_swiglu_f32", 1);
g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_sum6_f32", 1);
g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_pair_swiglu_fixed_route_f32", 1);
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_sum6_fixed_route_f32", 1);
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_f32", 1);
g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_f32", 1);
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_f32", 1);
g_moe_mul_mv_id_mxfp4_pair_swiglu_tp_static_pipeline_nsg1 =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_pair_swiglu_tp_static_f32", 1);
g_moe_mul_mv_id_mxfp4_sum6_tp_full_rows_static_pipeline_nsg1 =
ds4_gpu_new_mul_mv_tg_multiple_pipeline(
"kernel_mul_mv_id_mxfp4_sum6_tp_full_rows_static_f32", 1);
if (!g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1 ||
!g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1 ||
!g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple ||
!g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple ||
!g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 ||
!g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 ||
!g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 ||
!g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 ||
!g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 ||
!g_moe_mul_mv_id_mxfp4_pair_swiglu_tp_static_pipeline_nsg1 ||
!g_moe_mul_mv_id_mxfp4_sum6_tp_full_rows_static_pipeline_nsg1) {
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_slots6_mxfp4_pair_swiglu_f32", 2);
g_moe_mul_mv_slots6_mxfp4_sum6_pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_slots6_mxfp4_sum6_f32", 2);
if (!g_moe_mul_mv_id_mxfp4_pipeline ||
!g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline ||
!g_moe_mul_mv_id_mxfp4_sum6_pipeline ||
!g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline ||
!g_moe_mul_mv_slots6_mxfp4_sum6_pipeline) {
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_q4_gather_slots6"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_q4_gather_slots6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_q4_gather_slots6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_q4_gather_slots6 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_table_q4_pair_gate_encoder = [fn newArgumentEncoderWithBufferIndex:2];
g_moe_table_q4_pair_up_encoder = [fn newArgumentEncoderWithBufferIndex:3];
if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder) {
fprintf(stderr, "ds4: Metal Q4 expert-table pair argument encoder creation failed\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_pair_swiglu_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_table_q4_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_table_q4_sum_down_encoder = [fn newArgumentEncoderWithBufferIndex:1];
if (!g_moe_table_q4_sum_down_encoder) {
fprintf(stderr, "ds4: Metal Q4 expert-table down argument encoder creation failed\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_moe_mul_mv_table_q4_k_sum6_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_table_q4_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_table_q4_K_sum6_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_pair_swiglu_f32"
constantValues:moe_mv_id_constants
error:&error];
if (fn) {
g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 pipeline unavailable: %s\n",
[[error localizedDescription] UTF8String]);
}
} else {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_pair_swiglu_f32 function unavailable: %s\n",
[[error localizedDescription] UTF8String]);
}
error = nil;
fn = [library newFunctionWithName:@"kernel_mul_mv_addr_q4_K_sum6_f32"
constantValues:moe_mv_id_constants
error:&error];
if (fn) {
g_moe_mul_mv_addr_q4_k_sum6_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_moe_mul_mv_addr_q4_k_sum6_pipeline) {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 pipeline unavailable: %s\n",
[[error localizedDescription] UTF8String]);
}
} else {
fprintf(stderr, "ds4: Metal kernel_mul_mv_addr_q4_K_sum6_f32 function unavailable: %s\n",
[[error localizedDescription] UTF8String]);
}
fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_rope_tail_batch_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_rope_tail_batch_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_rope_tail_inplace_pair_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_rope_tail_inplace_pair_pipeline) {
fprintf(stderr,
"ds4: Metal kernel_dsv4_rope_tail_f32_inplace_pair pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_shared4"];
if (!fn) {
fprintf(stderr,
"ds4: optional Metal shared-head RoPE kernel unavailable; using per-head path\n");
} else {
g_rope_tail_inplace_pair_shared4_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_rope_tail_inplace_pair_shared4_pipeline) {
fprintf(stderr,
"ds4: optional Metal shared-head RoPE pipeline unavailable; using per-head path: %s\n",
[[error localizedDescription] UTF8String]);
}
}
error = nil;
fn = [library newFunctionWithName:@"kernel_dsv4_rope_tail_f32_inplace_pair_affine"];
if (fn) {
g_rope_tail_inplace_pair_affine_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_rope_tail_inplace_pair_affine_pipeline) {
fprintf(stderr,
"ds4: optional Metal affine-position RoPE pair pipeline unavailable: %s\n",
[[error localizedDescription] UTF8String]);
}
} else {
fprintf(stderr,
"ds4: optional Metal affine-position RoPE pair kernel unavailable\n");
}
error = nil;
fn = [library newFunctionWithName:@"kernel_dsv4_softmax_pool"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_softmax_pool_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_softmax_pool_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_softmax_pool pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_soft_max_f32"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_soft_max_f32 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_soft_max_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_soft_max_f32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_soft_max_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_soft_max_f32_4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_soft_max_f32_4_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_soft_max_f32_4_pipeline) {
fprintf(stderr, "ds4: Metal kernel_soft_max_f32_4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_argsort_f32_i32_desc"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_argsort_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_argsort_f32_i32_desc_pipeline) {
fprintf(stderr, "ds4: Metal kernel_argsort_f32_i32_desc pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_argsort_merge_f32_i32_desc"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_argsort_merge_f32_i32_desc_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_argsort_merge_f32_i32_desc_pipeline) {
fprintf(stderr, "ds4: Metal kernel_argsort_merge_f32_i32_desc pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *sum_rows_constants = [[MTLFunctionConstantValues alloc] init];
int16_t sum_rows_op = 10;
[sum_rows_constants setConstantValue:&sum_rows_op type:MTLDataTypeShort atIndex:1400];
error = nil;
fn = [library newFunctionWithName:@"kernel_sum_rows_f32_f32"
constantValues:sum_rows_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_sum_rows_f32_f32_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_sum_rows_f32_f32_pipeline) {
fprintf(stderr, "ds4: Metal kernel_sum_rows_f32_f32 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_topk_mask_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_topk_mask_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_topk_mask_scatter"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_topk_mask_scatter_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_topk_mask_scatter_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_topk_mask_scatter pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_indexer_weighted_sum"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_indexer_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_dsv4_indexer_weighted_sum_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_indexer_weighted_sum pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_sinkhorn"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_hc_split_sinkhorn_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_hc_split_sinkhorn_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_sinkhorn pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_hc_split_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_hc_split_weighted_sum_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_hc_split_weighted_sum_norm4"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_hc_split_weighted_sum_norm_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_hc_split_weighted_sum_norm_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_split_weighted_sum_norm4 pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_hc_weighted_sum"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_hc_weighted_sum_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_hc_weighted_sum_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_weighted_sum pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_dsv4_output_hc_weights4"];
if (fn) {
g_output_hc_weights4_pipeline =
[g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_output_hc_weights4_pipeline) {
fprintf(stderr,
"ds4: optional Metal output HC weights4 pipeline unavailable: %s\n",
[[error localizedDescription] UTF8String]);
}
} else {
fprintf(stderr,
"ds4: optional Metal output HC weights4 kernel unavailable\n");
}
MTLFunctionConstantValues *unary_sigmoid_constants = [[MTLFunctionConstantValues alloc] init];
int16_t unary_sigmoid_op = 102;
bool unary_cnt = false;
[unary_sigmoid_constants setConstantValue:&unary_sigmoid_op type:MTLDataTypeShort atIndex:1200];
[unary_sigmoid_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201];
error = nil;
fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
constantValues:unary_sigmoid_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_unary_sigmoid_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_unary_sigmoid_pipeline) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sigmoid pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *unary_silu_constants = [[MTLFunctionConstantValues alloc] init];
int16_t unary_silu_op = 106;
[unary_silu_constants setConstantValue:&unary_silu_op type:MTLDataTypeShort atIndex:1200];
[unary_silu_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201];
error = nil;
fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
constantValues:unary_silu_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_unary_silu_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_unary_silu_pipeline) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 silu pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *unary_softplus_constants = [[MTLFunctionConstantValues alloc] init];
int16_t unary_softplus_op = 115;
[unary_softplus_constants setConstantValue:&unary_softplus_op type:MTLDataTypeShort atIndex:1200];
[unary_softplus_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201];
error = nil;
fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
constantValues:unary_softplus_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_unary_softplus_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_unary_softplus_pipeline) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 softplus pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *unary_sqrt_constants = [[MTLFunctionConstantValues alloc] init];
int16_t unary_sqrt_op = 14;
[unary_sqrt_constants setConstantValue:&unary_sqrt_op type:MTLDataTypeShort atIndex:1200];
[unary_sqrt_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201];
error = nil;
fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
constantValues:unary_sqrt_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_unary_sqrt_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_unary_sqrt_pipeline) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 sqrt pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *unary_clamp_constants = [[MTLFunctionConstantValues alloc] init];
int16_t unary_clamp_op = 12;
[unary_clamp_constants setConstantValue:&unary_clamp_op type:MTLDataTypeShort atIndex:1200];
[unary_clamp_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201];
error = nil;
fn = [library newFunctionWithName:@"kernel_unary_f32_f32"
constantValues:unary_clamp_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_unary_clamp_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_unary_clamp_pipeline) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32 clamp pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *unary_scale_constants = [[MTLFunctionConstantValues alloc] init];
int16_t unary_scale_op = 10;
[unary_scale_constants setConstantValue:&unary_scale_op type:MTLDataTypeShort atIndex:1200];
[unary_scale_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201];
error = nil;
fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
constantValues:unary_scale_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_unary_scale_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_unary_scale_pipeline) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 scale pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
MTLFunctionConstantValues *unary_fill_constants = [[MTLFunctionConstantValues alloc] init];
int16_t unary_fill_op = 11;
[unary_fill_constants setConstantValue:&unary_fill_op type:MTLDataTypeShort atIndex:1200];
[unary_fill_constants setConstantValue:&unary_cnt type:MTLDataTypeBool atIndex:1201];
error = nil;
fn = [library newFunctionWithName:@"kernel_unary_f32_f32_4"
constantValues:unary_fill_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_unary_fill_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_unary_fill_pipeline) {
fprintf(stderr, "ds4: Metal kernel_unary_f32_f32_4 fill pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
error = nil;
fn = [library newFunctionWithName:@"kernel_unary_f16_f16"
constantValues:unary_fill_constants
error:&error];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill function not found: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_unary_fill_f16_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_unary_fill_f16_pipeline) {
fprintf(stderr, "ds4: Metal kernel_unary_f16_f16 fill pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
fn = [library newFunctionWithName:@"kernel_dsv4_hc_expand"];
if (!fn) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand function not found\n");
g_queue = nil;
g_device = nil;
return 0;
}
g_hc_expand_pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!g_hc_expand_pipeline) {
fprintf(stderr, "ds4: Metal kernel_dsv4_hc_expand pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
g_queue = nil;
g_device = nil;
return 0;
}
g_dsv4_indexer_score_one_direct_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexer_score_one_direct");
g_dsv4_compressor_store_one_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_compressor_store_one");
g_dsv4_compressor_pack_ratio4_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_compressor_pack_ratio4");
g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline =
ds4_gpu_get_pipeline(
"kernel_dsv4_compressor_pack_ratio4_decode_ggml");
g_dsv4_compressor_exact_softmax_product_pipeline =
ds4_gpu_get_pipeline(
"kernel_dsv4_compressor_exact_softmax_product_ratio4");
g_dsv4_compressor_exact_pool_ratio4_pipeline =
ds4_gpu_get_pipeline(
"kernel_dsv4_compressor_exact_pool_ratio4_decode_ggml");
g_dsv4_softmax_pool_ratio4_direct_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_softmax_pool_ratio4_direct");
g_rms_norm_scale_pipeline =
ds4_gpu_get_pipeline("kernel_rms_norm_scale_f32_4");
g_dsv4_sort_i32_rows_asc_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_sort_i32_rows_asc");
g_dsv4_indexed_attention_heads8_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8");
g_dsv4_indexed_attention_heads8_rb16_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_rb16");
g_dsv4_indexed_attention_heads16_dual_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads16_dual");
g_dsv4_indexed_attention_heads8_split_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_split");
g_dsv4_indexed_attention_heads8_split_reduce_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_indexed_mixed_attention_heads8_split_reduce");
g_dsv4_softplus_sqrt_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_softplus_sqrt_f32_4");
g_dsv4_router_finalize_one_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one");
g_dsv4_router_finalize_one_simd_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_one_simd");
g_dsv4_router_finalize_weights_one_simd_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_router_finalize_weights_one_simd");
g_dsv4_router_transform_finalize_weights_one_simd_pipeline =
ds4_gpu_get_pipeline(
"kernel_dsv4_router_transform_finalize_weights_one_simd");
g_dsv4_router_select_visual_batch_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_router_select_visual_batch");
g_dsv4_router_weights_one_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_router_weights_one");
g_glm_router_select_one_pipeline =
ds4_gpu_get_pipeline("kernel_glm_router_select_one");
g_glm_kv_lora_rms_norm_pipeline =
ds4_gpu_get_pipeline("kernel_glm_kv_lora_rms_norm");
g_glm_k_b_project_pipeline =
ds4_gpu_get_pipeline("kernel_glm_k_b_project_q8_0");
g_glm_store_compact_kv_pipeline =
ds4_gpu_get_pipeline("kernel_glm_store_compact_kv");
g_glm_qkv_norm_store_compact_kv_pipeline =
ds4_gpu_get_pipeline("kernel_glm_qkv_norm_store_compact_kv");
g_glm_store_indexer_k_pipeline =
ds4_gpu_get_pipeline("kernel_glm_store_indexer_k");
g_glm53_indexer_pool_update_pipeline =
ds4_gpu_get_pipeline("kernel_glm53_indexer_pool_update");
g_glm_build_kv_cache_pipeline =
ds4_gpu_get_pipeline("kernel_glm_build_kv_cache");
g_glm_build_kv_cache_decode_group4_pipeline =
ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_decode_group4");
g_glm_build_kv_cache_flash_pipeline =
ds4_gpu_get_pipeline("kernel_glm_build_kv_cache_flash");
g_glm_attention_full_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_full");
g_glm_fill_selected_range_pipeline =
ds4_gpu_get_pipeline("kernel_glm_fill_selected_range");
g_glm_fill_selected_range_batch_pipeline =
ds4_gpu_get_pipeline("kernel_glm_fill_selected_range_batch");
g_glm53_expand_pool_selection_pipeline =
ds4_gpu_get_pipeline("kernel_glm53_expand_pool_selection");
g_glm_indexer_rope_tail_pipeline =
ds4_gpu_get_pipeline("kernel_glm_indexer_rope_tail_f32");
g_glm_indexer_score_one_pipeline =
ds4_gpu_get_pipeline("kernel_glm_indexer_score_one");
g_glm_indexer_score_one_direct_pipeline =
ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_direct");
g_glm_indexer_score_one_wide_pipeline =
ds4_gpu_get_pipeline("kernel_glm_indexer_score_one_wide");
g_glm_indexer_scores_batch_pipeline =
ds4_gpu_get_pipeline("kernel_glm_indexer_scores_batch");
g_glm_indexer_scores_tiled_pipeline =
ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled");
g_glm_indexer_scores_tiled_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_indexer_scores_tiled_f32");
g_glm_qk_lowrank_pipeline =
ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0");
g_glm_qk_lowrank_glm52_pipeline =
ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52");
g_glm_qk_lowrank_glm52_sg_pipeline =
ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_glm52_sg");
g_glm_qk_lowrank_batch_pipeline =
ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch");
g_glm_qk_lowrank_batch_glm52_t4_pipeline =
ds4_gpu_get_pipeline("kernel_glm_qk_lowrank_q8_0_batch_glm52_t4");
g_glm_value_project_q8_0_pipeline =
ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0");
g_glm_value_project_q8_0_batch_heads_pipeline =
ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads");
g_glm_value_project_q8_0_batch_heads_mma_pipeline =
ds4_gpu_get_pipeline("kernel_glm_value_project_q8_0_batch_heads_mma");
g_glm_attention_indexed_decode_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode");
g_glm_attention_indexed_decode_split_group8_partial_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial");
g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads");
g_glm_attention_indexed_decode_split_group8_reduce_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce");
g_glm_attention_indexed_decode_split_group8_reduce16_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_decode_split_group8_reduce16");
g_glm_attention_indexed_batch_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch");
g_glm_attention_indexed_batch_group2_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group2");
g_glm_attention_indexed_batch_q2_group4_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_q2_group4");
g_glm_attention_indexed_batch_group8_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_group8");
g_glm_attention_indexed_batch_lora_group8_vec_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec");
g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid");
g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads");
g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal");
g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline =
ds4_gpu_get_pipeline("kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads");
g_glm_q4_k_pair_swiglu_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu_f32");
g_glm_q4_k_pair_swiglu2_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_f32");
g_glm_q4_k_pair_swiglu4_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu4_f32");
g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_f32");
g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_pair_swiglu2_mapped_row_f32");
g_glm_q2_k_pair_swiglu_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q2_K_pair_swiglu_f32");
g_glm_q2_k_addr_pair_swiglu2_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32");
g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_pair_swiglu2_f32_masked");
g_glm_q4_k_addr_pair_swiglu_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32");
g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_pair_swiglu_f32_masked");
g_glm_q2_k_down_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q2_K_down_f32");
g_glm_q4_k_down_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_down_simd_f32");
g_glm_q2_k_addr_down_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q2_K_addr_down_f32");
g_glm_q4_k_addr_down_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q4_K_addr_down_simd_f32");
g_glm_q5_k_pair_swiglu_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_f32");
g_glm_q5_k_pair_swiglu_mapped_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_f32");
g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q5_K_pair_swiglu_mapped_row_f32");
g_glm_q5_k_down_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q5_K_down_f32");
g_glm_q6_k_down_f32_pipeline =
ds4_gpu_get_pipeline("kernel_glm_q6_K_down_f32");
g_dsv4_router_weights_batch_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_router_weights_batch");
g_dsv4_hc_expand4_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_hc_expand4");
if (!g_dsv4_indexer_score_one_direct_pipeline ||
!g_dsv4_compressor_store_one_pipeline ||
!g_dsv4_sort_i32_rows_asc_pipeline ||
!g_dsv4_indexed_attention_heads8_pipeline ||
!g_dsv4_indexed_attention_heads8_rb16_pipeline ||
!g_dsv4_indexed_attention_heads16_dual_pipeline ||
!g_dsv4_indexed_attention_heads8_split_pipeline ||
!g_dsv4_indexed_attention_heads8_split_reduce_pipeline ||
!g_dsv4_softplus_sqrt_pipeline ||
!g_dsv4_router_finalize_one_pipeline ||
!g_dsv4_router_select_visual_batch_pipeline ||
!g_dsv4_router_weights_one_pipeline ||
!g_glm_router_select_one_pipeline ||
!g_glm_kv_lora_rms_norm_pipeline ||
!g_glm_k_b_project_pipeline ||
!g_glm_store_compact_kv_pipeline ||
!g_glm_qkv_norm_store_compact_kv_pipeline ||
!g_glm_store_indexer_k_pipeline ||
!g_glm53_indexer_pool_update_pipeline ||
!g_glm_build_kv_cache_pipeline ||
!g_glm_build_kv_cache_decode_group4_pipeline ||
!g_glm_build_kv_cache_flash_pipeline ||
!g_glm_attention_full_pipeline ||
!g_glm_fill_selected_range_pipeline ||
!g_glm_fill_selected_range_batch_pipeline ||
!g_glm53_expand_pool_selection_pipeline ||
!g_glm_indexer_rope_tail_pipeline ||
!g_glm_indexer_score_one_pipeline ||
!g_glm_indexer_score_one_direct_pipeline ||
!g_glm_indexer_score_one_wide_pipeline ||
!g_glm_indexer_scores_batch_pipeline ||
!g_glm_indexer_scores_tiled_pipeline ||
!g_glm_indexer_scores_tiled_f32_pipeline ||
!g_glm_qk_lowrank_pipeline ||
!g_glm_qk_lowrank_glm52_pipeline ||
!g_glm_qk_lowrank_glm52_sg_pipeline ||
!g_glm_qk_lowrank_batch_pipeline ||
!g_glm_qk_lowrank_batch_glm52_t4_pipeline ||
!g_glm_value_project_q8_0_pipeline ||
!g_glm_value_project_q8_0_batch_heads_pipeline ||
!g_glm_value_project_q8_0_batch_heads_mma_pipeline ||
!g_glm_attention_indexed_decode_pipeline ||
!g_glm_attention_indexed_decode_split_group8_partial_pipeline ||
!g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline ||
!g_glm_attention_indexed_decode_split_group8_reduce_pipeline ||
!g_glm_attention_indexed_decode_split_group8_reduce16_pipeline ||
!g_glm_attention_indexed_batch_pipeline ||
!g_glm_attention_indexed_batch_group2_pipeline ||
!g_glm_attention_indexed_batch_q2_group4_pipeline ||
!g_glm_attention_indexed_batch_group8_pipeline ||
!g_glm_attention_indexed_batch_lora_group8_vec_pipeline ||
!g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline ||
!g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline ||
!g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline ||
!g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline ||
!g_glm_q4_k_pair_swiglu_f32_pipeline ||
!g_glm_q4_k_pair_swiglu2_f32_pipeline ||
!g_glm_q4_k_pair_swiglu4_f32_pipeline ||
!g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline ||
!g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline ||
!g_glm_q2_k_pair_swiglu_f32_pipeline ||
!g_glm_q2_k_addr_pair_swiglu2_f32_pipeline ||
!g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline ||
!g_glm_q4_k_addr_pair_swiglu_f32_pipeline ||
!g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline ||
!g_glm_q2_k_down_f32_pipeline ||
!g_glm_q4_k_down_f32_pipeline ||
!g_glm_q2_k_addr_down_f32_pipeline ||
!g_glm_q4_k_addr_down_f32_pipeline ||
!g_glm_q5_k_pair_swiglu_f32_pipeline ||
!g_glm_q5_k_pair_swiglu_mapped_f32_pipeline ||
!g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline ||
!g_glm_q5_k_down_f32_pipeline ||
!g_glm_q6_k_down_f32_pipeline ||
!g_dsv4_hc_expand4_pipeline) {
g_queue = nil;
g_device = nil;
return 0;
}
g_initialized = 1;
}
return 1;
}
int ds4_gpu_test_mxfp4_down_half_lut(uint16_t *legacy_bits,
uint16_t *lut_bits) {
if (!legacy_bits || !lut_bits) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
const NSUInteger count = 4096u;
const NSUInteger bytes = count * sizeof(uint16_t);
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_test_mxfp4_down_half_lut");
id<MTLBuffer> legacy = [g_device newBufferWithLength:bytes
options:MTLResourceStorageModeShared];
id<MTLBuffer> lut = [g_device newBufferWithLength:bytes
options:MTLResourceStorageModeShared];
id<MTLCommandBuffer> cb = ds4_gpu_new_command_buffer();
if (!pipeline || !legacy || !lut || !cb) {
fprintf(stderr, "ds4: Metal MXFP4 half-LUT test setup failed\n");
return 0;
}
id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBuffer:legacy offset:0 atIndex:0];
[enc setBuffer:lut offset:0 atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake((count + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)];
[enc endEncoding];
[cb commit];
if (!ds4_gpu_wait_command_buffer(cb, "MXFP4 half-LUT raw-bit test")) {
return 0;
}
memcpy(legacy_bits, [legacy contents], bytes);
memcpy(lut_bits, [lut contents], bytes);
}
return 1;
}
void ds4_gpu_test_set_flags(uint32_t flags) {
g_test_flags = flags;
}
ds4_gpu_tensor *ds4_gpu_tensor_alloc(uint64_t bytes) {
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];
if (!tensor.buffer) {
return NULL;
}
tensor.offset = 0;
tensor.bytes = bytes;
tensor.owner = 1;
uint64_t live_snap = 0;
uint64_t peak_snap = 0;
pthread_mutex_lock(&g_tensor_mu);
const int tracked = ds4_gpu_tensor_track_alloc_locked(
(__bridge const void *)tensor,
bytes,
&live_snap,
&peak_snap);
pthread_mutex_unlock(&g_tensor_mu);
if (!tracked) {
fprintf(stderr, "ds4: failed to track Metal tensor allocation\n");
tensor.buffer = nil;
return NULL;
}
if (ds4_gpu_trace_allocs()) {
fprintf(stderr,
"ds4: Metal tensor alloc %.3f MiB live %.3f MiB peak %.3f MiB\n",
(double)bytes / (1024.0 * 1024.0),
(double)live_snap / (1024.0 * 1024.0),
(double)peak_snap / (1024.0 * 1024.0));
}
return (__bridge_retained ds4_gpu_tensor *)tensor;
}
}
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);
if (offset > base_obj.bytes || bytes > base_obj.bytes - offset) return NULL;
if (base_obj.offset > UINT64_MAX - offset) return NULL;
const uint64_t absolute_offset = base_obj.offset + offset;
if (absolute_offset > (uint64_t)NSUIntegerMax) return NULL;
@autoreleasepool {
DS4MetalTensor *view = [DS4MetalTensor new];
view.buffer = base_obj.buffer;
view.offset = absolute_offset;
view.bytes = bytes;
view.owner = 0;
pthread_mutex_lock(&g_tensor_mu);
const int tracked = ds4_gpu_tensor_track_view_locked((__bridge const void *)view);
pthread_mutex_unlock(&g_tensor_mu);
if (!tracked) {
fprintf(stderr, "ds4: failed to track Metal tensor view\n");
view.buffer = nil;
return NULL;
}
return (__bridge_retained ds4_gpu_tensor *)view;
}
}
void ds4_gpu_tensor_free(ds4_gpu_tensor *tensor) {
if (!tensor) return;
@autoreleasepool {
uint8_t owner = 0;
uint64_t bytes = 0;
uint64_t live_snap = 0;
uint64_t peak_snap = 0;
if (!ds4_gpu_tensor_prepare_free(tensor,
&owner,
&bytes,
&live_snap,
&peak_snap)) {
return;
}
DS4MetalTensor *obj = (__bridge_transfer DS4MetalTensor *)tensor;
if (owner) {
if (ds4_gpu_trace_allocs()) {
fprintf(stderr,
"ds4: Metal tensor free %.3f MiB live %.3f MiB peak %.3f MiB\n",
(double)bytes / (1024.0 * 1024.0),
(double)live_snap / (1024.0 * 1024.0),
(double)peak_snap / (1024.0 * 1024.0));
}
}
obj.buffer = nil;
obj.offset = 0;
obj.bytes = 0;
obj.owner = 0;
}
}
uint64_t ds4_gpu_tensor_bytes(const ds4_gpu_tensor *tensor) {
if (!tensor) return 0;
const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor);
return obj.bytes;
}
void *ds4_gpu_tensor_contents(ds4_gpu_tensor *tensor) {
if (!tensor) return NULL;
DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor);
return (uint8_t *)[obj.buffer contents] + obj.offset;
}
int ds4_gpu_tensor_fill_f32(ds4_gpu_tensor *tensor, float value, uint64_t count) {
if (!tensor || count > ds4_gpu_tensor_bytes(tensor) / sizeof(float)) return 0;
float *p = ds4_gpu_tensor_contents(tensor);
if (!p && count != 0) return 0;
for (uint64_t i = 0; i < count; i++) p[i] = value;
return 1;
}
int ds4_gpu_tensor_write(ds4_gpu_tensor *tensor, uint64_t offset, const void *data, uint64_t bytes) {
if (!tensor || (!data && bytes != 0)) return 0;
DS4MetalTensor *obj = ds4_gpu_tensor_obj(tensor);
if (offset > obj.bytes || bytes > obj.bytes - offset) return 0;
if (bytes != 0) {
memcpy((uint8_t *)[obj.buffer contents] + obj.offset + offset, data, (size_t)bytes);
}
return 1;
}
int ds4_gpu_tensor_read(const ds4_gpu_tensor *tensor, uint64_t offset, void *data, uint64_t bytes) {
if (!tensor || (!data && bytes != 0)) return 0;
const DS4MetalTensor *obj = ds4_gpu_tensor_const_obj(tensor);
if (offset > obj.bytes || bytes > obj.bytes - offset) return 0;
if (bytes != 0) {
memcpy(data, (const uint8_t *)[obj.buffer contents] + obj.offset + offset, (size_t)bytes);
}
return 1;
}
int ds4_gpu_tensor_copy(ds4_gpu_tensor *dst, uint64_t dst_offset,
const ds4_gpu_tensor *src, uint64_t src_offset,
uint64_t bytes) {
if (!dst || !src) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
DS4MetalTensor *d = ds4_gpu_tensor_obj(dst);
const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src);
if (dst_offset > d.bytes || bytes > d.bytes - dst_offset) return 0;
if (src_offset > s.bytes || bytes > s.bytes - src_offset) return 0;
if (bytes == 0) return 1;
if (!g_batch_cb) return 0;
/* C inference callers have no surrounding autorelease pool. In particular,
* validation-layer blit bookkeeping otherwise accumulates across tokens. */
@autoreleasepool {
ds4_gpu_close_batch_encoder();
g_batch_has_work = YES;
id<MTLBlitCommandEncoder> blit = ds4_gpu_blit_encoder(g_batch_cb, "tensor_copy", 1u);
if (!blit) return 0;
[blit copyFromBuffer:s.buffer
sourceOffset:(NSUInteger)(s.offset + src_offset)
toBuffer:d.buffer
destinationOffset:(NSUInteger)(d.offset + dst_offset)
size:(NSUInteger)bytes];
[blit endEncoding];
}
return 1;
}
int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset,
const ds4_gpu_tensor *src, uint64_t src_offset,
uint64_t count) {
if (!dst || !src) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
DS4MetalTensor *d = ds4_gpu_tensor_obj(dst);
const DS4MetalTensor *s = ds4_gpu_tensor_const_obj(src);
if (count == 0) return 1;
if (count > UINT64_MAX / sizeof(float) ||
count > UINT64_MAX / sizeof(uint16_t)) {
return 0;
}
const uint64_t src_bytes = count * sizeof(float);
const uint64_t dst_bytes = count * sizeof(uint16_t);
if (src_offset > s.bytes || src_bytes > s.bytes - src_offset ||
dst_offset > d.bytes || dst_bytes > d.bytes - dst_offset) {
return 0;
}
@autoreleasepool {
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
uint64_t done = 0;
int ok = 1;
while (done < count && ok) {
uint64_t chunk64 = count - done;
if (chunk64 > UINT32_MAX) chunk64 = UINT32_MAX;
const uint32_t chunk = (uint32_t)chunk64;
ok = ds4_gpu_encode_cpy_f32_f16_1d(
cb,
s.buffer,
(NSUInteger)(s.offset + src_offset + done * sizeof(float)),
d.buffer,
(NSUInteger)(d.offset + dst_offset + done * sizeof(uint16_t)),
chunk);
done += chunk;
}
if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "tensor f32 to f16 copy");
return ok;
}
}
int ds4_gpu_pack_slot_rows_f32_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *slots,
uint32_t n_rows,
uint32_t width,
uint32_t n_slots,
uint32_t slot_cap) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !slots || n_rows == 0 || width == 0 || n_slots == 0 ||
slot_cap == 0 || n_rows > slot_cap) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> slotsbuf = ds4_gpu_tensor_buffer(slots);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
uint64_t row_bytes = 0;
uint64_t slot_plane_bytes = 0;
uint64_t slots_bytes = 0;
uint64_t out_rows = 0;
uint64_t out_bytes = 0;
if ((uint64_t)width > UINT64_MAX / sizeof(float)) return 0;
row_bytes = (uint64_t)width * sizeof(float);
if ((uint64_t)slot_cap > UINT64_MAX / row_bytes) return 0;
slot_plane_bytes = (uint64_t)slot_cap * row_bytes;
if ((uint64_t)n_slots > UINT64_MAX / slot_plane_bytes) return 0;
slots_bytes = (uint64_t)n_slots * slot_plane_bytes;
if ((uint64_t)n_rows > UINT64_MAX / n_slots) return 0;
out_rows = (uint64_t)n_rows * n_slots;
if (out_rows > UINT64_MAX / row_bytes) return 0;
out_bytes = out_rows * row_bytes;
if (!slotsbuf || !outbuf ||
ds4_gpu_tensor_bytes(slots) < slots_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal slot-row pack received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb,
slotsbuf,
ds4_gpu_tensor_offset(slots),
outbuf,
ds4_gpu_tensor_offset(out),
width,
n_slots,
n_rows,
sizeof(float),
slot_plane_bytes,
row_bytes,
row_bytes,
(uint64_t)n_slots * row_bytes)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "slot-row pack")) return 0;
}
return 1;
}
int ds4_gpu_begin_commands(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
/* A failed concurrent FFN must never affect the next command batch. */
ds4_gpu_parallel_ffn_reset_state(YES);
if (g_batch_cb) return 0;
/* Refresh once per command batch so same-engine A/B runs can toggle the
* static PSO without paying for environment lookups in every layer. */
g_use_dsv4_head_rms_norm_rope_tail_pipeline =
!g_ssd_streaming_mode &&
ds4_gpu_device_is_pre_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_PRE_M5_HEAD_RMS_ROPE_PIPELINE_STATIC") == NULL;
g_batch_cb_created_ms = ds4_gpu_now_ms();
g_batch_cb = ds4_gpu_new_command_buffer();
g_batch_has_work = NO;
if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
ds4_gpu_timeline_attach(g_batch_cb);
return g_batch_cb != nil;
}
int ds4_gpu_flush_encoder(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
ds4_gpu_parallel_ffn_reset_state(YES);
if (!g_batch_cb) return 0;
ds4_gpu_close_batch_encoder();
return 1;
}
int ds4_gpu_flush_commands(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
ds4_gpu_parallel_ffn_reset_state(YES);
if (!g_batch_cb) return 0;
ds4_gpu_close_batch_encoder();
id<MTLCommandBuffer> cb = g_batch_cb;
g_batch_cb = nil;
g_batch_has_work = NO;
[cb commit];
[g_pending_cbs addObject:cb];
ds4_gpu_stream_expert_cache_note_batch_committed();
g_batch_cb = ds4_gpu_new_command_buffer();
g_batch_has_work = NO;
if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
ds4_gpu_timeline_attach(g_batch_cb);
if (!g_batch_cb) {
(void)ds4_gpu_wait_pending_command_buffers("command batch");
[g_transient_buffers removeAllObjects];
return 0;
}
return 1;
}
int ds4_gpu_commands_active(void) {
return g_batch_cb != nil;
}
/* Exact M5 full-FFN overlap inside one concurrent compute encoder. Shared
* gate/up and routed IQ2 pair-SwiGLU launch together; explicit level barriers
* precede the routed Q2 and shared Q8 down consumers. */
static id<MTLComputePipelineState> g_parallel_q8_pipeline;
static id<MTLBuffer> g_parallel_q8_weight;
static id<MTLBuffer> g_parallel_q8_x;
static id<MTLBuffer> g_parallel_q8_out;
static NSUInteger g_parallel_q8_weight_offset;
static NSUInteger g_parallel_q8_x_offset;
static NSUInteger g_parallel_q8_out_offset;
static ds4_gpu_q8_0_matvec_args g_parallel_q8_args;
static id<MTLComputePipelineState> g_parallel_gate_up_pipeline;
static id<MTLBuffer> g_parallel_gate_weight;
static id<MTLBuffer> g_parallel_up_weight;
static id<MTLBuffer> g_parallel_gate_x;
static id<MTLBuffer> g_parallel_gate_out;
static id<MTLBuffer> g_parallel_up_out;
static id<MTLBuffer> g_parallel_mid_out;
static NSUInteger g_parallel_gate_weight_offset;
static NSUInteger g_parallel_up_weight_offset;
static NSUInteger g_parallel_gate_x_offset;
static NSUInteger g_parallel_gate_out_offset;
static NSUInteger g_parallel_up_out_offset;
static NSUInteger g_parallel_mid_out_offset;
static ds4_gpu_q8_0_matvec_args g_parallel_gate_up_args;
static float g_parallel_gate_up_clamp;
static NSUInteger g_parallel_gate_up_nsg;
static NSUInteger g_parallel_gate_up_nr0;
static NSUInteger g_parallel_gate_up_smem;
static int g_parallel_ffn_mode; /* 2: gate/up + down */
static int g_parallel_ffn_stage;
static BOOL g_parallel_q8_pending;
static BOOL g_parallel_q8_encoded;
/* GPU-decided shared-expert lane split (see ds4_shared_split_range in
* metal/dense.metal): the split kernels read the selected expert ids and
* take a complementary lane range per rank sized to balance the bytes each
* rank streams; the host binds full weight ranges and the ids buffer. */
typedef struct {
int32_t tp_rank;
int32_t tp_world;
int32_t n_expert;
int32_t n_expert_used;
int32_t shared_dim;
int32_t lane_granule;
int32_t shift_q16;
} ds4_gpu_shared_split_args;
static BOOL g_parallel_split_active;
static ds4_gpu_shared_split_args g_parallel_split_args;
static id<MTLBuffer> g_parallel_split_ids;
static NSUInteger g_parallel_split_ids_offset;
/* Reset is deliberately idempotent. Closing the concurrent encoder preserves
* work already encoded, while clearing every admission/reference field keeps
* a failed FFN path from turning later ordinary dispatches concurrent. */
static void ds4_gpu_parallel_ffn_reset_state(BOOL close_encoder) {
if (close_encoder &&
(g_batch_encoder_concurrent || g_parallel_q8_pending ||
g_parallel_q8_encoded || g_parallel_ffn_mode != 0 ||
g_parallel_ffn_stage != 0)) {
ds4_gpu_close_batch_encoder();
}
g_batch_encoder_concurrent = NO;
g_parallel_q8_pending = NO;
g_parallel_q8_encoded = NO;
g_parallel_ffn_mode = 0;
g_parallel_ffn_stage = 0;
g_parallel_q8_pipeline = nil;
g_parallel_q8_weight = nil;
g_parallel_q8_x = nil;
g_parallel_q8_out = nil;
g_parallel_q8_weight_offset = 0;
g_parallel_q8_x_offset = 0;
g_parallel_q8_out_offset = 0;
g_parallel_q8_args = (ds4_gpu_q8_0_matvec_args){0};
g_parallel_gate_up_pipeline = nil;
g_parallel_gate_weight = nil;
g_parallel_up_weight = nil;
g_parallel_gate_x = nil;
g_parallel_gate_out = nil;
g_parallel_up_out = nil;
g_parallel_mid_out = nil;
g_parallel_gate_weight_offset = 0;
g_parallel_up_weight_offset = 0;
g_parallel_gate_x_offset = 0;
g_parallel_gate_out_offset = 0;
g_parallel_up_out_offset = 0;
g_parallel_mid_out_offset = 0;
g_parallel_gate_up_args = (ds4_gpu_q8_0_matvec_args){0};
g_parallel_gate_up_clamp = 0.0f;
g_parallel_gate_up_nsg = 0;
g_parallel_gate_up_nr0 = 0;
g_parallel_split_active = NO;
g_parallel_split_args = (ds4_gpu_shared_split_args){0};
g_parallel_split_ids = nil;
g_parallel_split_ids_offset = 0;
g_parallel_gate_up_smem = 0;
}
void ds4_gpu_parallel_ffn_abort(void) {
ds4_gpu_parallel_ffn_reset_state(YES);
}
/* Clang cleanup makes every early return from the large generic routed-MoE
* wrapper abort an armed concurrent FFN without touching its many
* established fallback branches. */
static void ds4_gpu_parallel_ffn_scope_cleanup(BOOL *armed) {
if (armed && *armed && g_parallel_q8_pending) {
ds4_gpu_parallel_ffn_abort();
}
}
/* DeepSeek V4.1 keeps the shared expert's bf16 boundaries inside the
* concurrent encoder: mid rounded by the fused gate/up, bf16 in and out on
* the down projection. */
static BOOL g_parallel_ffn_bf16;
static int ds4_gpu_parallel_ffn_start_range(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t model_dim,
uint32_t shared_dim,
uint32_t shared_lane_offset,
uint32_t shared_lane_count,
const ds4_gpu_tensor *x,
float clamp) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!g_batch_cb || g_parallel_q8_pending || g_batch_encoder_concurrent ||
!gate || !up || !mid || !shared_out || !x || !model_map ||
model_dim == 0 || shared_dim == 0 || shared_lane_count == 0 ||
shared_lane_offset > shared_dim ||
shared_lane_count > shared_dim - shared_lane_offset ||
(model_dim & 31u) != 0 || (shared_dim & 31u) != 0 ||
(shared_lane_offset & 31u) != 0 ||
(shared_lane_count & 31u) != 0 ||
!isfinite(clamp) || clamp < 0.0f) {
return 0;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(shared_out);
if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < (uint64_t)model_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(gate) < (uint64_t)shared_lane_count * sizeof(float) ||
ds4_gpu_tensor_bytes(up) < (uint64_t)shared_lane_count * sizeof(float) ||
ds4_gpu_tensor_bytes(mid) < (uint64_t)shared_lane_count * sizeof(float) ||
ds4_gpu_tensor_bytes(shared_out) < (uint64_t)model_dim * sizeof(float)) {
return 0;
}
const uint64_t gate_row_bytes = ((uint64_t)model_dim / 32u) * 34u;
const uint64_t gate_weight_bytes =
(uint64_t)shared_lane_count * gate_row_bytes;
const uint64_t gate_lane_bytes =
(uint64_t)shared_lane_offset * gate_row_bytes;
const uint64_t down_row_bytes = ((uint64_t)shared_dim / 32u) * 34u;
const uint64_t down_weight_bytes = (uint64_t)model_dim * down_row_bytes;
if (gate_offset > model_size || gate_lane_bytes > model_size - gate_offset ||
gate_weight_bytes > model_size - gate_offset - gate_lane_bytes ||
up_offset > model_size || gate_lane_bytes > model_size - up_offset ||
gate_weight_bytes > model_size - up_offset - gate_lane_bytes ||
down_offset > model_size || down_weight_bytes > model_size - down_offset) {
return 0;
}
uint64_t gate_inner = 0, up_inner = 0, down_inner = 0;
id<MTLBuffer> gate_wbuf = ds4_gpu_wrap_model_range(
model_map, model_size, gate_offset + gate_lane_bytes,
gate_weight_bytes, &gate_inner);
id<MTLBuffer> up_wbuf = ds4_gpu_wrap_model_range(
model_map, model_size, up_offset + gate_lane_bytes,
gate_weight_bytes, &up_inner);
id<MTLBuffer> down_wbuf = ds4_gpu_wrap_model_range(
model_map, model_size, down_offset, down_weight_bytes, &down_inner);
if (!gate_wbuf || !up_wbuf || !down_wbuf) return 0;
ds4_gpu_mv_dispatch gate_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
const char *gate_fn = g_parallel_ffn_bf16 ?
"kernel_dsv4_shared_mid_swiglu_q8_0_bf16" : "kernel_dsv4_shared_gate_up_swiglu_q8_0";
id<MTLComputePipelineState> gate_pipeline =
ds4_gpu_get_mul_mv_pipeline(gate_fn, gate_dispatch.nsg);
id<MTLComputePipelineState> down_pipeline =
ds4_gpu_get_mul_mv_pipeline(g_parallel_ffn_bf16 ?
"kernel_mul_mv_q8_0_f32_bf16io" : "kernel_mul_mv_q8_0_f32", 4);
if (!gate_pipeline || !down_pipeline ||
down_pipeline.maxTotalThreadsPerThreadgroup < 128u) {
return 0;
}
ds4_gpu_close_batch_encoder();
g_batch_encoder_concurrent = YES;
g_parallel_gate_up_pipeline = gate_pipeline;
g_parallel_gate_weight = gate_wbuf;
g_parallel_up_weight = up_wbuf;
g_parallel_gate_x = xbuf;
g_parallel_gate_out = gatebuf;
g_parallel_up_out = upbuf;
g_parallel_mid_out = midbuf;
g_parallel_gate_weight_offset = (NSUInteger)gate_inner;
g_parallel_up_weight_offset = (NSUInteger)up_inner;
g_parallel_gate_x_offset = ds4_gpu_tensor_offset(x);
g_parallel_gate_out_offset = ds4_gpu_tensor_offset(gate);
g_parallel_up_out_offset = ds4_gpu_tensor_offset(up);
g_parallel_mid_out_offset = ds4_gpu_tensor_offset(mid);
g_parallel_gate_up_args =
ds4_gpu_make_q8_0_mv_args(model_dim, shared_lane_count);
g_parallel_gate_up_args.nr0 = gate_dispatch.nr0;
g_parallel_gate_up_clamp = clamp;
g_parallel_gate_up_nsg = gate_dispatch.nsg;
g_parallel_gate_up_nr0 = gate_dispatch.nr0;
g_parallel_gate_up_smem = gate_dispatch.smem;
g_parallel_q8_pipeline = down_pipeline;
g_parallel_q8_weight = down_wbuf;
g_parallel_q8_x = midbuf;
g_parallel_q8_out = outbuf;
g_parallel_q8_weight_offset =
(NSUInteger)(down_inner +
((uint64_t)shared_lane_offset / 32u) * 34u);
g_parallel_q8_x_offset = ds4_gpu_tensor_offset(mid);
g_parallel_q8_out_offset = ds4_gpu_tensor_offset(shared_out);
g_parallel_q8_args =
ds4_gpu_make_q8_0_mv_args(shared_dim, model_dim);
g_parallel_q8_args.ne00 = (int32_t)shared_lane_count;
g_parallel_q8_args.ne10 = (int32_t)shared_lane_count;
g_parallel_q8_args.nr0 = 2;
g_parallel_ffn_mode = 2;
g_parallel_ffn_stage = 0;
g_parallel_q8_pending = YES;
g_parallel_q8_encoded = NO;
return 1;
}
int ds4_gpu_parallel_ffn_start(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t model_dim,
uint32_t shared_dim,
const ds4_gpu_tensor *x,
float clamp) {
return ds4_gpu_parallel_ffn_start_range(
gate, up, mid, shared_out, model_map, model_size,
gate_offset, up_offset, down_offset, model_dim, shared_dim,
0, shared_dim, x, clamp);
}
int ds4_gpu_parallel_ffn_start_bf16(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t model_dim,
uint32_t shared_dim,
const ds4_gpu_tensor *x,
float clamp) {
g_parallel_ffn_bf16 = YES;
const int ok = ds4_gpu_parallel_ffn_start_range(
gate, up, mid, shared_out, model_map, model_size,
gate_offset, up_offset, down_offset, model_dim, shared_dim,
0, shared_dim, x, clamp);
g_parallel_ffn_bf16 = NO;
return ok;
}
int ds4_gpu_parallel_ffn_start_sliced(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t model_dim,
uint32_t shared_dim,
uint32_t shared_lane_offset,
uint32_t shared_lane_count,
const ds4_gpu_tensor *x,
float clamp) {
return ds4_gpu_parallel_ffn_start_range(
gate, up, mid, shared_out, model_map, model_size,
gate_offset, up_offset, down_offset, model_dim, shared_dim,
shared_lane_offset, shared_lane_count, x, clamp);
}
int ds4_gpu_parallel_ffn_start_split(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t model_dim,
uint32_t shared_dim,
const ds4_gpu_tensor *x,
float clamp,
const ds4_gpu_tensor *selected,
uint32_t tp_rank,
uint32_t tp_world,
uint32_t n_expert,
uint32_t n_expert_used,
uint32_t shift_q16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!g_batch_cb || g_parallel_q8_pending || g_batch_encoder_concurrent ||
!gate || !up || !mid || !shared_out || !x || !selected || !model_map ||
model_dim == 0 || shared_dim == 0 ||
tp_world != 2 || tp_rank >= tp_world ||
n_expert == 0 || (n_expert % tp_world) != 0 || n_expert_used == 0 ||
n_expert_used > 64 || shift_q16 > 65536u ||
(model_dim & 31u) != 0 || (shared_dim & 63u) != 0 ||
!isfinite(clamp) || clamp < 0.0f) {
return 0;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(shared_out);
id<MTLBuffer> idsbuf = ds4_gpu_tensor_buffer(selected);
if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !idsbuf ||
ds4_gpu_tensor_bytes(x) < (uint64_t)model_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(gate) < (uint64_t)shared_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(up) < (uint64_t)shared_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(mid) < (uint64_t)shared_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(shared_out) < (uint64_t)model_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int32_t)) {
return 0;
}
const uint64_t gate_row_bytes = ((uint64_t)model_dim / 32u) * 34u;
const uint64_t gate_weight_bytes = (uint64_t)shared_dim * gate_row_bytes;
const uint64_t down_row_bytes = ((uint64_t)shared_dim / 32u) * 34u;
const uint64_t down_weight_bytes = (uint64_t)model_dim * down_row_bytes;
if (gate_offset > model_size || gate_weight_bytes > model_size - gate_offset ||
up_offset > model_size || gate_weight_bytes > model_size - up_offset ||
down_offset > model_size || down_weight_bytes > model_size - down_offset) {
return 0;
}
uint64_t gate_inner = 0, up_inner = 0, down_inner = 0;
id<MTLBuffer> gate_wbuf = ds4_gpu_wrap_model_range(
model_map, model_size, gate_offset, gate_weight_bytes, &gate_inner);
id<MTLBuffer> up_wbuf = ds4_gpu_wrap_model_range(
model_map, model_size, up_offset, gate_weight_bytes, &up_inner);
id<MTLBuffer> down_wbuf = ds4_gpu_wrap_model_range(
model_map, model_size, down_offset, down_weight_bytes, &down_inner);
if (!gate_wbuf || !up_wbuf || !down_wbuf) return 0;
ds4_gpu_mv_dispatch gate_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
id<MTLComputePipelineState> gate_pipeline =
ds4_gpu_get_mul_mv_pipeline(
"kernel_dsv4_shared_gate_up_swiglu_q8_0_split", gate_dispatch.nsg);
id<MTLComputePipelineState> down_pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_shared_down_q8_0_split", 4);
if (!gate_pipeline || !down_pipeline ||
down_pipeline.maxTotalThreadsPerThreadgroup < 128u) {
return 0;
}
ds4_gpu_close_batch_encoder();
g_batch_encoder_concurrent = YES;
g_parallel_gate_up_pipeline = gate_pipeline;
g_parallel_gate_weight = gate_wbuf;
g_parallel_up_weight = up_wbuf;
g_parallel_gate_x = xbuf;
g_parallel_gate_out = gatebuf;
g_parallel_up_out = upbuf;
g_parallel_mid_out = midbuf;
g_parallel_gate_weight_offset = (NSUInteger)gate_inner;
g_parallel_up_weight_offset = (NSUInteger)up_inner;
g_parallel_gate_x_offset = ds4_gpu_tensor_offset(x);
g_parallel_gate_out_offset = ds4_gpu_tensor_offset(gate);
g_parallel_up_out_offset = ds4_gpu_tensor_offset(up);
g_parallel_mid_out_offset = ds4_gpu_tensor_offset(mid);
/* Full width: the grid covers every shared row and the kernels trim to
* the rank's lane range. */
g_parallel_gate_up_args = ds4_gpu_make_q8_0_mv_args(model_dim, shared_dim);
g_parallel_gate_up_args.nr0 = gate_dispatch.nr0;
g_parallel_gate_up_clamp = clamp;
g_parallel_gate_up_nsg = gate_dispatch.nsg;
g_parallel_gate_up_nr0 = gate_dispatch.nr0;
g_parallel_gate_up_smem = gate_dispatch.smem;
g_parallel_q8_pipeline = down_pipeline;
g_parallel_q8_weight = down_wbuf;
g_parallel_q8_x = midbuf;
g_parallel_q8_out = outbuf;
g_parallel_q8_weight_offset = (NSUInteger)down_inner;
g_parallel_q8_x_offset = ds4_gpu_tensor_offset(mid);
g_parallel_q8_out_offset = ds4_gpu_tensor_offset(shared_out);
g_parallel_q8_args = ds4_gpu_make_q8_0_mv_args(shared_dim, model_dim);
g_parallel_q8_args.nr0 = 2;
g_parallel_split_active = YES;
g_parallel_split_args = (ds4_gpu_shared_split_args) {
.tp_rank = (int32_t)tp_rank,
.tp_world = (int32_t)tp_world,
.n_expert = (int32_t)n_expert,
.n_expert_used = (int32_t)n_expert_used,
.shared_dim = (int32_t)shared_dim,
.lane_granule = 64,
.shift_q16 = (int32_t)shift_q16,
};
g_parallel_split_ids = idsbuf;
g_parallel_split_ids_offset = ds4_gpu_tensor_offset(selected);
g_parallel_ffn_mode = 2;
g_parallel_ffn_stage = 0;
g_parallel_q8_pending = YES;
g_parallel_q8_encoded = NO;
return 1;
}
static void ds4_gpu_encode_parallel_q8_down(
id<MTLComputeCommandEncoder> enc) {
[enc setComputePipelineState:g_parallel_q8_pipeline];
[enc setBytes:&g_parallel_q8_args
length:sizeof(g_parallel_q8_args)
atIndex:0];
if (g_parallel_split_active) {
[enc setBytes:&g_parallel_split_args
length:sizeof(g_parallel_split_args)
atIndex:1];
[enc setBuffer:g_parallel_q8_weight
offset:g_parallel_q8_weight_offset
atIndex:2];
[enc setBuffer:g_parallel_q8_x offset:g_parallel_q8_x_offset atIndex:3];
[enc setBuffer:g_parallel_q8_out offset:g_parallel_q8_out_offset atIndex:4];
[enc setBuffer:g_parallel_split_ids
offset:g_parallel_split_ids_offset
atIndex:5];
} else {
[enc setBuffer:g_parallel_q8_weight
offset:g_parallel_q8_weight_offset
atIndex:1];
[enc setBuffer:g_parallel_q8_x offset:g_parallel_q8_x_offset atIndex:2];
[enc setBuffer:g_parallel_q8_out offset:g_parallel_q8_out_offset atIndex:3];
}
[enc setThreadgroupMemoryLength:32u * 2u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(
((NSUInteger)g_parallel_q8_args.ne0 + 1u) / 2u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
}
static int ds4_gpu_parallel_q8_matvec_encode_pending(
id<MTLCommandBuffer> cb,
id<MTLBuffer> routed_mid) {
if (!g_parallel_q8_pending || g_parallel_ffn_mode != 2 ||
g_parallel_ffn_stage != 0 || !g_batch_encoder_concurrent ||
!g_batch_cb || cb != g_batch_cb) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
if (!enc || enc.dispatchType != MTLDispatchTypeConcurrent) return 0;
[enc setComputePipelineState:g_parallel_gate_up_pipeline];
if (g_parallel_split_active) {
[enc setBytes:&g_parallel_gate_up_args
length:sizeof(g_parallel_gate_up_args)
atIndex:0];
[enc setBytes:&g_parallel_split_args
length:sizeof(g_parallel_split_args)
atIndex:1];
[enc setBuffer:g_parallel_gate_weight
offset:g_parallel_gate_weight_offset
atIndex:2];
[enc setBuffer:g_parallel_up_weight
offset:g_parallel_up_weight_offset
atIndex:3];
[enc setBuffer:g_parallel_gate_x
offset:g_parallel_gate_x_offset
atIndex:4];
[enc setBuffer:g_parallel_gate_out
offset:g_parallel_gate_out_offset
atIndex:5];
[enc setBuffer:g_parallel_up_out
offset:g_parallel_up_out_offset
atIndex:6];
[enc setBuffer:g_parallel_mid_out
offset:g_parallel_mid_out_offset
atIndex:7];
[enc setBytes:&g_parallel_gate_up_clamp
length:sizeof(g_parallel_gate_up_clamp)
atIndex:8];
[enc setBuffer:g_parallel_split_ids
offset:g_parallel_split_ids_offset
atIndex:9];
} else {
[enc setBytes:&g_parallel_gate_up_args
length:sizeof(g_parallel_gate_up_args)
atIndex:0];
[enc setBuffer:g_parallel_gate_weight
offset:g_parallel_gate_weight_offset
atIndex:1];
[enc setBuffer:g_parallel_up_weight
offset:g_parallel_up_weight_offset
atIndex:2];
[enc setBuffer:g_parallel_gate_x
offset:g_parallel_gate_x_offset
atIndex:3];
[enc setBuffer:g_parallel_gate_out
offset:g_parallel_gate_out_offset
atIndex:4];
[enc setBuffer:g_parallel_up_out
offset:g_parallel_up_out_offset
atIndex:5];
[enc setBuffer:g_parallel_mid_out
offset:g_parallel_mid_out_offset
atIndex:6];
[enc setBytes:&g_parallel_gate_up_clamp
length:sizeof(g_parallel_gate_up_clamp)
atIndex:7];
}
[enc setThreadgroupMemoryLength:2u * g_parallel_gate_up_smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(
((NSUInteger)g_parallel_gate_up_args.ne0 +
g_parallel_gate_up_nr0 - 1u) / g_parallel_gate_up_nr0,
1, 1)
threadsPerThreadgroup:MTLSizeMake(32, g_parallel_gate_up_nsg, 1)];
g_parallel_ffn_stage = 1;
/* Metal defines either barrier form as an execution barrier for every
* earlier dispatch in this concurrent encoder. The resource list narrows
* visibility, not completion. This is therefore the explicit level break:
* routed pair-SwiGLU and shared gate/up must finish before
* the routed Q2 sum6 consumer can start. */
id<MTLResource> barrier_resources[2] = {
routed_mid,
g_parallel_mid_out,
};
[enc memoryBarrierWithResources:barrier_resources count:2u];
return 1;
}
static int ds4_gpu_parallel_ffn_encode_second_stage(
id<MTLCommandBuffer> cb) {
if (!g_parallel_q8_pending || g_parallel_ffn_mode != 2 ||
g_parallel_ffn_stage != 1 || !g_batch_encoder_concurrent ||
!g_batch_cb || cb != g_batch_cb) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
if (!enc || enc.dispatchType != MTLDispatchTypeConcurrent) return 0;
ds4_gpu_encode_parallel_q8_down(enc);
/* No final memory barrier is required: finish() ends this concurrent
* encoder before the serial HC encoder is created, and a Metal command
* buffer is an ordered list of command encoders. */
g_parallel_ffn_stage = 2;
g_parallel_q8_encoded = YES;
return 1;
}
int ds4_gpu_parallel_ffn_finish(void) {
const int completed =
g_parallel_q8_pending && g_parallel_q8_encoded &&
g_parallel_ffn_stage == 2 && g_batch_encoder_concurrent;
/* Reset even on an incomplete join. This makes the error path just as
* safe and idempotent as an explicit abort. */
ds4_gpu_parallel_ffn_reset_state(YES);
return completed;
}
static int ds4_gpu_stream_expert_cache_wait_inflight(const char *label) {
const char *what = label ? label : "streaming expert cache in-flight";
if (g_batch_cb && ds4_gpu_flush_commands() == 0) return 0;
if ([g_pending_cbs count] != 0 &&
ds4_gpu_wait_pending_command_buffers(what) == 0) {
return 0;
}
return 1;
}
int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) {
if (!event_value) return 0;
*event_value = 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
ds4_gpu_parallel_ffn_reset_state(YES);
if (!g_batch_cb) return 0;
if (@available(macOS 12.0, *)) {
if (!g_selected_readback_event) {
g_selected_readback_event = [g_device newSharedEvent];
if (!g_selected_readback_event) {
fprintf(stderr, "ds4: failed to create Metal shared event for selected-id overlap\n");
return 0;
}
}
ds4_gpu_close_batch_encoder();
const uint64_t value = ++g_selected_readback_event_value;
[g_batch_cb encodeSignalEvent:g_selected_readback_event value:value];
g_batch_has_work = YES;
*event_value = value;
return 1;
}
fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n");
return 0;
}
int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label) {
if (!g_initialized && !ds4_gpu_init()) return 0;
ds4_gpu_parallel_ffn_reset_state(YES);
if (!g_batch_cb || event_value == 0) return 0;
if (@available(macOS 12.0, *)) {
if (!g_selected_readback_event) return 0;
ds4_gpu_close_batch_encoder();
id<MTLCommandBuffer> cb = g_batch_cb;
g_batch_cb = nil;
g_batch_has_work = NO;
[cb commit];
ds4_gpu_stream_expert_cache_note_batch_committed();
const char *what = label ? label : "selected-id overlap";
const BOOL signaled =
[g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000];
[g_pending_cbs addObject:cb];
if (!signaled) {
fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what);
(void)ds4_gpu_wait_pending_command_buffers(what);
[g_transient_buffers removeAllObjects];
return 0;
}
if (cb.status == MTLCommandBufferStatusError) {
fprintf(stderr, "ds4: Metal %s failed: %s\n",
what,
[[cb.error localizedDescription] UTF8String]);
(void)ds4_gpu_wait_pending_command_buffers(what);
[g_transient_buffers removeAllObjects];
return 0;
}
g_batch_cb = ds4_gpu_new_command_buffer();
g_batch_has_work = NO;
if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
if (!g_batch_cb) {
(void)ds4_gpu_wait_pending_command_buffers(what);
[g_transient_buffers removeAllObjects];
return 0;
}
return 1;
}
fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n");
return 0;
}
/*
* Tensor-parallel gates.
*
* A TP gate is a mid-command-stream rendezvous with the peer machine: the
* kernels ahead of the gate leave a partial block output in a slab slot,
* the GPU signals g_tp_gpu_event, and the pre-encoded combine kernel waits
* on g_tp_cpu_event. A dedicated service thread bridges the two: it spins
* until the GPU reaches the gate, runs the transport exchange (RDMA WRITE
* plus flag poll, or a TCP write/read pair — behind the callback), and
* CPU-signals the release. On exchange failure the release is signaled
* anyway so the GPU never deadlocks; the failure latches in g_tp_failed
* and the eval aborts at the next command-buffer boundary.
*
* Gate sequence values increase monotonically per encoded gate. Both ranks
* encode the identical graph, so the values agree by construction and slots
* never need resetting between tokens.
*/
typedef struct {
uint32_t layer;
uint32_t gate;
uint32_t rows; /* 0 = row gate; >0 = verify-block batch gate */
uint32_t batch_slot; /* slab batch slot of a verify-block gate */
uint32_t event_arrival;
uint64_t seq;
const volatile uint32_t *check; /* batch gate payload behind a checked flag */
uint32_t check_words;
/* Big batch gates (prefill): exchange big_bytes from big_out into
* big_in directly (CPU-visible bounce buffers), bypassing the slab. */
const void *big_out;
void *big_in;
uint64_t big_bytes;
uint32_t poll; /* release through the poll region */
uint32_t spin; /* inline gate: release through the coherent spin word */
} ds4_gpu_tp_request;
enum { DS4_GPU_TP_QUEUE = 1024 };
static id<MTLSharedEvent> g_tp_gpu_event; /* GPU -> service thread */
static id<MTLSharedEvent> g_tp_cpu_event; /* service thread -> GPU */
/* Batch (verify-block) gates run on their own sequence space and release
* event: the row-gate seq feeds the RDMA pre-posted recv accounting, which
* requires consecutive values, and a shared release event would make a
* small batch value satisfy waits armed against the larger row seq. */
static id<MTLSharedEvent> g_tp_batch_gpu_event;
static id<MTLSharedEvent> g_tp_batch_cpu_event;
static uint64_t g_tp_batch_seq;
/* Batch flag values are tagged so a stale row-gate seq in the reused FFN
* flag word can never satisfy a batch arrival spin (and vice versa). */
#define DS4_TP_BATCH_FLAG_TAG 0x80000000u
/* Expert-ownership split parameters for routed kernels. World 1 means TP is
* not bound; world 2 assigns each rank one contiguous expert range. */
static int32_t g_tp_split_rank;
static int32_t g_tp_split_world = 1;
static int32_t g_tp_session_batch_mode;
static int ds4_gpu_tp_world_is_two(void) {
return g_tp_split_world == 2;
}
/* Decode balances the routed experts by list position instead of by id
* range: every expert is bound and the kernels mask nothing. */
static int g_tp_full_expert_bind;
void ds4_gpu_tp_full_expert_bind(int on) { g_tp_full_expert_bind = on; }
/* Return the contiguous routed-expert range backed by this process. Rank 1
* owns the high range and receives any odd-count remainder. */
static void ds4_gpu_tp_expert_range(uint32_t n_total_expert,
uint32_t *first_expert,
uint32_t *n_expert) {
*first_expert = 0;
*n_expert = n_total_expert;
if (g_tp_split_world != 2 || g_tp_full_expert_bind) return;
const uint32_t low_experts = n_total_expert / 2u;
if (g_tp_split_rank == 1) {
*first_expert = low_experts;
*n_expert = n_total_expert - low_experts;
} else {
*n_expert = low_experts;
}
}
/* Attention head split for GLM batch prefill: each rank computes a
* contiguous half of the heads in the qk-low / attention-lora /
* value-project batch kernels; the caller zeroes the unowned head range
* of the heads buffer and combines the attn-output partials over the
* TP big-gate exchange. */
static int32_t g_tp_attn_head_split;
void ds4_gpu_tp_set_attn_head_split(int enabled) {
g_tp_attn_head_split = enabled ? 1 : 0;
}
static void ds4_gpu_tp_attn_head_range(uint32_t n_head,
uint32_t group,
uint32_t *head_base,
uint32_t *head_count) {
*head_base = 0;
*head_count = n_head;
if (!g_tp_attn_head_split || g_tp_split_world != 2) return;
const uint32_t half = n_head / 2u;
if (half == 0u || (half % group) != 0u || (n_head % 2u) != 0u) return;
*head_count = half;
*head_base = g_tp_split_rank == 1 ? half : 0u;
}
/* Flag gates (DS4_TP_FLAG_GATES): the GPU publishes gate arrival by storing
* the sequence number into a slab word instead of signaling the shared
* event; the service thread spin-reads it from shared memory, which wakes
* hundreds of microseconds earlier than signaledValue polling. The
* CPU->GPU release direction stays on the shared event. */
static bool g_tp_flag_gates;
/* Poll gates (default with flag gates): the GPU waits for a gate release by
* probing fresh cache lines at the head of a new command buffer instead of
* parking on g_tp_cpu_event. See kernel_dsv4_tp_poll_release. */
#define DS4_TP_POLL_LINES 8192u
#define DS4_TP_POLL_LINE_BYTES 128u
#define DS4_TP_POLL_MAX_SLOTS 1024u
#define DS4_TP_POLL_RING 16u /* regions cycled per gate so no line is re-read while cached */
static bool g_tp_poll_gates;
static id<MTLBuffer> g_tp_poll_buffer; /* release region + status words */
static volatile uint32_t *g_tp_poll_region; /* CPU view of the release lines */
static volatile uint32_t *g_tp_poll_status; /* per-slot poll result, GPU-written */
/* Inline decode gates: the partial and the flag are published with
* system-coherent stores and kernel_dsv4_tp_release_wait spins inside the open
* command buffer on [0] of this buffer ([1] spin count of the last wait, [2]
* timeouts), so a token ends no command buffer per gate.
* DS4_TP_DISABLE_DECODE_POLL_GATES keeps the flushing poll gates;
* DS4_TP_INLINE_SPIN_ITERS bounds the spin (default 8M). */
static id<MTLBuffer> g_tp_release_buffer;
static volatile uint32_t *g_tp_release_ptr;
static uint32_t g_tp_spin_max_iters;
static uint64_t g_tp_stat_spin_iters;
static int g_tp_poll_prev_valid;
static uint32_t g_tp_poll_prev_slot;
static uint64_t g_tp_poll_prev_seq;
/* Gate-time prefetch: the poll wait leaves the memory bus idle, so the next
* phase's weights are streamed into the cache first (see the notebook). */
#define DS4_TP_PREFETCH_MAX 8u
#define DS4_TP_PREFETCH_BYTES_PER_US 450000ull /* ~450 GB/s streaming */
#define DS4_TP_PREFETCH_MARGIN_US 8.0 /* flush + wire + pickup */
static double g_tp_exchange_ewma_us[2]; /* per gate kind, always on */
typedef struct { uint64_t offset; uint64_t bytes; } ds4_tp_prefetch_range;
static ds4_tp_prefetch_range g_tp_prefetch_plan[2][DS4_TP_PREFETCH_MAX]; /* per gate kind */
/* Gate flag folded into the payload producer (see
* kernel_dsv4_add2_f32_tp_flag_checked): the producer publishes the checked
* flag for the sequence number the next gate encode will use, and the gate
* encode skips its own flag kernel. */
static uint64_t g_tp_flag_prepublished_seq;
static id<MTLBuffer> g_tp_fold_ctl;
static uint32_t g_tp_prefetch_count[2];
static const void *g_tp_prefetch_map;
static uint64_t g_tp_prefetch_map_size;
static id<MTLBuffer> g_tp_prefetch_scratch;
static bool g_tp_gate_prefetch;
static uint64_t g_tp_stat_poll_lines; /* sum of poll hit lines (profile) */
static uint64_t g_tp_stat_poll_max_line;
static uint64_t g_tp_stat_poll_count;
static uint64_t g_tp_stat_encode_lead; /* sum of (encoded seq - gate seq) at arrival */
static double g_tp_stat_release_ms; /* time spent writing poll releases */
static double g_tp_stat_cbwait_ms; /* flag seen -> payload verified in memory */
static uint64_t g_tp_out_off; /* slab offset of out slot 0 */
static uint64_t g_tp_vec_bytes; /* partial vector bytes per slot */
static id<MTLBuffer> g_tp_check_buffer; /* per-slot payload checksums */
static volatile uint32_t *g_tp_check_words;
static volatile uint8_t *g_tp_slab_cpu; /* CPU view of the slab base */
/* Bound on gate command buffers in flight: the GPU only needs a few
* segments of runway and a deep backlog costs command-buffer memory. */
#define DS4_TP_POLL_MAX_INFLIGHT 32u
static id<MTLBuffer> g_tp_slab_buffer;
static NSUInteger g_tp_slab_buffer_off;
static volatile uint32_t *g_tp_gpu_flags; /* CPU view of the flag words */
static uint64_t g_tp_gpu_flags_off;
static uint64_t g_tp_seq;
static ds4_gpu_tp_exchange_fn g_tp_exchange_fn;
static ds4_gpu_tp_batch_exchange_fn g_tp_batch_exchange_fn;
static ds4_gpu_tp_big_exchange_fn g_tp_big_exchange_fn;
void ds4_gpu_tp_set_big_exchange(ds4_gpu_tp_big_exchange_fn fn) {
g_tp_big_exchange_fn = fn;
}
static void *g_tp_exchange_ud;
static pthread_t g_tp_thread;
static int g_tp_thread_running;
static int g_tp_shutdown;
static int g_tp_failed_flag;
static pthread_mutex_t g_tp_mutex = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t g_tp_cond = PTHREAD_COND_INITIALIZER;
static ds4_gpu_tp_request g_tp_queue[DS4_GPU_TP_QUEUE];
static uint32_t g_tp_queue_head;
static uint32_t g_tp_queue_count;
static uint64_t g_tp_stat_gates;
static uint64_t g_tp_stat_batch_gates; /* verify-block batch gates */
static double g_tp_stat_batch_gpu_wait_ms; /* arrival spin (GPU compute + park) */
static double g_tp_stat_batch_exchange_ms; /* exchange callback (peer wait + wire) */
static double g_tp_stat_batch_release_ms; /* exchange end -> cpu event signaled */
static double g_tp_stat_gpu_wait_ms;
static double g_tp_stat_exchange_ms;
enum { DS4_GPU_TP_GATE_ATTN = 0, DS4_GPU_TP_GATE_FFN = 1,
DS4_GPU_TP_GATES_PER_LAYER = 2 };
static uint64_t g_tp_stat_gate_count[DS4_GPU_TP_GATES_PER_LAYER];
static double g_tp_stat_gate_gpu_wait_ms[DS4_GPU_TP_GATES_PER_LAYER];
static double g_tp_stat_gate_exchange_ms[DS4_GPU_TP_GATES_PER_LAYER];
/* GPU keep-alive (see kernel_dsv4_tp_keepalive): its own queue and thread,
* alive exactly as long as the TP gate machinery. */
static id<MTLCommandQueue> g_tp_keepalive_queue;
static id<MTLBuffer> g_tp_keepalive_buffer;
static pthread_t g_tp_keepalive_thread;
static int g_tp_keepalive_running;
/* Nonzero while a verify block runs: the GPU is genuinely busy
* there, so the keep-alive is a pure parasite (~2.3ms per 5-row block
* measured against the single-machine verify). */
static volatile int g_tp_keepalive_paused;
void ds4_gpu_tp_keepalive_pause(int paused) {
g_tp_keepalive_paused = paused;
}
void ds4_gpu_tp_set_session_batch_mode(int enabled) {
g_tp_session_batch_mode = enabled ? 1 : 0;
}
int ds4_gpu_tp_decode_split_flush_safe(void) {
return g_tp_thread_running && g_tp_flag_gates &&
!g_tp_session_batch_mode;
}
int ds4_gpu_tp_gate_prefetch_plan(uint32_t gate,
const void *model_map, uint64_t model_size,
const uint64_t *offsets, const uint64_t *bytes,
uint32_t count) {
if (gate >= 2u) return 0;
if (!g_tp_gate_prefetch || !g_tp_poll_gates || !model_map) {
g_tp_prefetch_count[gate] = 0;
return 1;
}
if (count > DS4_TP_PREFETCH_MAX) count = DS4_TP_PREFETCH_MAX;
uint32_t n = 0;
for (uint32_t i = 0; i < count; i++) {
if (bytes[i] == 0 || offsets[i] > model_size ||
bytes[i] > model_size - offsets[i]) continue;
g_tp_prefetch_plan[gate][n].offset = offsets[i];
g_tp_prefetch_plan[gate][n].bytes = bytes[i];
n++;
}
g_tp_prefetch_map = model_map;
g_tp_prefetch_map_size = model_size;
g_tp_prefetch_count[gate] = n;
return 1;
}
/* Encode the touch dispatches of one gate's plan into the open batch. */
static uint64_t ds4_gpu_buffer_address(id<MTLBuffer> buffer, NSUInteger inner);
/* One dispatch of kernel_touch_u8_stride_table over up to DS4_TP_PREFETCH_MAX
* model ranges (priority order, cut at `budget` bytes), 128 B per thread. */
static void ds4_gpu_encode_touch_table(id<MTLComputeCommandEncoder> enc,
id<MTLComputePipelineState> touch,
const ds4_tp_prefetch_range *plan,
uint32_t n,
uint64_t budget) {
const uint64_t stride = 128;
uint64_t addr[DS4_TP_PREFETCH_MAX] = {0};
uint32_t lines[DS4_TP_PREFETCH_MAX] = {0};
id<MTLResource> views[DS4_TP_PREFETCH_MAX];
uint32_t count = 0;
uint64_t total_lines = 0;
for (uint32_t i = 0; i < n && budget != 0u; i++) {
uint64_t inner = 0;
uint64_t bytes = plan[i].bytes;
if (bytes > budget) bytes = budget;
budget -= bytes;
if (bytes == 0u) continue;
id<MTLBuffer> src = ds4_gpu_wrap_model_range(
g_tp_prefetch_map, g_tp_prefetch_map_size, plan[i].offset, bytes, &inner);
if (!src) continue;
const uint64_t range_lines = (bytes + stride - 1u) / stride;
if (total_lines + range_lines > (1u << 20)) break;
const uint64_t address = ds4_gpu_buffer_address(src, (NSUInteger)inner);
if (address == 0u) continue;
addr[count] = address;
lines[count] = (uint32_t)range_lines;
views[count] = src;
total_lines += range_lines;
count++;
}
if (count == 0u) return;
[enc setComputePipelineState:touch];
[enc useResources:views count:count usage:MTLResourceUsageRead];
[enc setBytes:addr length:sizeof(addr) atIndex:0];
[enc setBytes:lines length:sizeof(lines) atIndex:1];
[enc setBytes:&count length:sizeof(count) atIndex:2];
[enc setBuffer:g_tp_prefetch_scratch offset:0 atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)((total_lines + 255u) / 256u), 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
}
static int ds4_gpu_tp_encode_gate_prefetch(uint32_t gate) {
const uint32_t n = g_tp_prefetch_count[gate];
g_tp_prefetch_count[gate] = 0;
if (n == 0 || !g_batch_cb) return 1;
if (!g_tp_prefetch_scratch) {
g_tp_prefetch_scratch =
[g_device newBufferWithLength:(NSUInteger)(1u << 20)
options:MTLResourceStorageModePrivate];
if (!g_tp_prefetch_scratch) return 1;
}
id<MTLComputePipelineState> touch = ds4_gpu_get_pipeline("kernel_touch_u8_stride_table");
if (!touch) return 1;
/* Budget: what fits in this rank's typical wait at this gate kind. The
* first observations use a small default so a slow rank never delays
* its own release pickup. */
const double wait_us = g_tp_exchange_ewma_us[gate];
/* Diagnostics: DS4_TP_PREFETCH_MARGIN_US and DS4_TP_PREFETCH_BYTES_PER_US
* override the sizing constants for experiments. */
static double margin_us = -1.0;
static double bytes_per_us = -1.0;
if (margin_us < 0.0) {
margin_us = (double)ds4_gpu_env_u64("DS4_TP_PREFETCH_MARGIN_US",
(uint64_t)DS4_TP_PREFETCH_MARGIN_US, 0u, 1000u);
bytes_per_us = (double)ds4_gpu_env_u64("DS4_TP_PREFETCH_BYTES_PER_US",
DS4_TP_PREFETCH_BYTES_PER_US, 1000u, 10000000ull);
}
uint64_t budget = wait_us > margin_us
? (uint64_t)((wait_us - margin_us) * bytes_per_us)
: 0u;
if (wait_us == 0.0) budget = 4u << 20;
if (budget == 0u) return 1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(g_batch_cb);
ds4_gpu_encode_touch_table(enc, touch, g_tp_prefetch_plan[gate], n, budget);
ds4_gpu_end_compute_encoder(g_batch_cb, enc);
return 1;
}
static uint32_t ds4_gpu_tp_keepalive_tgs_from_env(void) {
uint32_t ka_tgs = 1;
const char *tgs_env = getenv("DS4_TP_KEEPALIVE_TGS");
if (tgs_env) {
int v = atoi(tgs_env);
if (v > 0 && v <= 2048) ka_tgs = (uint32_t)v;
}
return ka_tgs;
}
static void *ds4_gpu_tp_keepalive_thread(void *arg) {
(void)arg;
/* Swept on the M5 Max pair: too few iterations lets clocks sag. Current
* TP split-resident Flash runs show 1.2M is a small Q4/Q2 decode win over
* 800k, while two threadgroups waste work. */
uint32_t iters = 1200000;
const char *env = getenv("DS4_TP_KEEPALIVE_ITERS");
if (env) iters = (uint32_t)atoi(env);
/* One ALU-only threadgroup keeps the GPU from power-gating but does
* not push the frequency governor; solo-vs-engine kernel gaps
* (~1.7x) suggest decode runs well below max clocks. More TGs raise
* apparent utilization without eating memory bandwidth. */
uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env();
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_tp_keepalive");
if (!pipeline) {
fprintf(stderr, "ds4: TP keep-alive pipeline missing\n");
return NULL;
}
while (!g_tp_shutdown) {
if (g_tp_keepalive_paused) {
usleep(200);
continue;
}
@autoreleasepool {
id<MTLCommandBuffer> cb = [g_tp_keepalive_queue commandBuffer];
id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
[enc setComputePipelineState:pipeline];
[enc setBuffer:g_tp_keepalive_buffer offset:0 atIndex:0];
[enc setBytes:&iters length:sizeof(iters) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)ka_tgs, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
[enc endEncoding];
[cb commit];
[cb waitUntilCompleted];
}
}
return NULL;
}
static void *ds4_gpu_tp_service_thread(void *arg) {
(void)arg;
const bool profile = getenv("DS4_TP_GATE_PROFILE") != NULL;
while (1) {
pthread_mutex_lock(&g_tp_mutex);
while (g_tp_queue_count == 0 && !g_tp_shutdown)
pthread_cond_wait(&g_tp_cond, &g_tp_mutex);
if (g_tp_shutdown && g_tp_queue_count == 0) {
pthread_mutex_unlock(&g_tp_mutex);
break;
}
ds4_gpu_tp_request req = g_tp_queue[g_tp_queue_head];
g_tp_queue_head = (g_tp_queue_head + 1) % DS4_GPU_TP_QUEUE;
g_tp_queue_count--;
pthread_mutex_unlock(&g_tp_mutex);
/* Wait for the GPU to reach this gate. Tight spin: gate arrival is
* on the decode critical path and normally tens of microseconds
* out; yielding here measurably delays the release wake-up. */
const double t0 = profile ? ds4_gpu_now_ms() : 0.0;
uint32_t spins = 0;
if (req.big_bytes > 0) {
/* Big gates always signal arrival through the batch shared
* event (see ds4_gpu_tp_big_gate_kick): the event completion
* semantics are what guarantee the bounce payload is visible
* before the exchange reads it. */
/* A big-gate exchange takes milliseconds, so a tight spin on
* the shared event buys nothing; each signaledValue read is a
* driver round trip. Sleep-poll unless asked otherwise
* (DS4_TP_BIG_GATE_SPIN_US=0 restores the tight spin). */
static int64_t big_spin_us = -1;
if (big_spin_us < 0)
big_spin_us = (int64_t)ds4_gpu_env_u64("DS4_TP_BIG_GATE_SPIN_US", 20u, 0u, 5000u);
while (g_tp_batch_gpu_event.signaledValue < req.seq) {
if (g_tp_shutdown) break;
if (big_spin_us > 0) usleep((useconds_t)big_spin_us);
else if (++spins > (1u << 16)) sched_yield();
}
} else if (!req.event_arrival) {
const uint32_t slot =
req.layer * DS4_GPU_TP_GATES_PER_LAYER + req.gate;
uint32_t want = (uint32_t)req.seq;
if (req.rows > 0)
want = DS4_TP_BATCH_FLAG_TAG | (uint32_t)req.seq;
while (__atomic_load_n(&g_tp_gpu_flags[slot], __ATOMIC_ACQUIRE) != want) {
if (g_tp_shutdown) break;
if (++spins > (1u << 20)) {
sched_yield();
spins = 0;
}
}
} else if (req.rows > 0) {
while (g_tp_batch_gpu_event.signaledValue < req.seq) {
if (g_tp_shutdown) break;
if (++spins > (1u << 16)) sched_yield();
}
} else {
while (g_tp_gpu_event.signaledValue < req.seq) {
if (g_tp_shutdown) break;
if (++spins > (1u << 16)) sched_yield();
}
}
const double t1 = ds4_gpu_now_ms();
if (profile) g_tp_stat_encode_lead += g_tp_seq - req.seq;
if (req.poll && g_tp_check_words && g_tp_slab_cpu && (req.rows == 0 || req.check) &&
req.big_bytes == 0) {
/* The flag line may land before the payload lines: verify the
* partial against the checksum the flag kernel published. */
const uint32_t slot = req.layer * DS4_GPU_TP_GATES_PER_LAYER + req.gate;
const uint32_t words = req.rows ? req.check_words : (uint32_t)(g_tp_vec_bytes / 4u);
const volatile uint32_t *payload = req.rows ? req.check :
(const volatile uint32_t *)(g_tp_slab_cpu + g_tp_out_off + (uint64_t)slot * g_tp_vec_bytes);
const uint32_t mix = (req.rows ? DS4_TP_BATCH_FLAG_TAG | (uint32_t)req.seq :
(uint32_t)req.seq) * 0x9E3779B9u;
uint64_t tries = 0;
for (;;) {
const uint32_t want = __atomic_load_n(&g_tp_check_words[slot], __ATOMIC_ACQUIRE) ^ mix;
uint32_t sum = 0;
for (uint32_t i = 0; i < words; i++) sum += payload[i];
if (sum == want) break;
if (++tries > 20000000ull || g_tp_shutdown) {
if (!g_tp_failed_flag)
fprintf(stderr, "ds4: TP gate payload never matched its checksum (seq %llu)\n",
(unsigned long long)req.seq);
g_tp_failed_flag = 1;
break;
}
}
if (profile) g_tp_stat_cbwait_ms += ds4_gpu_now_ms() - t1;
}
int ok = 0;
if (!g_tp_shutdown && !g_tp_failed_flag) {
if (req.big_bytes > 0) {
if (g_tp_big_exchange_fn)
ok = g_tp_big_exchange_fn(g_tp_exchange_ud, req.layer,
req.seq, req.big_out,
req.big_in, req.big_bytes);
} else if (req.rows > 0) {
const double bx0 = profile ? ds4_gpu_now_ms() : 0.0;
if (g_tp_batch_exchange_fn)
ok = g_tp_batch_exchange_fn(g_tp_exchange_ud, req.batch_slot,
req.rows, req.seq);
if (profile) {
g_tp_stat_batch_gates++;
g_tp_stat_batch_gpu_wait_ms += t1 - t0;
g_tp_stat_batch_exchange_ms += ds4_gpu_now_ms() - bx0;
}
} else if (g_tp_exchange_fn) {
ok = g_tp_exchange_fn(g_tp_exchange_ud, req.layer, req.gate,
req.seq);
}
}
if (!ok && !g_tp_shutdown) {
if (!g_tp_failed_flag)
fprintf(stderr, "ds4: TP gate exchange failed (layer %u gate %u seq %llu)\n",
req.layer, req.gate, (unsigned long long)req.seq);
g_tp_failed_flag = 1;
}
const double t_ex = ds4_gpu_now_ms();
if (req.rows == 0 && req.big_bytes == 0 && req.gate < 2u) {
/* Sized the gate-time prefetch: how long this rank typically
* waits for the peer at this gate kind (flag seen -> exchange done). */
const double ex_us = (t_ex - t1) * 1000.0;
g_tp_exchange_ewma_us[req.gate] =
g_tp_exchange_ewma_us[req.gate] == 0.0 ? ex_us :
0.9 * g_tp_exchange_ewma_us[req.gate] + 0.1 * ex_us;
}
/* Release the GPU even on failure so end_commands can drain. */
if (req.spin && g_tp_release_ptr) {
const uint32_t prev_spin = g_tp_release_ptr[1];
__atomic_store_n(&g_tp_release_ptr[0], (uint32_t)req.seq, __ATOMIC_SEQ_CST);
/* a timed-out wait opened its gate without the peer partial */
if (g_tp_release_ptr[2] != 0 && !g_tp_shutdown) {
if (!g_tp_failed_flag)
fprintf(stderr, "ds4: TP release spin timed out before seq %llu (%u gates); failing the eval\n",
(unsigned long long)req.seq, g_tp_release_ptr[2]);
g_tp_failed_flag = 1;
}
if (profile) g_tp_stat_spin_iters += prev_spin;
} else if (req.rows > 0 && !req.poll) {
const double br0 = profile ? ds4_gpu_now_ms() : 0.0;
g_tp_batch_cpu_event.signaledValue = req.seq;
if (profile) g_tp_stat_batch_release_ms += ds4_gpu_now_ms() - br0;
} else if (req.poll) {
/* The flag we just consumed was published by the command buffer
* that ran the previous poll, so that poll's result is visible:
* a timeout there means the GPU has already run on stale data. */
const uint32_t slot = req.layer * DS4_GPU_TP_GATES_PER_LAYER + req.gate;
if (profile && g_tp_poll_prev_valid) {
const uint32_t hit = g_tp_poll_status[g_tp_poll_prev_slot * 2u];
if (hit != 0xffffffffu) {
g_tp_stat_poll_lines += hit;
if (hit > g_tp_stat_poll_max_line) g_tp_stat_poll_max_line = hit;
g_tp_stat_poll_count++;
}
}
if (g_tp_poll_prev_valid &&
g_tp_poll_status[g_tp_poll_prev_slot * 2u] == 0xffffffffu &&
!g_tp_shutdown) {
if (!g_tp_failed_flag)
fprintf(stderr,
"ds4: TP poll gate timed out (seq %llu, slot %u)\n",
(unsigned long long)g_tp_poll_prev_seq,
g_tp_poll_prev_slot);
g_tp_failed_flag = 1;
}
const uint32_t v = req.rows ? DS4_TP_BATCH_FLAG_TAG | (uint32_t)req.seq : (uint32_t)req.seq;
const double rel_t0 = profile ? ds4_gpu_now_ms() : 0.0;
volatile uint32_t *region =
g_tp_poll_region +
(size_t)(req.seq % DS4_TP_POLL_RING) *
DS4_TP_POLL_LINES * (DS4_TP_POLL_LINE_BYTES / 4u);
/* the poll kernel probes 32 lines per round and releases on any
* match: one marker per round is enough */
for (uint32_t j = 0; j < DS4_TP_POLL_LINES; j += 32u) {
__atomic_store_n(&region[j * (DS4_TP_POLL_LINE_BYTES / 4u)],
v, __ATOMIC_RELAXED);
}
__atomic_thread_fence(__ATOMIC_SEQ_CST);
if (profile) g_tp_stat_release_ms += ds4_gpu_now_ms() - rel_t0;
g_tp_poll_prev_valid = 1;
g_tp_poll_prev_slot = slot;
g_tp_poll_prev_seq = req.seq;
} else {
g_tp_cpu_event.signaledValue = req.seq;
}
if (!req.poll) g_tp_poll_prev_valid = 0;
if (profile) {
const double gpu_wait_ms = t1 - t0;
const double exchange_ms = t_ex - t1;
g_tp_stat_gpu_wait_ms += gpu_wait_ms;
g_tp_stat_exchange_ms += exchange_ms;
if (req.gate < DS4_GPU_TP_GATES_PER_LAYER) {
g_tp_stat_gate_count[req.gate]++;
g_tp_stat_gate_gpu_wait_ms[req.gate] += gpu_wait_ms;
g_tp_stat_gate_exchange_ms[req.gate] += exchange_ms;
}
if (++g_tp_stat_gates % 860 == 0) {
const double attn_count =
(double)(g_tp_stat_gate_count[DS4_GPU_TP_GATE_ATTN]
? g_tp_stat_gate_count[DS4_GPU_TP_GATE_ATTN] : 1u);
const double ffn_count =
(double)(g_tp_stat_gate_count[DS4_GPU_TP_GATE_FFN]
? g_tp_stat_gate_count[DS4_GPU_TP_GATE_FFN] : 1u);
fprintf(stderr,
"ds4: TP gates %llu: avg gpu-wait %.1f us, exchange %.1f us; "
"attn %.1f/%.1f us, ffn %.1f/%.1f us\n",
(unsigned long long)g_tp_stat_gates,
g_tp_stat_gpu_wait_ms / (double)g_tp_stat_gates * 1000.0,
g_tp_stat_exchange_ms / (double)g_tp_stat_gates * 1000.0,
g_tp_stat_gate_gpu_wait_ms[DS4_GPU_TP_GATE_ATTN] /
attn_count * 1000.0,
g_tp_stat_gate_exchange_ms[DS4_GPU_TP_GATE_ATTN] /
attn_count * 1000.0,
g_tp_stat_gate_gpu_wait_ms[DS4_GPU_TP_GATE_FFN] /
ffn_count * 1000.0,
g_tp_stat_gate_exchange_ms[DS4_GPU_TP_GATE_FFN] /
ffn_count * 1000.0);
if (g_tp_stat_batch_gates)
fprintf(stderr, "ds4: TP batch gates: %llu, arrival wait %.1f us, exchange %.1f us, release %.1f us per gate\n",
(unsigned long long)g_tp_stat_batch_gates,
g_tp_stat_batch_gpu_wait_ms * 1e3 / (double)g_tp_stat_batch_gates,
g_tp_stat_batch_exchange_ms * 1e3 / (double)g_tp_stat_batch_gates,
g_tp_stat_batch_release_ms * 1e3 / (double)g_tp_stat_batch_gates);
fprintf(stderr,
"ds4: TP gates: encode lead %.2f gates; verify %.1f us; release %.1f us; poll hit line avg %.1f max %llu over %llu\n",
(double)g_tp_stat_encode_lead / (double)g_tp_stat_gates,
g_tp_stat_cbwait_ms / (double)g_tp_stat_gates * 1000.0,
g_tp_stat_release_ms / (double)g_tp_stat_gates * 1000.0,
g_tp_stat_poll_count ?
(double)g_tp_stat_poll_lines / (double)g_tp_stat_poll_count : 0.0,
(unsigned long long)g_tp_stat_poll_max_line,
(unsigned long long)g_tp_stat_poll_count);
}
}
}
return NULL;
}
int ds4_gpu_tp_init(uint32_t rank,
ds4_gpu_tensor *slab, uint64_t gpu_flags_off,
uint64_t out_off, uint64_t vec_bytes,
ds4_gpu_tp_exchange_fn fn, void *ud) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (g_tp_thread_running || rank > 1) return 0;
g_tp_split_rank = (int32_t)rank;
g_tp_split_world = 2;
g_tp_slab_buffer = slab ? ds4_gpu_tensor_buffer(slab) : nil;
g_tp_slab_buffer_off = slab ? ds4_gpu_tensor_offset(slab) : 0;
g_tp_gpu_flags_off = gpu_flags_off;
g_tp_out_off = out_off;
g_tp_vec_bytes = vec_bytes;
g_tp_slab_cpu = slab ? (volatile uint8_t *)ds4_gpu_tensor_contents(slab) : NULL;
g_tp_check_buffer = [g_device newBufferWithLength:(NSUInteger)DS4_TP_POLL_MAX_SLOTS * sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (g_tp_check_buffer) {
memset([g_tp_check_buffer contents], 0, (size_t)DS4_TP_POLL_MAX_SLOTS * sizeof(uint32_t));
g_tp_check_words = (volatile uint32_t *)[g_tp_check_buffer contents];
} else {
g_tp_check_words = NULL;
}
g_tp_gpu_flags = slab ?
(volatile uint32_t *)((uint8_t *)ds4_gpu_tensor_contents(slab) + gpu_flags_off) : NULL;
/* Flag arrival is the default: the slab-word publish detects in ~1-3us
* where signaledValue polling costs 10-20, worth +2.3 t/s on the pair
* (A/B 2026-07-06, byte-identical output). DS4_TP_EVENT_GATES falls
* back to the shared-event arrival path. */
g_tp_flag_gates = g_tp_gpu_flags != NULL && getenv("DS4_TP_EVENT_GATES") == NULL;
g_tp_poll_gates = false;
g_tp_poll_buffer = nil;
g_tp_poll_region = NULL;
g_tp_poll_status = NULL;
g_tp_poll_prev_valid = 0;
g_tp_gate_prefetch = getenv("DS4_TP_DISABLE_GATE_PREFETCH") == NULL;
memset(g_tp_prefetch_count, 0, sizeof(g_tp_prefetch_count));
/* Default with flag gates: byte-identical to the event wait and about
* 20 us cheaper per gate on the M5 Max pair (A/B 2026-09-01, README
* ctx 128: 46.6-47.2 vs 43.6 t/s). DS4_TP_DISABLE_POLL_GATES restores
* the shared-event wait for diagnosis. */
if (g_tp_flag_gates && getenv("DS4_TP_DISABLE_POLL_GATES") == NULL) {
const NSUInteger region_bytes =
(NSUInteger)DS4_TP_POLL_LINES * DS4_TP_POLL_LINE_BYTES * DS4_TP_POLL_RING;
const NSUInteger status_bytes =
(NSUInteger)DS4_TP_POLL_MAX_SLOTS * 2u * sizeof(uint32_t);
g_tp_poll_buffer =
[g_device newBufferWithLength:region_bytes + status_bytes
options:MTLResourceStorageModeShared];
if (g_tp_poll_buffer) {
memset([g_tp_poll_buffer contents], 0, region_bytes + status_bytes);
g_tp_poll_region = (volatile uint32_t *)[g_tp_poll_buffer contents];
g_tp_poll_status =
(volatile uint32_t *)((uint8_t *)[g_tp_poll_buffer contents] + region_bytes);
g_tp_poll_gates = true;
} else {
fprintf(stderr, "ds4: TP poll gate buffer failed; using event gates\n");
}
}
g_tp_release_buffer = nil;
g_tp_release_ptr = NULL;
g_tp_stat_spin_iters = 0;
if (g_tp_flag_gates && getenv("DS4_TP_DISABLE_DECODE_POLL_GATES") == NULL &&
ds4_gpu_get_pipeline("kernel_dsv4_tp_release_wait") != nil &&
ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_publish") != nil &&
ds4_gpu_get_pipeline("kernel_dsv4_tp_payload_publish") != nil) {
g_tp_release_buffer = [g_device newBufferWithLength:8 * sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (g_tp_release_buffer) {
memset(g_tp_release_buffer.contents, 0, 8 * sizeof(uint32_t));
g_tp_release_ptr = (volatile uint32_t *)g_tp_release_buffer.contents;
g_tp_spin_max_iters =
(uint32_t)ds4_gpu_env_u64("DS4_TP_INLINE_SPIN_ITERS", 8000000u, 100000u, 200000000u);
}
}
g_tp_gpu_event = [g_device newSharedEvent];
g_tp_cpu_event = [g_device newSharedEvent];
g_tp_batch_gpu_event = [g_device newSharedEvent];
g_tp_batch_cpu_event = [g_device newSharedEvent];
if (!g_tp_gpu_event || !g_tp_cpu_event ||
!g_tp_batch_gpu_event || !g_tp_batch_cpu_event) {
fprintf(stderr, "ds4: failed to create TP shared events\n");
return 0;
}
g_tp_exchange_fn = fn;
g_tp_exchange_ud = ud;
g_tp_seq = 0;
g_tp_batch_seq = 0;
g_tp_shutdown = 0;
g_tp_failed_flag = 0;
g_tp_queue_head = 0;
g_tp_queue_count = 0;
g_tp_stat_gates = 0;
g_tp_stat_gpu_wait_ms = 0.0;
g_tp_stat_exchange_ms = 0.0;
memset(g_tp_stat_gate_count, 0, sizeof(g_tp_stat_gate_count));
memset(g_tp_stat_gate_gpu_wait_ms, 0,
sizeof(g_tp_stat_gate_gpu_wait_ms));
memset(g_tp_stat_gate_exchange_ms, 0,
sizeof(g_tp_stat_gate_exchange_ms));
pthread_attr_t attr;
pthread_attr_init(&attr);
pthread_attr_set_qos_class_np(&attr, QOS_CLASS_USER_INTERACTIVE, 0);
if (pthread_create(&g_tp_thread, &attr, ds4_gpu_tp_service_thread, NULL) != 0) {
pthread_attr_destroy(&attr);
fprintf(stderr, "ds4: failed to start TP gate service thread\n");
return 0;
}
pthread_attr_destroy(&attr);
g_tp_thread_running = 1;
/* The spin kernel keeps the GPU busy; the keep-alive kernels would
* timeslice against it. */
if (getenv("DS4_TP_NO_KEEPALIVE") == NULL && !ds4_gpu_tp_decode_inline_gates()) {
uint32_t ka_tgs = ds4_gpu_tp_keepalive_tgs_from_env();
g_tp_keepalive_queue = [g_device newCommandQueue];
g_tp_keepalive_buffer = [g_device newBufferWithLength:(NSUInteger)ka_tgs * 256u * sizeof(float)
options:MTLResourceStorageModeShared];
if (g_tp_keepalive_queue && g_tp_keepalive_buffer &&
pthread_create(&g_tp_keepalive_thread, NULL,
ds4_gpu_tp_keepalive_thread, NULL) == 0) {
g_tp_keepalive_running = 1;
} else {
fprintf(stderr, "ds4: TP keep-alive setup failed (continuing without)\n");
}
}
return 1;
}
void ds4_gpu_tp_shutdown(void) {
if (!g_tp_thread_running) return;
pthread_mutex_lock(&g_tp_mutex);
g_tp_shutdown = 1;
pthread_cond_broadcast(&g_tp_cond);
pthread_mutex_unlock(&g_tp_mutex);
pthread_join(g_tp_thread, NULL);
g_tp_thread_running = 0;
if (g_tp_keepalive_running) {
pthread_join(g_tp_keepalive_thread, NULL);
g_tp_keepalive_running = 0;
g_tp_keepalive_queue = nil;
g_tp_keepalive_buffer = nil;
}
g_tp_exchange_fn = NULL;
g_tp_batch_exchange_fn = NULL;
g_tp_exchange_ud = NULL;
g_tp_split_rank = 0;
g_tp_split_world = 1;
g_tp_session_batch_mode = 0;
g_tp_poll_gates = false;
g_tp_poll_buffer = nil;
g_tp_poll_region = NULL;
g_tp_poll_status = NULL;
g_tp_release_buffer = nil;
g_tp_release_ptr = NULL;
}
void ds4_gpu_tp_suspend_expert_sharding(int suspend) {
if (!g_tp_thread_running) return;
g_tp_split_world = suspend ? 1 : 2;
}
/* Fold request for producers that write the TP partial with a generic matvec
* (the attention output K-slice): the layer code registers (layer, gate)
* right before the producer call; the producer takes the request when its
* output is that slot and publishes the checked flag itself. */
static int g_tp_fold_req_active;
static uint32_t g_tp_fold_req_layer;
static uint32_t g_tp_fold_req_gate;
static bool ds4_gpu_tp_flag_fold_ok(uint32_t slot, uint64_t bytes) {
return g_initialized && g_batch_cb && g_tp_thread_running &&
!g_tp_session_batch_mode && g_tp_flag_gates && g_tp_poll_gates &&
!g_ssd_streaming_mode && g_tp_check_words != NULL &&
slot < DS4_TP_POLL_MAX_SLOTS && g_tp_vec_bytes == bytes &&
getenv("DS4_TP_DISABLE_FLAG_FOLD") == NULL;
}
static bool ds4_gpu_tp_fold_ctl_ready(void) {
if (g_tp_fold_ctl) return true;
g_tp_fold_ctl = [g_device newBufferWithLength:(NSUInteger)DS4_TP_POLL_MAX_SLOTS * 8u
options:MTLResourceStorageModeShared];
if (!g_tp_fold_ctl) return false;
memset(g_tp_fold_ctl.contents, 0, (size_t)DS4_TP_POLL_MAX_SLOTS * 8u);
return true;
}
int ds4_gpu_tp_decode_inline_gates(void) {
return g_tp_flag_gates && !g_tp_session_batch_mode && g_tp_release_ptr != NULL;
}
void ds4_gpu_tp_flag_fold_request(uint32_t layer, uint32_t gate) {
g_tp_fold_req_active = 0;
if (ds4_gpu_tp_decode_inline_gates()) return; /* the inline gate publishes coherently */
const uint32_t slot = layer * DS4_GPU_TP_GATES_PER_LAYER + gate;
static int debug_left = -1;
if (debug_left < 0) debug_left = getenv("DS4_TP_FOLD_DEBUG") ? 8 : 0;
if (debug_left > 0) {
debug_left--;
fprintf(stderr, "ds4: fold request slot %u ok=%d init=%d cb=%d thr=%d batch=%d flag=%d poll=%d ssd=%d chk=%d vec=%llu\n",
slot, (int)ds4_gpu_tp_flag_fold_ok(slot, g_tp_vec_bytes), (int)g_initialized,
(int)(g_batch_cb != nil), (int)g_tp_thread_running, (int)g_tp_session_batch_mode,
(int)g_tp_flag_gates, (int)g_tp_poll_gates, (int)g_ssd_streaming_mode,
(int)(g_tp_check_words != NULL), (unsigned long long)g_tp_vec_bytes);
}
if (!ds4_gpu_tp_flag_fold_ok(slot, g_tp_vec_bytes) || g_tp_vec_bytes == 0) return;
g_tp_fold_req_active = 1;
g_tp_fold_req_layer = layer;
g_tp_fold_req_gate = gate;
}
/* Returns the slot when `out` (out_bytes long) is the requested partial. */
static int ds4_gpu_tp_flag_fold_take(const ds4_gpu_tensor *out, uint64_t out_bytes,
uint32_t *slot_out, uint32_t *value_out) {
if (!g_tp_fold_req_active || !out) return 0;
const uint32_t slot = g_tp_fold_req_layer * DS4_GPU_TP_GATES_PER_LAYER + g_tp_fold_req_gate;
static int debug_left = -1;
if (debug_left < 0) debug_left = getenv("DS4_TP_FOLD_DEBUG") ? 8 : 0;
if (debug_left > 0) {
debug_left--;
fprintf(stderr, "ds4: fold take slot %u ok=%d buf_eq=%d off=%llu want=%llu bytes=%llu vec=%llu\n",
slot, (int)ds4_gpu_tp_flag_fold_ok(slot, out_bytes),
(int)(ds4_gpu_tensor_buffer(out) == g_tp_slab_buffer),
(unsigned long long)ds4_gpu_tensor_offset(out),
(unsigned long long)(g_tp_slab_buffer_off + g_tp_out_off + (uint64_t)slot * g_tp_vec_bytes),
(unsigned long long)out_bytes, (unsigned long long)g_tp_vec_bytes);
}
if (!ds4_gpu_tp_flag_fold_ok(slot, out_bytes) ||
ds4_gpu_tensor_buffer(out) != g_tp_slab_buffer ||
ds4_gpu_tensor_offset(out) !=
(uint64_t)(g_tp_slab_buffer_off + g_tp_out_off + (uint64_t)slot * g_tp_vec_bytes) ||
!ds4_gpu_tp_fold_ctl_ready()) {
return 0;
}
g_tp_fold_req_active = 0;
*slot_out = slot;
*value_out = (uint32_t)(g_tp_seq + 1u);
return 1;
}
int ds4_gpu_add_tensor_tp_flag(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
uint32_t n,
uint32_t layer,
uint32_t gate) {
const uint32_t slot = layer * DS4_GPU_TP_GATES_PER_LAYER + gate;
const bool fold =
g_initialized && g_batch_cb && g_tp_thread_running &&
!g_tp_session_batch_mode && g_tp_flag_gates && g_tp_poll_gates &&
!ds4_gpu_tp_decode_inline_gates() && /* the inline gate publishes coherently */
!g_ssd_streaming_mode && g_tp_check_words != NULL &&
slot < DS4_TP_POLL_MAX_SLOTS && g_tp_vec_bytes == (uint64_t)n * 4u &&
out && a && b && n != 0 && (n % 256u) == 0u &&
ds4_gpu_tensor_buffer(out) == g_tp_slab_buffer &&
ds4_gpu_tensor_offset(out) ==
(uint64_t)(g_tp_slab_buffer_off + g_tp_out_off + (uint64_t)slot * g_tp_vec_bytes) &&
getenv("DS4_TP_DISABLE_FLAG_FOLD") == NULL;
if (!fold) return ds4_gpu_add_tensor(out, a, b, n);
if (!ds4_gpu_tp_fold_ctl_ready()) return ds4_gpu_add_tensor(out, a, b, n);
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_add2_f32_tp_flag_checked");
if (!pipeline) return ds4_gpu_add_tensor(out, a, b, n);
id<MTLBuffer> abuf = ds4_gpu_tensor_buffer(a);
id<MTLBuffer> bbuf = ds4_gpu_tensor_buffer(b);
if (!abuf || !bbuf) return ds4_gpu_add_tensor(out, a, b, n);
const uint32_t value = (uint32_t)(g_tp_seq + 1u);
const uint32_t ntg = n / 256u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(g_batch_cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&n length:sizeof(n) atIndex:0];
[enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1];
[enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2];
[enc setBuffer:g_tp_slab_buffer offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setBuffer:g_tp_slab_buffer
offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u)
atIndex:4];
[enc setBuffer:g_tp_check_buffer offset:(NSUInteger)slot * 4u atIndex:5];
[enc setBytes:&value length:sizeof(value) atIndex:6];
[enc setBuffer:g_tp_fold_ctl offset:(NSUInteger)slot * 8u atIndex:7];
[enc setBytes:&ntg length:sizeof(ntg) atIndex:8];
[enc setThreadgroupMemoryLength:8 * sizeof(uint32_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(ntg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
g_tp_flag_prepublished_seq = g_tp_seq + 1u;
return 1;
}
/* Inline gate: publish the partial (n_words at slab offset out_off) and the
* flag word of `slot` coherently, stream the next phase's weights while the
* exchange runs, then spin on the release word inside this command buffer. */
static int ds4_gpu_tp_spin_gate_encode(uint32_t slot, uint32_t value,
uint64_t out_off, uint32_t n_words,
uint32_t prefetch_gate) {
id<MTLComputePipelineState> publish = ds4_gpu_get_pipeline("kernel_dsv4_tp_payload_publish");
id<MTLComputePipelineState> flag = ds4_gpu_get_pipeline("kernel_dsv4_tp_flag_publish");
id<MTLComputePipelineState> wait = ds4_gpu_get_pipeline("kernel_dsv4_tp_release_wait");
if (!publish || !flag || !wait || !g_tp_release_buffer) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(g_batch_cb);
if (g_batch_encoder_concurrent) [enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
[enc setComputePipelineState:publish];
[enc setBuffer:g_tp_slab_buffer offset:(NSUInteger)(g_tp_slab_buffer_off + out_off) atIndex:0];
[enc setBytes:&n_words length:sizeof(n_words) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake((n_words + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
[enc setComputePipelineState:flag];
[enc setBuffer:g_tp_slab_buffer
offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u)
atIndex:0];
[enc setBytes:&value length:sizeof(value) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
ds4_gpu_end_compute_encoder(g_batch_cb, enc);
ds4_gpu_close_batch_encoder();
if (prefetch_gate < 2u) (void)ds4_gpu_tp_encode_gate_prefetch(prefetch_gate);
enc = ds4_gpu_compute_encoder(g_batch_cb);
[enc setComputePipelineState:wait];
[enc setBuffer:g_tp_release_buffer offset:0 atIndex:0];
[enc setBytes:&value length:sizeof(value) atIndex:1];
[enc setBytes:&g_tp_spin_max_iters length:sizeof(g_tp_spin_max_iters) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
/* Hazard tracking cannot see the CPU-side dependency that orders the
* combine reads after the spin. */
[enc memoryBarrierWithScope:MTLBarrierScopeBuffers];
ds4_gpu_end_compute_encoder(g_batch_cb, enc);
return 1;
}
int ds4_gpu_tp_gate_encode(uint32_t layer, uint32_t gate) {
/* The C graph has no outer pool for flag/poll encoder temporaries. */
@autoreleasepool {
if (!g_batch_cb) {
fprintf(stderr, "ds4: TP gate encode without an open command batch (layer %u gate %u)\n",
layer, gate);
return 0;
}
if (!g_tp_thread_running) {
fprintf(stderr, "ds4: TP gate encode without the gate service (layer %u)\n", layer);
return 0;
}
g_tp_fold_req_active = 0;
const uint64_t seq = ++g_tp_seq;
const uint32_t slot = layer * DS4_GPU_TP_GATES_PER_LAYER + gate;
const bool spin_gate = ds4_gpu_tp_decode_inline_gates() && !g_ssd_streaming_mode &&
g_tp_vec_bytes != 0;
const bool event_arrival = !g_tp_flag_gates || (g_tp_session_batch_mode && !spin_gate);
if (spin_gate) {
if (!ds4_gpu_tp_spin_gate_encode(slot, (uint32_t)seq,
g_tp_out_off + (uint64_t)slot * g_tp_vec_bytes,
(uint32_t)(g_tp_vec_bytes / 4u), gate))
return 0;
g_tp_flag_prepublished_seq = 0;
} else if (!event_arrival) {
/* Publish arrival through the slab word; the buffer hazard against
* the partial-output kernels orders the store after the payload. */
const uint32_t value = (uint32_t)seq;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) return 0;
const bool checked =
g_tp_poll_gates && !g_ssd_streaming_mode && g_tp_check_words != NULL &&
slot < DS4_TP_POLL_MAX_SLOTS && g_tp_vec_bytes != 0;
/* The payload producer may already have published this flag. */
const bool prepublished = checked && g_tp_flag_prepublished_seq == seq;
g_tp_flag_prepublished_seq = 0;
if (prepublished) {
ds4_gpu_close_batch_encoder();
} else {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
checked ? "kernel_dsv4_tp_flag_set_checked" : "kernel_dsv4_tp_flag_set");
if (!pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBuffer:g_tp_slab_buffer
offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u)
atIndex:0];
if (checked) {
const uint32_t words = (uint32_t)(g_tp_vec_bytes / 4u);
[enc setBuffer:g_tp_check_buffer offset:(NSUInteger)slot * 4u atIndex:1];
[enc setBytes:&value length:sizeof(value) atIndex:2];
[enc setBuffer:g_tp_slab_buffer
offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_out_off + (uint64_t)slot * g_tp_vec_bytes)
atIndex:3];
[enc setBytes:&words length:sizeof(words) atIndex:4];
[enc setThreadgroupMemoryLength:8 * sizeof(uint32_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
} else {
[enc setBytes:&value length:sizeof(value) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
ds4_gpu_close_batch_encoder();
}
} else {
ds4_gpu_close_batch_encoder();
[g_batch_cb encodeSignalEvent:g_tp_gpu_event value:seq];
}
const bool poll_gate =
!spin_gate && !event_arrival && g_tp_poll_gates && !g_ssd_streaming_mode &&
slot < DS4_TP_POLL_MAX_SLOTS;
if (poll_gate) {
/* Ending the command buffer is what makes the flag and the partial
* CPU-visible; the next buffer, already queued, opens with the poll
* so the GPU resumes within a memory round trip of the release. */
if (!ds4_gpu_flush_commands()) return 0;
/* Keep the queue shallow: the firmware only chains what the driver
* has handed it, and a deep backlog of many-encoder buffers is fed
* from the CPU with scheduling latency. */
if ([g_pending_cbs count] > DS4_TP_POLL_MAX_INFLIGHT) {
id<MTLCommandBuffer> old_cb =
g_pending_cbs[[g_pending_cbs count] - 1u - DS4_TP_POLL_MAX_INFLIGHT];
[old_cb waitUntilCompleted];
}
/* The flag is already published by the flushed buffer; stream the
* next phase's weights into the cache while the exchange runs. */
(void)ds4_gpu_tp_encode_gate_prefetch(gate);
id<MTLComputePipelineState> poll_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_tp_poll_release");
if (!poll_pipeline) return 0;
const uint32_t value = (uint32_t)seq;
const uint32_t nlines = DS4_TP_POLL_LINES;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(g_batch_cb);
[enc setComputePipelineState:poll_pipeline];
[enc setBuffer:g_tp_poll_buffer
offset:(NSUInteger)(seq % DS4_TP_POLL_RING) *
DS4_TP_POLL_LINES * DS4_TP_POLL_LINE_BYTES
atIndex:0];
[enc setBytes:&value length:sizeof(value) atIndex:1];
[enc setBytes:&nlines length:sizeof(nlines) atIndex:2];
[enc setBuffer:g_tp_poll_buffer
offset:(NSUInteger)DS4_TP_POLL_LINES * DS4_TP_POLL_LINE_BYTES *
DS4_TP_POLL_RING +
(NSUInteger)slot * 2u * sizeof(uint32_t)
atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(g_batch_cb, enc);
} else if (!spin_gate) {
[g_batch_cb encodeWaitForEvent:g_tp_cpu_event value:seq];
}
pthread_mutex_lock(&g_tp_mutex);
if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) {
pthread_mutex_unlock(&g_tp_mutex);
fprintf(stderr, "ds4: TP gate queue overflow\n");
return 0;
}
uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE;
g_tp_queue[tail].layer = layer;
g_tp_queue[tail].batch_slot = layer;
g_tp_queue[tail].gate = gate;
g_tp_queue[tail].rows = 0;
g_tp_queue[tail].poll = poll_gate ? 1u : 0u;
g_tp_queue[tail].spin = spin_gate ? 1u : 0u;
g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u;
g_tp_queue[tail].seq = seq;
g_tp_queue[tail].big_out = NULL;
g_tp_queue[tail].big_in = NULL;
g_tp_queue[tail].big_bytes = 0;
g_tp_queue_count++;
pthread_cond_signal(&g_tp_cond);
pthread_mutex_unlock(&g_tp_mutex);
return 1;
}
}
void ds4_gpu_tp_set_batch_exchange(ds4_gpu_tp_batch_exchange_fn fn) {
g_tp_batch_exchange_fn = fn;
}
/* Verify-block batch gate: same arrival/release machinery as the row gate
* (the FFN flag word and event pair are reused — a decode gate and a batch
* gate are never in flight together, and seq values stay globally unique),
* but the service thread runs the multi-row exchange callback. */
/* A verify-block gate publishes its arrival on the (layer, gate) flag word
* and exchanges the rows of slab batch slot `batch_slot`. */
static int ds4_gpu_tp_batch_gate_encode_slot(uint32_t layer, uint32_t gate,
uint32_t batch_slot, uint32_t rows,
const ds4_gpu_tensor *payload) {
@autoreleasepool {
if (!g_batch_cb) return 0;
if (!g_tp_thread_running || rows == 0) return 0;
const uint64_t seq = ++g_tp_batch_seq;
const bool event_arrival = g_tp_session_batch_mode || !g_tp_flag_gates;
const uint32_t slot = layer * DS4_GPU_TP_GATES_PER_LAYER + gate;
const uint32_t value = DS4_TP_BATCH_FLAG_TAG | (uint32_t)seq;
const uint32_t words = (uint32_t)(rows * g_tp_vec_bytes / 4u);
/* With the payload at hand the gate spins like a decode gate: a checked
* flag on arrival, kernel_dsv4_tp_poll_release instead of the event wake. */
static int poll_off = -1;
if (poll_off < 0) poll_off = getenv("DS4_TP_DISABLE_BATCH_POLL_GATES") != NULL;
const bool poll_gate = !event_arrival && !poll_off && payload && g_tp_poll_gates &&
!g_ssd_streaming_mode && g_tp_check_words != NULL &&
slot < DS4_TP_POLL_MAX_SLOTS && words != 0;
if (!event_arrival) {
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) return 0;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
poll_gate ? "kernel_dsv4_tp_flag_set_checked" : "kernel_dsv4_tp_flag_set");
if (!pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBuffer:g_tp_slab_buffer
offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)slot * 4u)
atIndex:0];
if (poll_gate) {
[enc setBuffer:g_tp_check_buffer offset:(NSUInteger)slot * 4u atIndex:1];
[enc setBytes:&value length:sizeof(value) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer((ds4_gpu_tensor *)payload)
offset:ds4_gpu_tensor_offset((ds4_gpu_tensor *)payload) atIndex:3];
[enc setBytes:&words length:sizeof(words) atIndex:4];
[enc setThreadgroupMemoryLength:8 * sizeof(uint32_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
} else {
[enc setBytes:&value length:sizeof(value) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
ds4_gpu_close_batch_encoder();
} else {
ds4_gpu_close_batch_encoder();
[g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq];
}
if (poll_gate) {
if (!ds4_gpu_flush_commands()) return 0;
if ([g_pending_cbs count] > DS4_TP_POLL_MAX_INFLIGHT) {
id<MTLCommandBuffer> old_cb =
g_pending_cbs[[g_pending_cbs count] - 1u - DS4_TP_POLL_MAX_INFLIGHT];
[old_cb waitUntilCompleted];
}
id<MTLComputePipelineState> poll_pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_tp_poll_release");
if (!poll_pipeline) return 0;
const uint32_t nlines = DS4_TP_POLL_LINES;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(g_batch_cb);
[enc setComputePipelineState:poll_pipeline];
[enc setBuffer:g_tp_poll_buffer
offset:(NSUInteger)(seq % DS4_TP_POLL_RING) *
DS4_TP_POLL_LINES * DS4_TP_POLL_LINE_BYTES
atIndex:0];
[enc setBytes:&value length:sizeof(value) atIndex:1];
[enc setBytes:&nlines length:sizeof(nlines) atIndex:2];
[enc setBuffer:g_tp_poll_buffer
offset:(NSUInteger)DS4_TP_POLL_LINES * DS4_TP_POLL_LINE_BYTES *
DS4_TP_POLL_RING +
(NSUInteger)slot * 2u * sizeof(uint32_t)
atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(g_batch_cb, enc);
} else {
[g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq];
}
pthread_mutex_lock(&g_tp_mutex);
if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) {
pthread_mutex_unlock(&g_tp_mutex);
fprintf(stderr, "ds4: TP gate queue overflow\n");
return 0;
}
uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE;
g_tp_queue[tail].layer = layer;
g_tp_queue[tail].gate = gate;
g_tp_queue[tail].batch_slot = batch_slot;
g_tp_queue[tail].rows = rows;
g_tp_queue[tail].spin = 0;
g_tp_queue[tail].poll = poll_gate ? 1u : 0u;
g_tp_queue[tail].check = poll_gate ?
(const volatile uint32_t *)ds4_gpu_tensor_contents((ds4_gpu_tensor *)payload) : NULL;
g_tp_queue[tail].check_words = words;
g_tp_queue[tail].event_arrival = event_arrival ? 1u : 0u;
g_tp_queue[tail].seq = seq;
g_tp_queue[tail].big_out = NULL;
g_tp_queue[tail].big_in = NULL;
g_tp_queue[tail].big_bytes = 0;
g_tp_queue_count++;
pthread_cond_signal(&g_tp_cond);
pthread_mutex_unlock(&g_tp_mutex);
return 1;
}
}
int ds4_gpu_tp_batch_gate_encode(uint32_t layer, uint32_t rows) {
return ds4_gpu_tp_batch_gate_encode_slot(layer, DS4_GPU_TP_GATE_FFN, layer, rows, NULL);
}
/* Two gates per layer: slot layer*2+gate of the doubled batch region, whose
* rows `payload` (the slot's out view) holds. */
int ds4_gpu_tp_batch_gate_encode_kind(uint32_t layer, uint32_t gate, uint32_t rows,
const ds4_gpu_tensor *payload) {
return ds4_gpu_tp_batch_gate_encode_slot(layer, gate,
layer * DS4_GPU_TP_GATES_PER_LAYER + gate, rows, payload);
}
/* Prefill batch gate: the service thread exchanges big_bytes directly
* between two CPU-visible bounce buffers. The shared-event signal orders
* all preceding GPU stores before the service thread reads the payload. */
static uint64_t ds4_gpu_tp_big_gate_kick(uint32_t layer, uint32_t rows,
const ds4_gpu_tensor *out_t,
ds4_gpu_tensor *in_t,
uint64_t bytes) {
if (!g_batch_cb) return 0;
if (!g_tp_thread_running || rows == 0 || bytes == 0) return 0;
const void *out_ptr = ds4_gpu_tensor_contents((ds4_gpu_tensor *)out_t);
void *in_ptr = ds4_gpu_tensor_contents(in_t);
if (!out_ptr || !in_ptr) {
fprintf(stderr, "ds4: TP big gate needs CPU-visible bounce buffers\n");
return 0;
}
const uint64_t seq = ++g_tp_batch_seq;
ds4_gpu_close_batch_encoder();
[g_batch_cb encodeSignalEvent:g_tp_batch_gpu_event value:seq];
pthread_mutex_lock(&g_tp_mutex);
if (g_tp_queue_count >= DS4_GPU_TP_QUEUE) {
pthread_mutex_unlock(&g_tp_mutex);
fprintf(stderr, "ds4: TP gate queue overflow\n");
return 0;
}
uint32_t tail = (g_tp_queue_head + g_tp_queue_count) % DS4_GPU_TP_QUEUE;
g_tp_queue[tail].layer = layer;
g_tp_queue[tail].batch_slot = layer;
g_tp_queue[tail].gate = 1u;
g_tp_queue[tail].rows = rows;
g_tp_queue[tail].event_arrival = 1u;
g_tp_queue[tail].seq = seq;
g_tp_queue[tail].big_out = out_ptr;
g_tp_queue[tail].big_in = in_ptr;
g_tp_queue[tail].big_bytes = bytes;
g_tp_queue[tail].spin = 0;
g_tp_queue[tail].poll = 0;
g_tp_queue_count++;
pthread_cond_signal(&g_tp_cond);
pthread_mutex_unlock(&g_tp_mutex);
return seq;
}
static int ds4_gpu_tp_big_gate_wait(uint64_t seq) {
if (!g_batch_cb || seq == 0) return 0;
ds4_gpu_close_batch_encoder();
[g_batch_cb encodeWaitForEvent:g_tp_batch_cpu_event value:seq];
return 1;
}
int ds4_gpu_tp_big_gate_encode(uint32_t layer, uint32_t rows,
const ds4_gpu_tensor *out_t,
ds4_gpu_tensor *in_t,
uint64_t bytes) {
const uint64_t seq = ds4_gpu_tp_big_gate_kick(layer, rows, out_t, in_t, bytes);
if (seq == 0) return 0;
return ds4_gpu_tp_big_gate_wait(seq);
}
/* Run one trivial command buffer so the process's first-submission costs
* (command queue and GPU context setup, page-table validation of the
* mapped model views) are paid at load time rather than inside the first
* prefill. TP sharding skips the load-time model-view warm, which used to
* leave a 600-800 ms stall at the head of the first prompt on each rank. */
int ds4_gpu_warm_command_queue(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (g_batch_cb) return 1;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv4_tp_keepalive");
if (!pipeline) return 0;
const double t0 = ds4_gpu_now_ms();
@autoreleasepool {
id<MTLBuffer> scratch = [g_device newBufferWithLength:256u * sizeof(float)
options:MTLResourceStorageModeShared];
if (!scratch) return 0;
memset(scratch.contents, 0, 256u * sizeof(float));
id<MTLCommandBuffer> cb = ds4_gpu_new_command_buffer();
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
const uint32_t iters = 1u;
[enc setComputePipelineState:pipeline];
[enc setBuffer:scratch offset:0 atIndex:0];
[enc setBytes:&iters length:sizeof(iters) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
[enc endEncoding];
[cb commit];
[cb waitUntilCompleted];
if (cb.status != MTLCommandBufferStatusCompleted) return 0;
}
if (getenv("DS4_METAL_CB_TIMES"))
fprintf(stderr, "ds4: command queue warm in %.1f ms\n", ds4_gpu_now_ms() - t0);
ds4_gpu_queue_keepalive_start();
return 1;
}
/* Main-queue idle keepalive. Measured on M5 Max (macOS 26): once the
* engine's command queue has been idle for about 3 s, the next command
* buffer on it starts 600-800 ms late (the driver drops and rebuilds the
* process's GPU mappings; 76-145 GB of no-copy model views here). Work on
* other queues does not prevent it. One 1 KB kernel per second on this
* queue keeps the mappings alive at no measurable cost.
* DS4_METAL_DISABLE_QUEUE_KEEPALIVE=1 opts out. */
static pthread_t g_queue_keepalive_thread;
static volatile int g_queue_keepalive_running;
static volatile int g_queue_keepalive_stop;
static void *ds4_gpu_queue_keepalive_thread(void *arg) {
(void)arg;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv4_tp_keepalive");
id<MTLBuffer> scratch = [g_device newBufferWithLength:256u * sizeof(float)
options:MTLResourceStorageModeShared];
if (!pipeline || !scratch) return NULL;
memset(scratch.contents, 0, 256u * sizeof(float));
const uint32_t iters = 1u;
uint64_t period_ms = ds4_gpu_env_u64("DS4_METAL_QUEUE_KEEPALIVE_MS", 1000u, 50u, 10000u);
while (!g_queue_keepalive_stop) {
@autoreleasepool {
id<MTLCommandBuffer> cb = [g_queue commandBuffer];
id<MTLComputeCommandEncoder> enc = [cb computeCommandEncoder];
[enc setComputePipelineState:pipeline];
[enc setBuffer:scratch offset:0 atIndex:0];
[enc setBytes:&iters length:sizeof(iters) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
[enc endEncoding];
[cb commit];
}
for (uint64_t slept = 0; slept < period_ms && !g_queue_keepalive_stop; slept += 50)
usleep(50000);
}
return NULL;
}
static void ds4_gpu_queue_keepalive_start(void) {
if (g_queue_keepalive_running) return;
if (getenv("DS4_METAL_DISABLE_QUEUE_KEEPALIVE")) return;
g_queue_keepalive_stop = 0;
if (pthread_create(&g_queue_keepalive_thread, NULL, ds4_gpu_queue_keepalive_thread, NULL) == 0)
g_queue_keepalive_running = 1;
}
static void ds4_gpu_queue_keepalive_stop_thread(void) {
if (!g_queue_keepalive_running) return;
g_queue_keepalive_stop = 1;
pthread_join(g_queue_keepalive_thread, NULL);
g_queue_keepalive_running = 0;
}
int ds4_gpu_tp_failed(void) {
return g_tp_failed_flag;
}
int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (event_value == 0) return 0;
if (@available(macOS 12.0, *)) {
if (!g_selected_readback_event) return 0;
const char *what = label ? label : "selected-id readback";
const BOOL signaled =
[g_selected_readback_event waitUntilSignaledValue:event_value timeoutMS:60000];
if (!signaled) {
fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n", what);
return 0;
}
return 1;
}
fprintf(stderr, "ds4: selected-id overlap requires MTLSharedEvent support\n");
return 0;
}
static int ds4_gpu_signal_batch_and_wait_event(const char *label) {
if (!g_initialized && !ds4_gpu_init()) return 0;
ds4_gpu_parallel_ffn_reset_state(YES);
if (!g_batch_cb) return 0;
if (@available(macOS 12.0, *)) {
if (!g_selected_readback_event) {
g_selected_readback_event = [g_device newSharedEvent];
if (!g_selected_readback_event) {
fprintf(stderr, "ds4: failed to create Metal shared event for %s\n",
label ? label : "selected-id readback");
return 0;
}
}
ds4_gpu_close_batch_encoder();
id<MTLCommandBuffer> cb = g_batch_cb;
g_batch_cb = nil;
const uint64_t value = ++g_selected_readback_event_value;
[cb encodeSignalEvent:g_selected_readback_event value:value];
g_batch_has_work = YES;
[cb commit];
ds4_gpu_stream_expert_cache_note_batch_committed();
const BOOL signaled = [g_selected_readback_event waitUntilSignaledValue:value timeoutMS:60000];
[g_pending_cbs addObject:cb];
if (!signaled) {
fprintf(stderr, "ds4: timeout waiting for Metal shared event in %s\n",
label ? label : "selected-id readback");
(void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback");
[g_transient_buffers removeAllObjects];
return 0;
}
if (cb.status == MTLCommandBufferStatusError) {
fprintf(stderr, "ds4: Metal %s failed: %s\n",
label ? label : "selected-id readback",
[[cb.error localizedDescription] UTF8String]);
(void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback");
[g_transient_buffers removeAllObjects];
return 0;
}
g_batch_cb = ds4_gpu_new_command_buffer();
g_batch_has_work = NO;
if (g_batch_cb) ds4_gpu_stream_expert_cache_note_batch_created();
if (!g_batch_cb) {
(void)ds4_gpu_wait_pending_command_buffers(label ? label : "selected-id readback");
[g_transient_buffers removeAllObjects];
return 0;
}
return 1;
} else {
ds4_gpu_close_batch_encoder();
id<MTLCommandBuffer> cb = g_batch_cb;
g_batch_cb = nil;
g_batch_has_work = NO;
if (ds4_gpu_finish_command_buffer(cb, 1, label ? label : "selected-id readback") == 0) {
return 0;
}
return ds4_gpu_begin_commands();
}
}
int ds4_gpu_end_commands(void) {
if (!g_batch_cb) {
ds4_gpu_parallel_ffn_reset_state(YES);
return 0;
}
ds4_gpu_parallel_ffn_reset_state(YES);
ds4_gpu_close_batch_encoder();
id<MTLCommandBuffer> cb = g_batch_cb;
g_batch_cb = nil;
g_batch_has_work = NO;
g_stream_expert_cache_owned_seq = g_stream_expert_cache_batch_seq;
g_stream_expert_cache_batch_seq = 0;
return ds4_gpu_finish_command_buffer(cb, 1, "command batch");
}
static int ds4_gpu_flash_attn_stage_profile_boundary(
id<MTLCommandBuffer> __strong *cbp,
const char *mode,
const char *stage,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t n_keys,
uint32_t n_head,
uint32_t head_dim,
uint32_t window,
uint32_t ratio,
double *stage_t0) {
if (!cbp || !*cbp || !stage_t0 || !stage) return 0;
if (ds4_gpu_end_commands() == 0) return 0;
const double now_ms = ds4_gpu_now_ms();
const char *filter = getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE_FILTER");
const int print_stage =
!filter || !filter[0] ||
strstr(stage, filter) != NULL ||
(mode && strstr(mode, filter) != NULL);
if (print_stage) {
fprintf(stderr,
"ds4: Metal FlashAttention prefill stage mode=%s tokens=%u comp=%u "
"keys=%u heads=%u dim=%u window=%u ratio=%u %s=%.3f ms\n",
mode ? mode : "unknown",
n_tokens,
n_comp,
n_keys,
n_head,
head_dim,
window,
ratio,
stage,
now_ms - *stage_t0);
}
*stage_t0 = now_ms;
if (ds4_gpu_begin_commands() == 0) return 0;
int owned = 0;
*cbp = ds4_gpu_command_buffer(&owned);
return *cbp != nil && owned == 0;
}
int ds4_gpu_synchronize(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (g_batch_cb) return ds4_gpu_end_commands();
ds4_gpu_parallel_ffn_reset_state(YES);
if ([g_pending_cbs count] != 0) {
int ok = ds4_gpu_wait_pending_command_buffers("synchronize");
[g_transient_buffers removeAllObjects];
ds4_gpu_model_buffer_cache_maybe_evict("synchronize");
return ok;
}
id<MTLCommandBuffer> cb = ds4_gpu_new_command_buffer();
if (!cb) return 0;
return ds4_gpu_finish_command_buffer(cb, 1, "synchronize");
}
static void qwen4_nax_release_scratch(void);
static void qwen4_batch_release_scratch(void);
void ds4_gpu_cleanup(void) {
if (!g_initialized) return;
ds4_gpu_queue_keepalive_stop_thread();
@autoreleasepool {
ds4_gpu_decode_pipeline_fast_cache_reset();
ds4_gpu_parallel_ffn_reset_state(YES);
if (g_batch_cb) {
ds4_gpu_close_batch_encoder();
[g_batch_cb commit];
[g_batch_cb waitUntilCompleted];
g_batch_cb = nil;
if (g_stream_expert_cache_batch_seq > g_stream_expert_cache_done_seq) {
g_stream_expert_cache_done_seq = g_stream_expert_cache_batch_seq;
}
g_stream_expert_cache_batch_seq = 0;
}
(void)ds4_gpu_wait_pending_command_buffers("cleanup");
if (ds4_gpu_stream_expert_timing_summary_enabled() &&
getenv("DS4_METAL_MEMORY_REPORT") == NULL) {
ds4_gpu_print_memory_report("at cleanup");
}
g_selected_readback_event = nil;
g_selected_readback_event_value = 0;
[g_transient_buffers removeAllObjects];
ds4_gpu_stream_expert_pread_pool_shutdown();
ds4_gpu_stream_expert_cache_clear_all(1);
for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) {
g_stream_expert_cache_gate_addr_buffers[layer] = nil;
g_stream_expert_cache_up_addr_buffers[layer] = nil;
g_stream_expert_cache_down_addr_buffers[layer] = nil;
g_stream_compact_gate_addr_buffers[layer] = nil;
g_stream_compact_up_addr_buffers[layer] = nil;
g_stream_compact_down_addr_buffers[layer] = nil;
g_stream_compact_selected_buffers[layer] = nil;
g_stream_selected_id_buffers[layer] = nil;
}
g_set_rows_f32_i32_pipeline = nil;
g_get_rows_f32_pipeline = nil;
g_get_rows_f16_pipeline = nil;
g_get_rows_i32_pipeline = nil;
g_get_rows_q8_0_pipeline = nil;
g_get_rows_q4_0_pipeline = nil;
g_get_rows_q4_K_pipeline = nil;
g_repeat_f32_pipeline = nil;
g_concat_pipeline = nil;
g_cpy_f32_f32_pipeline = nil;
g_cpy_f32_f16_pipeline = nil;
g_cpy_contig_f32_f16_pipeline = nil;
g_cpy_f16_f32_pipeline = nil;
g_cpy_f16_f16_pipeline = nil;
g_cpy_contig_f16_f32_pipeline = nil;
g_cpy_contig_f16_f16_pipeline = nil;
g_flash_kv_stage_f16_pipeline = nil;
g_swiglu_pipeline = nil;
g_swiglu_flat_pipeline = nil;
g_add_pipeline = nil;
g_add2_pipeline = nil;
g_add3_pipeline = nil;
g_moe_sum6_pipeline = nil;
g_moe_sum8_pipeline = nil;
g_mul_pipeline = nil;
g_bin_mul_scalar_pipeline = nil;
g_bin_div_row_pipeline = nil;
g_unary_sigmoid_pipeline = nil;
g_unary_silu_pipeline = nil;
g_unary_softplus_pipeline = nil;
g_unary_sqrt_pipeline = nil;
g_unary_clamp_pipeline = nil;
g_unary_scale_pipeline = nil;
g_unary_fill_pipeline = nil;
g_unary_fill_f16_pipeline = nil;
g_rms_norm_pipeline = nil;
g_rms_norm_plain_pipeline = nil;
g_add_rms_norm_pipeline = nil;
g_rms_norm_scale_pipeline = nil;
g_dsv4_qkv_rms_norm_pipeline = nil;
g_dsv4_head_rms_norm_rope_tail_pipeline = nil;
g_use_dsv4_head_rms_norm_rope_tail_pipeline = false;
g_hc_split_sinkhorn_pipeline = nil;
g_hc_split_weighted_sum_pipeline = nil;
g_hc_split_weighted_sum_norm_pipeline = nil;
g_dsv4_hc_producer_pre_norm_pipeline = nil;
g_hc_weighted_sum_pipeline = nil;
g_output_hc_weights4_pipeline = nil;
g_hc_expand_pipeline = nil;
g_moe_mul_mv_id_iq2_xxs_pipeline = nil;
g_moe_mul_mv_id_iq2_xxs_pair_pipeline = nil;
g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline = nil;
g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline = nil;
g_moe_mul_mv_id_q2_k_pipeline = nil;
g_moe_mul_mv_id_q2_k_sum6_pipeline = nil;
g_moe_mul_mv_id_iq2_xxs_sum6_pipeline = nil;
g_moe_mul_mv_id_q4_k_pipeline = nil;
g_moe_mul_mv_id_q4_k_pair_pipeline = nil;
g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline = nil;
g_moe_mul_mv_id_q4_k_sum6_pipeline = nil;
g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline = nil;
g_moe_mul_mv_group_q4_k_sum6_pipeline = nil;
g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline = nil;
g_moe_mul_mv_group6_q4_k_sum6_pipeline = nil;
g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline = nil;
g_moe_mul_mv_group8_q4_k_sum6_pipeline = nil;
g_moe_mul_mv_group24_q4_k_id_pipeline = nil;
g_moe_mul_mv_group24_q4_k_sum6_pipeline = nil;
g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline = nil;
g_moe_mul_mv_slots6_q2_k_sum6_pipeline = nil;
g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline = nil;
g_moe_mul_mv_slots6_q4_k_sum6_pipeline = nil;
g_moe_mul_mv_id_mxfp4_pipeline = nil;
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline = nil;
g_moe_mul_mv_id_mxfp4_sum6_pipeline = nil;
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1 = nil;
g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1 = nil;
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple = nil;
g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple = nil;
g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 = nil;
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 = nil;
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 = nil;
g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 = nil;
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 = nil;
g_moe_mul_mv_id_mxfp4_pair_swiglu_tp_static_pipeline_nsg1 = nil;
g_moe_mul_mv_id_mxfp4_sum6_tp_full_rows_static_pipeline_nsg1 = nil;
g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline = nil;
g_moe_mul_mv_slots6_mxfp4_sum6_pipeline = nil;
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline = nil;
g_moe_mul_mv_addr_iq2_xxs_pipeline = nil;
g_moe_mul_mv_addr_q2_k_sum6_pipeline = nil;
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline = nil;
g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline = nil;
g_moe_stream_expert_cache_validate_pipeline = nil;
g_moe_q4_gather_slots6_pipeline = nil;
g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline = nil;
g_moe_mul_mv_table_q4_k_sum6_pipeline = nil;
g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline = nil;
g_moe_mul_mv_addr_q4_k_sum6_pipeline = nil;
g_moe_table_q4_pair_gate_encoder = nil;
g_moe_table_q4_pair_up_encoder = nil;
g_moe_table_q4_sum_down_encoder = nil;
g_rope_tail_batch_pipeline = nil;
g_rope_tail_inplace_pair_pipeline = nil;
g_rope_tail_inplace_pair_shared4_pipeline = nil;
g_rope_tail_inplace_pair_affine_pipeline = nil;
g_dsv4_fp8_kv_quantize_pipeline = nil;
g_dsv4_indexer_qat_pipeline = nil;
g_dsv4_kv_fp8_store_pipeline = nil;
g_dsv4_kv_rope_fp8_store_pipeline = nil;
g_dsv4_ratio4_shift_pipeline = nil;
g_dsv4_compressor_pack_ratio4_pipeline = nil;
g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline = nil;
g_dsv4_compressor_exact_softmax_product_pipeline = nil;
g_dsv4_compressor_exact_pool_ratio4_pipeline = nil;
g_dsv4_softmax_pool_ratio4_direct_pipeline = nil;
g_dsv4_softmax_pool_pipeline = nil;
g_soft_max_f32_pipeline = nil;
g_soft_max_f32_4_pipeline = nil;
g_argsort_f32_i32_desc_pipeline = nil;
g_argsort_merge_f32_i32_desc_pipeline = nil;
g_sum_rows_f32_f32_pipeline = nil;
g_dsv4_topk_mask_pipeline = nil;
g_dsv4_topk_mask_scatter_pipeline = nil;
g_dsv4_indexer_weighted_sum_pipeline = nil;
g_dsv4_indexer_score_one_direct_pipeline = nil;
g_dsv4_compressor_store_one_pipeline = nil;
g_dsv4_sort_i32_rows_asc_pipeline = nil;
g_dsv4_indexed_attention_heads8_pipeline = nil;
g_dsv4_indexed_attention_heads8_rb16_pipeline = nil;
g_dsv4_indexed_attention_heads16_dual_pipeline = nil;
g_dsv4_indexed_attention_heads8_split_pipeline = nil;
g_dsv4_indexed_attention_heads8_split_reduce_pipeline = nil;
g_dsv4_softplus_sqrt_pipeline = nil;
g_dsv4_router_finalize_one_pipeline = nil;
g_dsv4_router_finalize_one_simd_pipeline = nil;
g_dsv4_router_finalize_weights_one_simd_pipeline = nil;
g_dsv4_router_transform_finalize_weights_one_simd_pipeline = nil;
g_dsv4_router_select_visual_batch_pipeline = nil;
g_dsv4_hc_producer_last_completion = nil;
g_dsv4_hc_producer_last_mix_buffer = nil;
g_dsv4_hc_producer_last_mix_offset = 0;
[g_dsv4_completion_cache removeAllObjects];
g_dsv4_completion_cache = nil;
g_dsv4_router_weights_one_pipeline = nil;
g_glm_router_select_one_pipeline = nil;
g_glm_kv_lora_rms_norm_pipeline = nil;
g_glm_k_b_project_pipeline = nil;
g_glm_store_compact_kv_pipeline = nil;
g_glm_qkv_norm_store_compact_kv_pipeline = nil;
g_glm_store_indexer_k_pipeline = nil;
g_glm53_indexer_pool_update_pipeline = nil;
g_glm_build_kv_cache_pipeline = nil;
g_glm_build_kv_cache_decode_group4_pipeline = nil;
g_glm_build_kv_cache_flash_pipeline = nil;
g_glm_attention_full_pipeline = nil;
g_glm_fill_selected_range_pipeline = nil;
g_glm_fill_selected_range_batch_pipeline = nil;
g_glm53_expand_pool_selection_pipeline = nil;
g_glm_indexer_rope_tail_pipeline = nil;
g_glm_indexer_score_one_pipeline = nil;
g_glm_indexer_score_one_direct_pipeline = nil;
g_glm_indexer_scores_batch_pipeline = nil;
g_glm_indexer_scores_tiled_pipeline = nil;
g_glm_indexer_scores_tiled_f32_pipeline = nil;
g_glm_qk_lowrank_pipeline = nil;
g_glm_qk_lowrank_glm52_pipeline = nil;
g_glm_qk_lowrank_glm52_sg_pipeline = nil;
g_glm_qk_lowrank_batch_pipeline = nil;
g_glm_qk_lowrank_batch_glm52_t4_pipeline = nil;
g_glm_value_project_q8_0_pipeline = nil;
g_glm_value_project_q8_0_batch_heads_pipeline = nil;
g_glm_value_project_q8_0_batch_heads_mma_pipeline = nil;
g_glm_attention_indexed_decode_pipeline = nil;
g_glm_attention_indexed_decode_split_group8_partial_pipeline = nil;
g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline = nil;
g_glm_attention_indexed_decode_split_group8_reduce_pipeline = nil;
g_glm_attention_indexed_decode_split_group8_reduce16_pipeline = nil;
g_glm_attention_indexed_batch_pipeline = nil;
g_glm_attention_indexed_batch_group2_pipeline = nil;
g_glm_attention_indexed_batch_q2_group4_pipeline = nil;
g_glm_attention_indexed_batch_group8_pipeline = nil;
g_glm_attention_indexed_batch_lora_group8_vec_pipeline = nil;
g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline = nil;
g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline = nil;
g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline = nil;
g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline = nil;
g_glm_q4_k_pair_swiglu_f32_pipeline = nil;
g_glm_q4_k_pair_swiglu2_f32_pipeline = nil;
g_glm_q4_k_pair_swiglu4_f32_pipeline = nil;
g_glm_q4_k_pair_swiglu2_mapped_f32_pipeline = nil;
g_glm_q4_k_pair_swiglu2_mapped_row_f32_pipeline = nil;
g_glm_q2_k_pair_swiglu_f32_pipeline = nil;
g_glm_q2_k_addr_pair_swiglu2_f32_pipeline = nil;
g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline = nil;
g_glm_q4_k_addr_pair_swiglu_f32_pipeline = nil;
g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline = nil;
g_glm_q2_k_down_f32_pipeline = nil;
g_glm_q4_k_down_f32_pipeline = nil;
g_glm_q2_k_addr_down_f32_pipeline = nil;
g_glm_q4_k_addr_down_f32_pipeline = nil;
g_glm_q5_k_pair_swiglu_f32_pipeline = nil;
g_glm_q5_k_pair_swiglu_mapped_f32_pipeline = nil;
g_glm_q5_k_pair_swiglu_mapped_row_f32_pipeline = nil;
g_glm_q5_k_down_f32_pipeline = nil;
g_glm_q6_k_down_f32_pipeline = nil;
g_dsv4_router_weights_batch_pipeline = nil;
g_dsv4_hc_expand4_pipeline = nil;
g_flash_attn_mask_buffer = nil;
g_flash_attn_zero_mask_buffer = nil;
g_flash_attn_pad_buffer = nil;
g_flash_attn_tmp_buffer = nil;
g_flash_attn_blk_buffer = nil;
ds4_gpu_clear_zero_prefix_prefill_mask_cache();
g_flash_attn_ring_buffer = nil;
g_flash_attn_kv_buffer = nil;
g_glm_flash_attn_mask_buffer = nil;
g_compressor_pool_kv_buffer = nil;
g_compressor_pool_score_buffer = nil;
g_compressor_pool_score_cont_buffer = nil;
g_compressor_pool_softmax_buffer = nil;
g_compressor_pool_product_buffer = nil;
g_compressor_store_ape_buffer = nil;
g_compressor_store_score_buffer = nil;
g_embed_rows_buffer = nil;
g_router_selection_buffer = nil;
g_router_weight_sum_buffer = nil;
g_indexer_head_scores_buffer = nil;
g_indexer_topk_buffer = nil;
g_indexed_topk_buffer = nil;
g_stream_expert_validate_status_buffer = nil;
g_f16_round_scratch_buffer = nil;
g_q8_prefill_scratch_buffer = nil;
g_qwen4_conv_halo_buffer = nil;
g_qwen4_conv_halo_bytes = 0;
g_raw_store_round_buffer = nil;
g_moe_gate_scratch_buffer = nil;
g_moe_down_scratch_buffer = nil;
g_moe_id_map_buffer = nil;
g_moe_packed_rhs_buffer = nil;
g_moe_packed_rhs_bytes = 0;
g_moe_q4_gate_slots_buffer = nil;
g_moe_q4_up_slots_buffer = nil;
g_moe_q4_down_slots_buffer = nil;
g_attn_out_group_ids_buffer = nil;
g_model_fd = -1;
g_model_map_ptr = NULL;
g_model_map_size = 0;
g_model_mapped_offset = 0;
g_model_mapped_size = 0;
g_model_mapped_max_tensor_bytes = 0;
qwen4_nax_release_scratch();
qwen4_batch_release_scratch();
ds4_gpu_tensor_tracking_reset();
g_flash_attn_mask_bytes = 0;
g_flash_attn_zero_mask_bytes = 0;
g_flash_attn_pad_bytes = 0;
g_flash_attn_tmp_bytes = 0;
g_flash_attn_blk_bytes = 0;
g_flash_attn_ring_bytes = 0;
g_flash_attn_kv_bytes = 0;
g_glm_flash_attn_mask_bytes = 0;
g_glm_flash_attn_mask_valid = 0;
g_glm_flash_attn_mask_pos0 = 0;
g_glm_flash_attn_mask_tokens = 0;
g_glm_flash_attn_mask_cache_len = 0;
g_compressor_pool_kv_bytes = 0;
g_compressor_pool_score_bytes = 0;
g_compressor_pool_score_cont_bytes = 0;
g_compressor_pool_softmax_bytes = 0;
g_compressor_pool_product_bytes = 0;
g_compressor_store_ape_bytes = 0;
g_compressor_store_score_bytes = 0;
g_embed_rows_bytes = 0;
g_router_selection_bytes = 0;
g_router_weight_sum_bytes = 0;
g_indexer_head_scores_bytes = 0;
g_indexer_topk_bytes = 0;
g_indexed_topk_bytes = 0;
g_f16_round_scratch_bytes = 0;
g_q8_prefill_scratch_bytes = 0;
g_raw_store_round_bytes = 0;
g_moe_gate_scratch_bytes = 0;
g_moe_down_scratch_bytes = 0;
g_moe_id_map_bytes = 0;
g_moe_q4_gate_slots_bytes = 0;
g_moe_q4_up_slots_bytes = 0;
g_moe_q4_down_slots_bytes = 0;
g_attn_out_group_ids_bytes = 0;
g_model_wrap_count = 0;
g_model_wrap_bytes = 0;
g_model_wrap_max_bytes = 0;
g_model_buffer_cache_bytes = 0;
g_model_buffer_cache_evictions = 0;
g_model_buffer_cache_over_limit = 0;
ds4_gpu_model_residency_clear();
ds4_gpu_model_views_clear();
[g_pipeline_cache removeAllObjects];
g_pipeline_cache = nil;
[g_q4_expert_layer_residency_cache removeAllObjects];
g_q4_expert_layer_residency_cache = nil;
[g_q4_expert_table_cache removeAllObjects];
g_q4_expert_table_cache = nil;
[g_model_buffer_cache removeAllObjects];
g_model_buffer_cache = nil;
g_transient_buffers = nil;
g_pending_cbs = nil;
g_library = nil;
g_queue = nil;
g_device = nil;
g_initialized = 0;
}
}
static int ds4_gpu_encode_get_rows_f16(
id<MTLCommandBuffer> cb,
id<MTLBuffer> weight,
NSUInteger weight_offset,
id<MTLBuffer> tokens,
NSUInteger tokens_offset,
id<MTLBuffer> out,
NSUInteger out_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd) {
if (!cb || !weight || !tokens || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
return 0;
}
const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t);
const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
ds4_gpu_get_rows_args args = {
.ne00t = (int32_t)n_embd,
.ne00 = (int32_t)n_embd,
.nb01 = src_row_bytes,
.nb02 = (uint64_t)n_vocab * src_row_bytes,
.nb03 = (uint64_t)n_vocab * src_row_bytes,
.ne10 = (int32_t)n_tokens,
.nb10 = sizeof(int32_t),
.nb11 = token_bytes,
.nb12 = token_bytes,
.nb1 = dst_row_bytes,
.nb2 = (uint64_t)n_tokens * dst_row_bytes,
.nb3 = (uint64_t)n_tokens * dst_row_bytes,
};
NSUInteger nth = (NSUInteger)n_embd;
const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth == 0) nth = 1;
const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_get_rows_f16_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weight offset:weight_offset atIndex:1];
[enc setBuffer:tokens offset:tokens_offset atIndex:2];
[enc setBuffer:out offset:out_offset atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static uint64_t ds4_gpu_q8_0_row_bytes(uint32_t n_embd) {
return (((uint64_t)n_embd + 31u) / 32u) * 34u;
}
static int ds4_gpu_quant_row_bytes(
uint32_t type,
uint32_t n_embd,
uint64_t *row_bytes_out) {
if (!row_bytes_out || n_embd == 0) return 0;
switch (type) {
case DS4_METAL_TENSOR_Q8_0:
*row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 34u;
return 1;
case DS4_METAL_TENSOR_Q4_0:
*row_bytes_out = (((uint64_t)n_embd + 31u) / 32u) * 18u;
return 1;
case DS4_METAL_TENSOR_Q4_K:
if ((n_embd % 256u) != 0) return 0;
*row_bytes_out = ((uint64_t)n_embd / 256u) * 144u;
return 1;
default:
return 0;
}
}
static int ds4_gpu_q8_0_table_bytes(
uint32_t n_vocab,
uint32_t n_embd,
uint64_t *bytes_out) {
if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0;
const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd);
if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0;
*bytes_out = (uint64_t)n_vocab * row_bytes;
return 1;
}
static int ds4_gpu_quant_table_bytes(
uint32_t type,
uint32_t n_vocab,
uint32_t n_embd,
uint64_t *bytes_out) {
if (!bytes_out || n_vocab == 0 || n_embd == 0) return 0;
uint64_t row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(type, n_embd, &row_bytes)) return 0;
if (row_bytes != 0 && (uint64_t)n_vocab > UINT64_MAX / row_bytes) return 0;
*bytes_out = (uint64_t)n_vocab * row_bytes;
return 1;
}
static int ds4_gpu_encode_get_rows_q8_0(
id<MTLCommandBuffer> cb,
id<MTLBuffer> weight,
NSUInteger weight_offset,
id<MTLBuffer> tokens,
NSUInteger tokens_offset,
const int32_t *single_token,
id<MTLBuffer> out,
NSUInteger out_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd) {
if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
return 0;
}
if (!tokens && (!single_token || n_tokens != 1)) {
return 0;
}
ds4_gpu_get_rows_q8_0_args args = {
.n_embd = (int32_t)n_embd,
.n_vocab = (int32_t)n_vocab,
.n_tokens = (int32_t)n_tokens,
.src_row_bytes = ds4_gpu_q8_0_row_bytes(n_embd),
.dst_row_bytes = (uint64_t)n_embd * sizeof(float),
.token_stride = sizeof(int32_t),
};
NSUInteger nth = 32u;
const NSUInteger max_threads = g_get_rows_q8_0_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth == 0) nth = 1;
const NSUInteger nblocks = ((NSUInteger)n_embd + 31u) / 32u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_get_rows_q8_0_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weight offset:weight_offset atIndex:1];
if (tokens) {
[enc setBuffer:tokens offset:tokens_offset atIndex:2];
} else {
[enc setBytes:single_token length:sizeof(*single_token) atIndex:2];
}
[enc setBuffer:out offset:out_offset atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_get_rows_quant(
id<MTLCommandBuffer> cb,
id<MTLBuffer> weight,
NSUInteger weight_offset,
uint32_t weight_type,
id<MTLBuffer> tokens,
NSUInteger tokens_offset,
const int32_t *single_token,
id<MTLBuffer> out,
NSUInteger out_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd) {
if (weight_type == DS4_METAL_TENSOR_Q8_0) {
return ds4_gpu_encode_get_rows_q8_0(cb,
weight,
weight_offset,
tokens,
tokens_offset,
single_token,
out,
out_offset,
n_vocab,
n_tokens,
n_embd);
}
if (!cb || !weight || !out || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
return 0;
}
if (!tokens && (!single_token || n_tokens != 1)) {
return 0;
}
uint64_t src_row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(weight_type, n_embd, &src_row_bytes)) return 0;
ds4_gpu_get_rows_q8_0_args args = {
.n_embd = (int32_t)n_embd,
.n_vocab = (int32_t)n_vocab,
.n_tokens = (int32_t)n_tokens,
.src_row_bytes = src_row_bytes,
.dst_row_bytes = (uint64_t)n_embd * sizeof(float),
.token_stride = sizeof(int32_t),
};
id<MTLComputePipelineState> pipeline = nil;
NSUInteger block_width = 0;
if (weight_type == DS4_METAL_TENSOR_Q4_0) {
pipeline = g_get_rows_q4_0_pipeline;
block_width = 32u;
} else if (weight_type == DS4_METAL_TENSOR_Q4_K) {
pipeline = g_get_rows_q4_K_pipeline;
block_width = 256u;
}
if (!pipeline || block_width == 0) return 0;
NSUInteger nth = block_width;
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth == 0) nth = 1;
const NSUInteger nblocks = ((NSUInteger)n_embd + block_width - 1u) / block_width;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weight offset:weight_offset atIndex:1];
if (tokens) {
[enc setBuffer:tokens offset:tokens_offset atIndex:2];
} else {
[enc setBytes:single_token length:sizeof(*single_token) atIndex:2];
}
[enc setBuffer:out offset:out_offset atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(nblocks, n_tokens, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_repeat_hc_embedding(
id<MTLCommandBuffer> cb,
id<MTLBuffer> rows,
NSUInteger rows_offset,
id<MTLBuffer> out,
NSUInteger out_offset,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_hc) {
if (!cb || !rows || !out || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0;
const uint64_t embd_bytes = (uint64_t)n_embd * sizeof(float);
ds4_gpu_repeat_args args = {
.ne00 = (int32_t)n_embd,
.ne01 = 1,
.ne02 = (int32_t)n_tokens,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = embd_bytes,
.nb02 = embd_bytes,
.nb03 = (uint64_t)n_tokens * embd_bytes,
.ne0 = (int32_t)n_embd,
.ne1 = (int32_t)n_hc,
.ne2 = (int32_t)n_tokens,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = embd_bytes,
.nb2 = (uint64_t)n_hc * embd_bytes,
.nb3 = (uint64_t)n_tokens * n_hc * embd_bytes,
};
NSUInteger nth = (NSUInteger)n_embd;
const NSUInteger max_threads = g_repeat_f32_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth == 0) nth = 1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_repeat_f32_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:rows offset:rows_offset atIndex:1];
[enc setBuffer:out offset:out_offset atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(n_hc, n_tokens, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
int ds4_gpu_embed_token_q8_0_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t token,
uint32_t n_embd) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float);
if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal Q8_0 embedding received undersized output buffer\n");
return 0;
}
uint64_t weight_bytes = 0;
if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) ||
weight_offset > model_size ||
weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal Q8_0 embedding range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
uint32_t token_for_kernel = token;
id<MTLBuffer> wbuf = nil;
const bool exact_token_row =
getenv("DS4_METAL_DISABLE_TOKEN_EMBED_EXACT_VIEW") == NULL;
if (exact_token_row) {
const uint64_t row_bytes = ds4_gpu_q8_0_row_bytes(n_embd);
const uint64_t token_rel = (uint64_t)token * row_bytes;
if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) {
fprintf(stderr, "ds4: Metal Q8_0 embedding token row is outside the mapped table\n");
return 0;
}
wbuf = ds4_gpu_wrap_model_exact_range(model_map,
model_size,
weight_offset + token_rel,
row_bytes,
&inner_offset);
token_for_kernel = 0;
} else {
wbuf = ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
}
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const int32_t token_i32 = (int32_t)token_for_kernel;
if (!ds4_gpu_encode_get_rows_q8_0(cb,
wbuf,
(NSUInteger)inner_offset,
nil,
0,
&token_i32,
outbuf,
ds4_gpu_tensor_offset(out),
n_vocab,
1,
n_embd)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed token")) return 0;
}
return 1;
}
int ds4_gpu_embed_tokens_q8_0_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *tokens,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> tokbuf = ds4_gpu_tensor_buffer(tokens);
const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float);
const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
if (!outbuf || !tokbuf ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(tokens) < token_bytes) {
fprintf(stderr, "ds4: Metal Q8_0 batched embedding received undersized buffers\n");
return 0;
}
uint64_t weight_bytes = 0;
if (!ds4_gpu_q8_0_table_bytes(n_vocab, n_embd, &weight_bytes) ||
weight_offset > model_size ||
weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal Q8_0 batched embedding range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_get_rows_q8_0(cb,
wbuf,
(NSUInteger)inner_offset,
tokbuf,
ds4_gpu_tensor_offset(tokens),
NULL,
outbuf,
ds4_gpu_tensor_offset(out),
n_vocab,
n_tokens,
n_embd)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "q8_0 embed tokens")) return 0;
}
return 1;
}
int ds4_gpu_embed_token_quant_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_vocab,
uint32_t token,
uint32_t n_embd) {
if (weight_type == DS4_METAL_TENSOR_Q8_0) {
return ds4_gpu_embed_token_q8_0_tensor(out,
model_map,
model_size,
weight_offset,
n_vocab,
token,
n_embd);
}
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float);
if (!outbuf || ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal quant embedding received undersized output buffer\n");
return 0;
}
uint64_t weight_bytes = 0;
uint64_t row_bytes = 0;
if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) ||
!ds4_gpu_quant_row_bytes(weight_type, n_embd, &row_bytes) ||
weight_offset > model_size ||
weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal quant embedding range is outside the mapped model\n");
return 0;
}
const uint64_t token_rel = (uint64_t)token * row_bytes;
if (token_rel > weight_bytes || row_bytes > weight_bytes - token_rel) {
fprintf(stderr, "ds4: Metal quant embedding token row is outside the mapped table\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf =
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
weight_offset + token_rel,
row_bytes,
&inner_offset);
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const int32_t token_i32 = 0;
if (!ds4_gpu_encode_get_rows_quant(cb,
wbuf,
(NSUInteger)inner_offset,
weight_type,
nil,
0,
&token_i32,
outbuf,
ds4_gpu_tensor_offset(out),
n_vocab,
1,
n_embd)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed token")) return 0;
}
return 1;
}
int ds4_gpu_embed_tokens_quant_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *tokens,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd) {
if (weight_type == DS4_METAL_TENSOR_Q8_0) {
return ds4_gpu_embed_tokens_q8_0_tensor(out,
tokens,
model_map,
model_size,
weight_offset,
n_vocab,
n_tokens,
n_embd);
}
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> tokbuf = ds4_gpu_tensor_buffer(tokens);
const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * sizeof(float);
const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
if (!outbuf || !tokbuf ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(tokens) < token_bytes) {
fprintf(stderr, "ds4: Metal quant batched embedding received undersized buffers\n");
return 0;
}
uint64_t weight_bytes = 0;
if (!ds4_gpu_quant_table_bytes(weight_type, n_vocab, n_embd, &weight_bytes) ||
weight_offset > model_size ||
weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal quant batched embedding range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_get_rows_quant(cb,
wbuf,
(NSUInteger)inner_offset,
weight_type,
tokbuf,
ds4_gpu_tensor_offset(tokens),
NULL,
outbuf,
ds4_gpu_tensor_offset(out),
n_vocab,
n_tokens,
n_embd)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "quant embed tokens")) return 0;
}
return 1;
}
int ds4_gpu_embed_token_hc_tensor(
ds4_gpu_tensor *out_hc,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t token,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out_hc || !model_map || n_vocab == 0 || token >= n_vocab || n_embd == 0 || n_hc == 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
const uint64_t out_bytes = (uint64_t)n_embd * n_hc * sizeof(float);
if (!outbuf || ds4_gpu_tensor_bytes(out_hc) < out_bytes) {
fprintf(stderr, "ds4: Metal graph embedding received undersized HC output buffer\n");
return 0;
}
const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t);
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal graph embedding range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) return 0;
const NSUInteger row_bytes = (NSUInteger)n_embd * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer,
&g_embed_rows_bytes,
row_bytes,
"ds4_embed_rows")) {
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const int32_t token_i32 = (int32_t)token;
const uint64_t src_row_bytes = (uint64_t)n_embd * sizeof(uint16_t);
const uint64_t dst_row_bytes = (uint64_t)n_embd * sizeof(float);
ds4_gpu_get_rows_args args = {
.ne00t = (int32_t)n_embd,
.ne00 = (int32_t)n_embd,
.nb01 = src_row_bytes,
.nb02 = (uint64_t)n_vocab * src_row_bytes,
.nb03 = (uint64_t)n_vocab * src_row_bytes,
.ne10 = 1,
.nb10 = sizeof(int32_t),
.nb11 = sizeof(int32_t),
.nb12 = sizeof(int32_t),
.nb1 = dst_row_bytes,
.nb2 = dst_row_bytes,
.nb3 = dst_row_bytes,
};
NSUInteger nth = (NSUInteger)n_embd;
const NSUInteger max_threads = g_get_rows_f16_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth == 0) nth = 1;
const NSUInteger nw0 = ((NSUInteger)n_embd + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_get_rows_f16_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBytes:&token_i32 length:sizeof(token_i32) atIndex:2];
[enc setBuffer:g_embed_rows_buffer offset:0 atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(nw0, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_encode_repeat_hc_embedding(cb,
g_embed_rows_buffer,
0,
outbuf,
ds4_gpu_tensor_offset(out_hc),
1,
n_embd,
n_hc)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed token")) return 0;
}
return 1;
}
int ds4_gpu_embed_tokens_hc_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *tokens,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_vocab,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out_hc || !tokens || !model_map || n_vocab == 0 || n_tokens == 0 || n_embd == 0 || n_hc == 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
id<MTLBuffer> tokbuf = ds4_gpu_tensor_buffer(tokens);
const uint64_t out_bytes = (uint64_t)n_tokens * n_embd * n_hc * sizeof(float);
const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
if (!outbuf || !tokbuf ||
ds4_gpu_tensor_bytes(out_hc) < out_bytes ||
ds4_gpu_tensor_bytes(tokens) < token_bytes) {
fprintf(stderr, "ds4: Metal graph batched embedding received undersized buffers\n");
return 0;
}
const uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t);
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal graph batched embedding range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) return 0;
const NSUInteger rows_bytes = (NSUInteger)n_tokens * n_embd * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_embed_rows_buffer,
&g_embed_rows_bytes,
rows_bytes,
"ds4_embed_rows")) {
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_get_rows_f16(cb,
wbuf,
(NSUInteger)inner_offset,
tokbuf,
ds4_gpu_tensor_offset(tokens),
g_embed_rows_buffer,
0,
n_vocab,
n_tokens,
n_embd) ||
!ds4_gpu_encode_repeat_hc_embedding(cb,
g_embed_rows_buffer,
0,
outbuf,
ds4_gpu_tensor_offset(out_hc),
n_tokens,
n_embd,
n_hc)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph embed tokens")) return 0;
}
return 1;
}
int ds4_gpu_set_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) return 0;
if (map_offset > model_size || map_size == 0 || map_size > model_size - map_offset) return 0;
max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(map_size, max_tensor_bytes);
@autoreleasepool {
if (g_model_map_ptr == model_map &&
g_model_map_size == model_size &&
g_model_mapped_offset == map_offset &&
g_model_mapped_size == map_size &&
g_model_mapped_max_tensor_bytes == max_tensor_bytes) {
return 1;
}
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 &&
map_offset >= g_model_views[i].model_offset &&
map_offset + map_size <= g_model_views[i].model_offset + g_model_views[i].bytes) {
return 1;
}
}
ds4_gpu_model_residency_clear();
if (!ds4_gpu_map_model_views(model_map, model_size, map_offset, map_size, max_tensor_bytes)) {
ds4_gpu_model_residency_clear();
return 0;
}
g_model_map_ptr = model_map;
g_model_map_size = model_size;
g_model_mapped_offset = map_offset;
g_model_mapped_size = map_size;
g_model_mapped_max_tensor_bytes = max_tensor_bytes;
if (ds4_gpu_model_map_log_enabled()) {
fprintf(stderr,
"ds4: Metal mapped mmaped model as %u overlapping shared buffers\n",
g_model_view_count);
}
return 1;
}
}
static int ds4_gpu_model_views_cover_spans(
const void *model_map,
uint64_t model_size,
const uint64_t *offsets,
const uint64_t *sizes,
uint32_t count);
int ds4_gpu_set_model_map_spans(
const void *model_map,
uint64_t model_size,
const uint64_t *offsets,
const uint64_t *sizes,
uint32_t count,
uint64_t max_tensor_bytes) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!model_map || model_size == 0 || !offsets || !sizes || count == 0) return 0;
if (count == 1 && !g_ssd_streaming_mode) {
return ds4_gpu_set_model_map_range(model_map,
model_size,
offsets[0],
sizes[0],
max_tensor_bytes);
}
if (ds4_gpu_model_views_cover_spans(model_map, model_size, offsets, sizes, count)) {
return 1;
}
@autoreleasepool {
const double t0 = ds4_gpu_now_ms();
max_tensor_bytes = ds4_gpu_effective_model_max_tensor_bytes(model_size, max_tensor_bytes);
ds4_gpu_model_residency_clear();
/* Streaming changes the target's active layer, not the auxiliary
* vision/drafter models. Single spans need the same bounded lifetime
* as disjoint spans instead of accumulating every visited layer. */
if (g_ssd_streaming_mode) ds4_gpu_model_views_remove_map(model_map);
else ds4_gpu_model_views_clear();
uint64_t mapped_total = 0;
uint64_t first_offset = UINT64_MAX;
for (uint32_t i = 0; i < count; i++) {
if (offsets[i] > model_size || sizes[i] == 0 || sizes[i] > model_size - offsets[i]) {
fprintf(stderr, "ds4: Metal model span %u is outside the GGUF mapping\n", i);
ds4_gpu_model_residency_clear();
ds4_gpu_model_views_clear();
return 0;
}
if (offsets[i] < first_offset) first_offset = offsets[i];
uint64_t effective_max = max_tensor_bytes;
if (effective_max > sizes[i]) effective_max = sizes[i];
if (!ds4_gpu_add_model_view_range(model_map,
model_size,
offsets[i],
sizes[i],
effective_max,
true,
&mapped_total)) {
ds4_gpu_model_residency_clear();
ds4_gpu_model_views_clear();
return 0;
}
}
if (!ds4_gpu_finish_model_views(t0, mapped_total, first_offset)) {
ds4_gpu_model_residency_clear();
ds4_gpu_model_views_clear();
return 0;
}
g_model_map_ptr = model_map;
g_model_map_size = model_size;
g_model_mapped_offset = first_offset == UINT64_MAX ? 0 : first_offset;
g_model_mapped_size = mapped_total;
g_model_mapped_max_tensor_bytes = max_tensor_bytes;
if (ds4_gpu_model_map_log_enabled()) {
fprintf(stderr,
"ds4: Metal mapped mmaped model as %u disjoint shared buffers across %u tensor spans\n",
g_model_view_count,
count);
}
return 1;
}
}
int ds4_gpu_set_model_map(const void *model_map, uint64_t model_size) {
return ds4_gpu_set_model_map_range(model_map, model_size, 0, model_size, 0);
}
int ds4_gpu_set_model_fd(int fd) {
g_model_fd = fd;
return 1;
}
static int ds4_gpu_model_views_cover_range(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t size) {
if (!model_map || model_size == 0 || size == 0 ||
offset > model_size || size > model_size - offset) {
return 0;
}
const uint64_t end = offset + size;
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) {
continue;
}
const uint64_t view_start = g_model_views[i].model_offset;
const uint64_t view_end = view_start + g_model_views[i].bytes;
if (offset >= view_start && end <= view_end) return 1;
}
return 0;
}
static int ds4_gpu_model_views_cover_spans(
const void *model_map,
uint64_t model_size,
const uint64_t *offsets,
const uint64_t *sizes,
uint32_t count) {
if (!offsets || !sizes || count == 0) return 0;
for (uint32_t i = 0; i < count; i++) {
if (!ds4_gpu_model_views_cover_range(model_map,
model_size,
offsets[i],
sizes[i])) {
return 0;
}
}
return 1;
}
int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) {
(void)fd;
(void)model_map;
return 1;
}
static id<MTLBuffer> ds4_gpu_wrap_model_range(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t len,
uint64_t *inner_offset) {
(void)model_map;
if (model_size == 0 || offset > model_size || len > model_size - offset) {
fprintf(stderr, "ds4: Metal model range is outside the mapped model\n");
return nil;
}
const uint64_t end = offset + len;
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) {
continue;
}
const uint64_t view_start = g_model_views[i].model_offset;
const uint64_t view_end = view_start + g_model_views[i].bytes;
if (offset >= view_start && end <= view_end) {
*inner_offset = offset - view_start;
return g_model_views[i].buffer;
}
}
fprintf(stderr,
"ds4: Metal model range %.2f..%.2f GiB is not covered by mapped model views\n",
ds4_gpu_gib(offset),
ds4_gpu_gib(end));
return nil;
}
typedef enum {
DS4_GPU_EXACT_VIEW_CACHED,
DS4_GPU_EXACT_VIEW_TRANSIENT,
DS4_GPU_EXACT_VIEW_OWNED,
} ds4_gpu_exact_view_lifetime;
static id<MTLBuffer> ds4_gpu_wrap_model_exact_range_impl(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t len,
uint64_t *inner_offset,
ds4_gpu_exact_view_lifetime lifetime) {
const bool cache_view = lifetime == DS4_GPU_EXACT_VIEW_CACHED;
const bool transient_view = lifetime == DS4_GPU_EXACT_VIEW_TRANSIENT;
if (!model_map || !g_device ||
(cache_view && !g_model_buffer_cache) ||
(transient_view && !g_transient_buffers) ||
model_size == 0 || offset > model_size || len > model_size - offset) {
fprintf(stderr, "ds4: Metal exact model range is outside the mapped model\n");
return nil;
}
const uint64_t page = (uint64_t)getpagesize();
const uint64_t page_offset = offset & ~(page - 1);
const uint64_t leading = offset - page_offset;
if (len > UINT64_MAX - leading ||
leading + len > UINT64_MAX - (page - 1)) {
fprintf(stderr, "ds4: Metal exact model range overflows page alignment\n");
return nil;
}
uint64_t view_bytes = round_up_u64(leading + len, page);
if (view_bytes > model_size - page_offset) view_bytes = model_size - page_offset;
if (leading + len > view_bytes) {
fprintf(stderr, "ds4: Metal exact model range alignment exceeds mapped model\n");
return nil;
}
if (view_bytes > (uint64_t)[g_device maxBufferLength]) {
fprintf(stderr,
"ds4: Metal exact model range %.2f GiB exceeds maxBufferLength %.2f GiB\n",
ds4_gpu_gib(view_bytes),
ds4_gpu_gib((uint64_t)[g_device maxBufferLength]));
return nil;
}
NSString *key = nil;
id<MTLBuffer> buffer = nil;
if (cache_view) {
key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu",
model_map,
(unsigned long long)model_size,
(unsigned long long)page_offset,
(unsigned long long)view_bytes];
buffer = [g_model_buffer_cache objectForKey:key];
}
if (!buffer) {
const uintptr_t base = (uintptr_t)model_map;
buffer = [g_device newBufferWithBytesNoCopy:(void *)(base + page_offset)
length:(NSUInteger)view_bytes
options:ds4_gpu_model_resource_options()
deallocator:nil];
if (!buffer) {
fprintf(stderr,
"ds4: Metal could not wrap exact mmaped model range at %.2f GiB, size %.2f MiB\n",
ds4_gpu_gib(page_offset),
ds4_gpu_mib(view_bytes));
return nil;
}
if (cache_view) {
buffer.label = @"ds4_model_exact_view";
} else if (transient_view) {
buffer.label = @"ds4_model_exact_transient_view";
} else {
buffer.label = @"ds4_model_exact_owned_view";
}
if (cache_view) {
[g_model_buffer_cache setObject:buffer forKey:key];
ds4_gpu_model_buffer_cache_note_insert(view_bytes);
} else if (transient_view) {
[g_transient_buffers addObject:buffer];
}
}
if (inner_offset) *inner_offset = leading;
return buffer;
}
static id<MTLBuffer> ds4_gpu_wrap_model_exact_range(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t len,
uint64_t *inner_offset) {
return ds4_gpu_wrap_model_exact_range_impl(model_map,
model_size,
offset,
len,
inner_offset,
DS4_GPU_EXACT_VIEW_CACHED);
}
static id<MTLBuffer> ds4_gpu_wrap_model_exact_range_transient(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t len,
uint64_t *inner_offset) {
return ds4_gpu_wrap_model_exact_range_impl(model_map,
model_size,
offset,
len,
inner_offset,
DS4_GPU_EXACT_VIEW_TRANSIENT);
}
static id<MTLBuffer> ds4_gpu_wrap_model_exact_range_owned(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t len,
uint64_t *inner_offset) {
return ds4_gpu_wrap_model_exact_range_impl(model_map,
model_size,
offset,
len,
inner_offset,
DS4_GPU_EXACT_VIEW_OWNED);
}
uint32_t ds4_gpu_stream_expert_cache_configured_count(void) {
uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
if (budget > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {
budget = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES;
}
return budget;
}
uint32_t ds4_gpu_stream_expert_cache_current_count(void) {
return g_stream_expert_cache_entry_count;
}
uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size(
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
if (!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes)) {
return 0;
}
return ds4_gpu_stream_expert_cache_configured_budget();
}
static int ds4_gpu_stream_expert_cache_note_expert_size(
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
if (gate_expert_bytes == 0 || down_expert_bytes == 0) return 0;
if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) {
fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n");
return 0;
}
/*
* The cache is a single-size-class slab allocator: the expert byte size is
* frozen on first sight (or pre-seeded at startup from the model's slab
* class) and off-size layers are REJECTED rather than adopted. A rejected
* layer (mixed-precision boost: Q4_K experts among IQ2 layers) falls back
* to the mapped-model per-expert path; last-writer-wins here would instead
* poison the slab size class and deadlock slab reuse.
*/
const uint64_t bytes = gate_expert_bytes * 2ull + down_expert_bytes;
if (g_stream_expert_cache_expert_bytes == 0) {
g_stream_expert_cache_expert_bytes = bytes;
return 1;
}
return bytes == g_stream_expert_cache_expert_bytes;
}
static uint32_t ds4_gpu_stream_expert_cache_requested_budget(void) {
if (!g_ssd_streaming_mode) return 0;
if (g_stream_expert_cache_budget_override != 0) {
return g_stream_expert_cache_budget_override;
}
return 0;
}
static uint32_t ds4_gpu_stream_expert_cache_configured_budget(void) {
uint32_t budget = ds4_gpu_stream_expert_cache_requested_budget();
if (budget != 0 &&
g_stream_expert_cache_mlock_budget_cap != 0 &&
budget > g_stream_expert_cache_mlock_budget_cap) {
budget = g_stream_expert_cache_mlock_budget_cap;
}
return budget;
}
static uint32_t ds4_gpu_stream_expert_cache_effective_cap(
uint32_t layer,
uint32_t n_total_expert,
uint32_t n_selected) {
(void)layer;
if (ds4_gpu_stream_expert_cache_configured_budget() == 0) return 0;
/*
* The residency policy is global: every layer can use any expert slot it
* routes to, and global pruning decides which existing entry is least
* valuable. A per-layer cap made cache size depend on model depth rather
* than the actual byte budget.
*/
uint32_t cap = n_total_expert;
if (cap < n_selected) cap = n_selected;
if (cap > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
cap = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
}
return cap;
}
static int ds4_gpu_stream_expert_timing_summary_enabled(void) {
static int checked = 0;
static int enabled = 0;
if (!checked) {
enabled =
(getenv("DS4_METAL_STREAMING_EXPERT_TIMING_SUMMARY") != NULL ||
getenv("DS4_METAL_STREAMING_EXPERT_PROFILE_SUMMARY") != NULL) &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_TIMING_SUMMARY") == NULL;
checked = 1;
}
return enabled;
}
static uint32_t ds4_gpu_stream_expert_popcount(uint32_t mask) {
return (uint32_t)__builtin_popcount(mask);
}
static int ds4_gpu_stream_expert_split_worthwhile(
uint32_t resident_mask,
uint32_t missing_mask) {
if (resident_mask == 0 || missing_mask == 0) return 0;
/*
* The split path pays an extra command stage and a second routed-expert
* bind. It is worthwhile when several experts are missing and their SSD
* reads can be hidden by resident expert work. With one or two misses,
* especially in large caches, a single unsplit routed pass is faster.
*/
return ds4_gpu_stream_expert_popcount(missing_mask) >= 3u;
}
static void ds4_gpu_stream_expert_timing_note_selected(
double sync_ms,
double copy_ms,
double bind_ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_selected_calls++;
g_stream_expert_timing_selected_sync_ms += sync_ms;
g_stream_expert_timing_selected_copy_ms += copy_ms;
g_stream_expert_timing_selected_read_ms += sync_ms + copy_ms;
g_stream_expert_timing_selected_bind_ms += bind_ms;
}
static void ds4_gpu_stream_expert_timing_note_split(
uint32_t resident_mask,
uint32_t missing_mask,
double resident_ms,
double missing_ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_split_layers++;
g_stream_expert_timing_split_resident_experts +=
ds4_gpu_stream_expert_popcount(resident_mask);
g_stream_expert_timing_split_missing_experts +=
ds4_gpu_stream_expert_popcount(missing_mask);
g_stream_expert_timing_split_resident_ms += resident_ms;
g_stream_expert_timing_split_missing_ms += missing_ms;
}
static void ds4_gpu_stream_expert_timing_note_split_missing_detail(
double load_ms,
double slot_ms,
double prune_ms,
double addr_ms,
double wait_ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_split_missing_load_ms += load_ms;
g_stream_expert_timing_split_missing_slot_ms += slot_ms;
g_stream_expert_timing_split_missing_prune_ms += prune_ms;
g_stream_expert_timing_split_missing_addr_ms += addr_ms;
g_stream_expert_timing_split_missing_wait_ms += wait_ms;
}
static void ds4_gpu_stream_expert_timing_note_load_detail(
double prepare_ms,
double pread_ms,
double install_ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_load_calls++;
g_stream_expert_timing_load_prepare_ms += prepare_ms;
g_stream_expert_timing_load_pread_ms += pread_ms;
g_stream_expert_timing_load_install_ms += install_ms;
}
static void ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(double ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_prepare_batch_reuse_calls++;
g_stream_expert_timing_prepare_batch_reuse_ms += ms;
}
static void ds4_gpu_stream_expert_timing_note_prepare_buffer(double ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_prepare_buffer_calls++;
g_stream_expert_timing_prepare_buffer_ms += ms;
}
static void ds4_gpu_stream_expert_timing_note_prepare_task(
uint32_t experts,
double ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_prepare_task_experts += experts;
g_stream_expert_timing_prepare_task_ms += ms;
}
static void ds4_gpu_stream_expert_timing_note_reuse_scan(
uint64_t entries,
double ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_reuse_scan_calls++;
g_stream_expert_timing_reuse_scan_entries += entries;
g_stream_expert_timing_reuse_scan_ms += ms;
}
static void ds4_gpu_stream_expert_timing_note_reuse_clear(double ms) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
g_stream_expert_timing_reuse_clear_ms += ms;
}
static void ds4_gpu_stream_expert_timing_note_cache_class(
uint32_t resident_mask,
uint32_t missing_mask) {
if (!ds4_gpu_stream_expert_timing_summary_enabled()) return;
const uint32_t resident = ds4_gpu_stream_expert_popcount(resident_mask);
const uint32_t missing = ds4_gpu_stream_expert_popcount(missing_mask);
if (missing == 0) {
g_stream_expert_timing_cache_all_resident_layers++;
} else if (resident == 0) {
g_stream_expert_timing_cache_all_missing_layers++;
} else {
g_stream_expert_timing_cache_mixed_layers++;
}
g_stream_expert_timing_cache_resident_experts += resident;
g_stream_expert_timing_cache_missing_experts += missing;
}
static int ds4_gpu_stream_expert_readahead_enabled(void) {
return g_ssd_streaming_mode &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_READAHEAD") == NULL;
}
static void ds4_gpu_stream_expert_readahead_range(uint64_t offset, uint64_t len) {
if (!ds4_gpu_stream_expert_readahead_enabled() || g_model_fd < 0 || len == 0) {
return;
}
#if defined(F_RDADVISE)
const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
const double t0 = timing ? ds4_gpu_now_ms() : 0.0;
uint64_t pos = offset;
uint64_t rem = len;
while (rem > 0) {
const uint64_t chunk64 =
rem > (uint64_t)INT_MAX ? (uint64_t)INT_MAX : rem;
if (pos > (uint64_t)LLONG_MAX) break;
struct radvisory ra;
ra.ra_offset = (off_t)pos;
ra.ra_count = (int)chunk64;
(void)fcntl(g_model_fd, F_RDADVISE, &ra);
pos += chunk64;
rem -= chunk64;
}
if (timing) {
g_stream_expert_timing_readahead_calls++;
g_stream_expert_timing_readahead_bytes += len;
g_stream_expert_timing_readahead_ms += ds4_gpu_now_ms() - t0;
}
#else
(void)offset;
(void)len;
#endif
}
typedef struct {
uint64_t offset;
uint64_t len;
uint8_t *dst;
uint64_t read_bytes;
double ms;
int ok;
} ds4_gpu_stream_expert_pread_task;
typedef struct {
int active;
const void *model_map;
uint64_t model_size;
uint32_t layer;
uint32_t n_total_expert;
uint32_t n_selected;
uint32_t missing_mask;
uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint32_t n_loads;
uint32_t n_tasks;
uint64_t gate_expert_bytes;
uint64_t down_expert_bytes;
int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint64_t gate_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint64_t up_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint64_t down_abs_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
__strong id<MTLBuffer> gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
__strong id<MTLBuffer> up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
__strong id<MTLBuffer> down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u];
double start_ms;
double prepare_ms;
} ds4_gpu_stream_expert_pending_load;
static ds4_gpu_stream_expert_pending_load g_stream_expert_pending_load;
typedef struct {
int active;
const void *model_map;
uint64_t model_size;
uint32_t layer;
uint32_t n_total_expert;
uint32_t n_selected;
uint64_t gate_offset;
uint64_t up_offset;
uint64_t down_offset;
uint64_t gate_expert_bytes;
uint64_t down_expert_bytes;
int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
} ds4_gpu_glm_stream_selected_prefetch;
static ds4_gpu_glm_stream_selected_prefetch g_glm_stream_selected_prefetch;
typedef struct {
ds4_gpu_stream_expert_pread_task *tasks;
uint32_t n_tasks;
uint32_t worker_index;
uint32_t n_workers;
} ds4_gpu_stream_expert_pread_worker_args;
static void ds4_gpu_stream_expert_cache_note_pread(
uint32_t layer,
uint64_t bytes,
double ms) {
if (g_stream_expert_cache_pread_bytes > UINT64_MAX - bytes) {
g_stream_expert_cache_pread_bytes = UINT64_MAX;
} else {
g_stream_expert_cache_pread_bytes += bytes;
}
g_stream_expert_cache_pread_ms += ms;
if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) {
if (g_stream_expert_cache_layer_pread_bytes[layer] > UINT64_MAX - bytes) {
g_stream_expert_cache_layer_pread_bytes[layer] = UINT64_MAX;
} else {
g_stream_expert_cache_layer_pread_bytes[layer] += bytes;
}
g_stream_expert_cache_layer_pread_ms[layer] += ms;
}
}
static uint32_t ds4_gpu_stream_expert_pread_thread_limit(void) {
uint32_t threads = 9;
const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_THREADS");
if (env && env[0]) {
char *end = NULL;
unsigned long v = strtoul(env, &end, 10);
if (end != env && *end == '\0') {
threads = v > UINT32_MAX ? UINT32_MAX : (uint32_t)v;
}
}
if (threads == 0) threads = 1;
if (threads > 18) threads = 18;
return threads;
}
static uint32_t ds4_gpu_stream_expert_pread_thread_count(uint32_t n_tasks) {
if (n_tasks <= 1) return n_tasks;
uint32_t threads = ds4_gpu_stream_expert_pread_thread_limit();
if (threads > n_tasks) threads = n_tasks;
return threads;
}
static int ds4_gpu_stream_expert_pread_into(
uint64_t offset,
uint64_t len,
uint8_t *dst,
uint64_t *read_bytes,
double *ms_out) {
if (read_bytes) *read_bytes = 0;
if (ms_out) *ms_out = 0.0;
if (g_model_fd < 0 ||
!dst ||
len == 0 ||
offset > (uint64_t)LLONG_MAX ||
len > (uint64_t)LLONG_MAX - offset) {
return 0;
}
const double t0 = ds4_gpu_now_ms();
uint64_t pos = 0;
int ok = 1;
while (pos < len) {
const uint64_t rem = len - pos;
const size_t want = rem > (uint64_t)SSIZE_MAX ? (size_t)SSIZE_MAX : (size_t)rem;
ssize_t nread;
do {
nread = pread(g_model_fd, dst + pos, want, (off_t)(offset + pos));
} while (nread < 0 && errno == EINTR);
if (nread <= 0) {
ok = 0;
break;
}
pos += (uint64_t)nread;
}
const double dt = ds4_gpu_now_ms() - t0;
if (read_bytes) *read_bytes = pos;
if (ms_out) *ms_out = dt;
if (!ok || pos != len) {
fprintf(stderr,
"ds4: Metal streaming expert explicit pread failed offset=%.2f GiB len=%.2f MiB read=%.2f MiB\n",
ds4_gpu_gib(offset),
ds4_gpu_mib(len),
ds4_gpu_mib(pos));
return 0;
}
return 1;
}
static void *ds4_gpu_stream_expert_pread_worker(void *arg) {
ds4_gpu_stream_expert_pread_worker_args *wa =
(ds4_gpu_stream_expert_pread_worker_args *)arg;
for (uint32_t i = wa->worker_index; i < wa->n_tasks; i += wa->n_workers) {
ds4_gpu_stream_expert_pread_task *task = &wa->tasks[i];
task->ok = ds4_gpu_stream_expert_pread_into(task->offset,
task->len,
task->dst,
&task->read_bytes,
&task->ms);
}
return NULL;
}
static pthread_mutex_t g_stream_expert_pread_pool_mutex = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t g_stream_expert_pread_pool_start_cond = PTHREAD_COND_INITIALIZER;
static pthread_cond_t g_stream_expert_pread_pool_done_cond = PTHREAD_COND_INITIALIZER;
static pthread_t g_stream_expert_pread_pool_threads[18];
static uint32_t g_stream_expert_pread_pool_thread_count;
static uint32_t g_stream_expert_pread_pool_active_workers;
static uint32_t g_stream_expert_pread_pool_remaining_workers;
static uint32_t g_stream_expert_pread_pool_n_tasks;
static uint32_t g_stream_expert_pread_pool_next_task;
static uint64_t g_stream_expert_pread_pool_generation;
static ds4_gpu_stream_expert_pread_task *g_stream_expert_pread_pool_tasks;
static int g_stream_expert_pread_pool_initialized;
static int g_stream_expert_pread_pool_stopping;
static int ds4_gpu_stream_expert_pread_pool_enabled(void) {
const char *env = getenv("DS4_METAL_STREAMING_EXPERT_PREAD_POOL");
return !(env && strcmp(env, "0") == 0);
}
static void *ds4_gpu_stream_expert_pread_pool_worker(void *arg) {
const uint32_t worker_index = (uint32_t)(uintptr_t)arg;
uint64_t seen_generation = 0;
for (;;) {
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
while (!g_stream_expert_pread_pool_stopping &&
g_stream_expert_pread_pool_generation == seen_generation) {
pthread_cond_wait(&g_stream_expert_pread_pool_start_cond,
&g_stream_expert_pread_pool_mutex);
}
if (g_stream_expert_pread_pool_stopping) {
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
break;
}
seen_generation = g_stream_expert_pread_pool_generation;
if (worker_index >= g_stream_expert_pread_pool_active_workers) {
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
continue;
}
for (;;) {
const uint32_t task_index =
g_stream_expert_pread_pool_next_task++;
if (task_index >= g_stream_expert_pread_pool_n_tasks) break;
ds4_gpu_stream_expert_pread_task *task =
&g_stream_expert_pread_pool_tasks[task_index];
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
task->ok = ds4_gpu_stream_expert_pread_into(task->offset,
task->len,
task->dst,
&task->read_bytes,
&task->ms);
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
}
if (g_stream_expert_pread_pool_remaining_workers > 0 &&
--g_stream_expert_pread_pool_remaining_workers == 0) {
g_stream_expert_pread_pool_tasks = NULL;
g_stream_expert_pread_pool_n_tasks = 0;
g_stream_expert_pread_pool_active_workers = 0;
pthread_cond_signal(&g_stream_expert_pread_pool_done_cond);
}
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
}
return NULL;
}
static int ds4_gpu_stream_expert_pread_pool_init(uint32_t n_threads) {
if (g_stream_expert_pread_pool_initialized) return 1;
if (!ds4_gpu_stream_expert_pread_pool_enabled() || n_threads <= 1) return 0;
if (n_threads > 18) n_threads = 18;
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
g_stream_expert_pread_pool_thread_count = n_threads;
g_stream_expert_pread_pool_stopping = 0;
g_stream_expert_pread_pool_generation = 0;
g_stream_expert_pread_pool_tasks = NULL;
g_stream_expert_pread_pool_n_tasks = 0;
g_stream_expert_pread_pool_next_task = 0;
g_stream_expert_pread_pool_active_workers = 0;
g_stream_expert_pread_pool_remaining_workers = 0;
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
uint32_t started = 0;
for (uint32_t i = 0; i < n_threads; i++) {
const int rc = pthread_create(&g_stream_expert_pread_pool_threads[i],
NULL,
ds4_gpu_stream_expert_pread_pool_worker,
(void *)(uintptr_t)i);
if (rc != 0) {
fprintf(stderr,
"ds4: Metal streaming expert pread pool thread creation failed: %s\n",
strerror(rc));
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
g_stream_expert_pread_pool_stopping = 1;
g_stream_expert_pread_pool_generation++;
pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond);
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
for (uint32_t j = 0; j < started; j++) {
(void)pthread_join(g_stream_expert_pread_pool_threads[j], NULL);
}
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
g_stream_expert_pread_pool_thread_count = 0;
g_stream_expert_pread_pool_stopping = 0;
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
return 0;
}
started++;
}
g_stream_expert_pread_pool_initialized = 1;
return 1;
}
static int ds4_gpu_stream_expert_pread_pool_begin(
ds4_gpu_stream_expert_pread_task *tasks,
uint32_t n_tasks,
uint32_t n_workers) {
if (n_workers <= 1) return 0;
const uint32_t limit = ds4_gpu_stream_expert_pread_thread_limit();
if (!ds4_gpu_stream_expert_pread_pool_init(limit)) return 0;
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
if (!g_stream_expert_pread_pool_initialized ||
g_stream_expert_pread_pool_stopping ||
g_stream_expert_pread_pool_thread_count == 0) {
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
return 0;
}
if (n_workers > g_stream_expert_pread_pool_thread_count) {
n_workers = g_stream_expert_pread_pool_thread_count;
}
if (n_workers == 0) {
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
return 0;
}
if (g_stream_expert_pread_pool_remaining_workers != 0 ||
g_stream_expert_pread_pool_tasks != NULL) {
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
return 0;
}
g_stream_expert_pread_pool_tasks = tasks;
g_stream_expert_pread_pool_n_tasks = n_tasks;
g_stream_expert_pread_pool_next_task = 0;
g_stream_expert_pread_pool_active_workers = n_workers;
g_stream_expert_pread_pool_remaining_workers = n_workers;
g_stream_expert_pread_pool_generation++;
pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond);
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
return 1;
}
static int ds4_gpu_stream_expert_pread_pool_wait(void) {
if (!g_stream_expert_pread_pool_initialized) return 0;
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
while (g_stream_expert_pread_pool_remaining_workers != 0) {
pthread_cond_wait(&g_stream_expert_pread_pool_done_cond,
&g_stream_expert_pread_pool_mutex);
}
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
return 1;
}
static int ds4_gpu_stream_expert_pread_pool_dispatch(
ds4_gpu_stream_expert_pread_task *tasks,
uint32_t n_tasks,
uint32_t n_workers) {
if (!ds4_gpu_stream_expert_pread_pool_begin(tasks, n_tasks, n_workers)) {
return 0;
}
return ds4_gpu_stream_expert_pread_pool_wait();
}
static void ds4_gpu_stream_expert_pread_pool_shutdown(void) {
if (!g_stream_expert_pread_pool_initialized) return;
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
g_stream_expert_pread_pool_stopping = 1;
g_stream_expert_pread_pool_generation++;
pthread_cond_broadcast(&g_stream_expert_pread_pool_start_cond);
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
const uint32_t n_threads = g_stream_expert_pread_pool_thread_count;
for (uint32_t i = 0; i < n_threads; i++) {
(void)pthread_join(g_stream_expert_pread_pool_threads[i], NULL);
}
pthread_mutex_lock(&g_stream_expert_pread_pool_mutex);
g_stream_expert_pread_pool_thread_count = 0;
g_stream_expert_pread_pool_active_workers = 0;
g_stream_expert_pread_pool_remaining_workers = 0;
g_stream_expert_pread_pool_n_tasks = 0;
g_stream_expert_pread_pool_next_task = 0;
g_stream_expert_pread_pool_tasks = NULL;
g_stream_expert_pread_pool_initialized = 0;
g_stream_expert_pread_pool_stopping = 0;
pthread_mutex_unlock(&g_stream_expert_pread_pool_mutex);
}
static int ds4_gpu_stream_expert_pread_tasks(
ds4_gpu_stream_expert_pread_task *tasks,
uint32_t n_tasks,
uint64_t *total_bytes,
double *wall_ms) {
if (total_bytes) *total_bytes = 0;
if (wall_ms) *wall_ms = 0.0;
if (!tasks || n_tasks == 0) return 1;
const uint32_t n_workers =
ds4_gpu_stream_expert_pread_thread_count(n_tasks);
const double t0 = ds4_gpu_now_ms();
int ok = 1;
if (n_workers <= 1) {
ds4_gpu_stream_expert_pread_worker_args wa = {
.tasks = tasks,
.n_tasks = n_tasks,
.worker_index = 0,
.n_workers = 1,
};
(void)ds4_gpu_stream_expert_pread_worker(&wa);
} else if (!ds4_gpu_stream_expert_pread_pool_dispatch(tasks,
n_tasks,
n_workers)) {
pthread_t threads[18];
ds4_gpu_stream_expert_pread_worker_args args[18];
uint32_t started = 0;
for (uint32_t i = 0; i < n_workers; i++) {
args[i].tasks = tasks;
args[i].n_tasks = n_tasks;
args[i].worker_index = i;
args[i].n_workers = n_workers;
const int rc = pthread_create(&threads[i],
NULL,
ds4_gpu_stream_expert_pread_worker,
&args[i]);
if (rc != 0) {
fprintf(stderr,
"ds4: Metal streaming expert pread thread creation failed: %s\n",
strerror(rc));
ok = 0;
break;
}
started++;
}
for (uint32_t i = 0; i < started; i++) {
if (pthread_join(threads[i], NULL) != 0) ok = 0;
}
}
const double dt = ds4_gpu_now_ms() - t0;
uint64_t bytes = 0;
for (uint32_t i = 0; i < n_tasks; i++) {
if (!tasks[i].ok) ok = 0;
if (bytes > UINT64_MAX - tasks[i].read_bytes) {
bytes = UINT64_MAX;
} else {
bytes += tasks[i].read_bytes;
}
}
if (total_bytes) *total_bytes = bytes;
if (wall_ms) *wall_ms = dt;
return ok;
}
static void ds4_gpu_stream_expert_cache_warn_mlock_failure(
uint64_t failed_len,
int err) {
if (g_stream_expert_cache_mlock_warned) return;
g_stream_expert_cache_mlock_warned = 1;
const uint64_t gib = 1024ull * 1024ull * 1024ull;
const uint32_t budget = ds4_gpu_stream_expert_cache_requested_budget();
uint64_t requested = 0;
if (budget != 0 && g_stream_expert_cache_expert_bytes != 0) {
requested =
budget > UINT64_MAX / g_stream_expert_cache_expert_bytes ?
UINT64_MAX :
(uint64_t)budget * g_stream_expert_cache_expert_bytes;
}
uint64_t suggested_gib = g_stream_expert_cache_mlock_bytes / gib;
if (suggested_gib > 1) suggested_gib--;
fprintf(stderr,
"ds4: warning: streaming expert cache could not mlock all buffers\n");
if (requested != 0) {
fprintf(stderr,
"ds4: requested cache: %u experts / %.2f GiB\n",
budget,
ds4_gpu_gib(requested));
} else {
fprintf(stderr,
"ds4: requested cache: %u experts\n",
budget);
}
fprintf(stderr,
"ds4: locked so far: %.2f GiB\n",
ds4_gpu_gib(g_stream_expert_cache_mlock_bytes));
fprintf(stderr,
"ds4: failed buffer: %.2f MiB (%s)\n",
ds4_gpu_mib(failed_len),
err != 0 ? strerror(err) : "mlock unavailable");
fprintf(stderr,
"ds4: macOS may page unlocked expert buffers, causing poor or unstable speed\n");
if (g_stream_expert_cache_mlock_budget_cap != 0 &&
g_stream_expert_cache_expert_bytes != 0) {
const uint64_t capped_bytes =
g_stream_expert_cache_mlock_budget_cap >
UINT64_MAX / g_stream_expert_cache_expert_bytes ?
UINT64_MAX :
(uint64_t)g_stream_expert_cache_mlock_budget_cap *
g_stream_expert_cache_expert_bytes;
fprintf(stderr,
"ds4: using locked cache cap: %u experts / %.2f GiB\n",
g_stream_expert_cache_mlock_budget_cap,
ds4_gpu_gib(capped_bytes));
}
if (suggested_gib != 0) {
fprintf(stderr,
"ds4: try: --ssd-streaming-cache-experts %" PRIu64 "GB\n",
suggested_gib);
} else {
fprintf(stderr,
"ds4: try a smaller --ssd-streaming-cache-experts NGB budget\n");
}
}
static uint32_t ds4_gpu_stream_expert_slab_locked_slot_count(void) {
uint32_t count = 0;
uint32_t total = g_stream_expert_cache_slab_total_slots;
if (total > DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {
total = DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES;
}
for (uint32_t slot = 0; slot < total; slot++) {
if (g_stream_expert_cache_slab_slot_locked[slot]) count++;
}
return count;
}
static void ds4_gpu_stream_expert_cache_cap_budget_to_locked(void) {
uint32_t cap = g_stream_expert_cache_entry_count;
const uint32_t locked_slots = ds4_gpu_stream_expert_slab_locked_slot_count();
if (locked_slots != 0 && locked_slots < cap) cap = locked_slots;
const uint64_t gib = 1024ull * 1024ull * 1024ull;
uint64_t safe_gib = g_stream_expert_cache_mlock_bytes / gib;
if (safe_gib > 1) safe_gib--;
if (safe_gib != 0 && g_stream_expert_cache_expert_bytes != 0) {
uint64_t safe_bytes =
safe_gib > UINT64_MAX / gib ? UINT64_MAX : safe_gib * gib;
uint64_t safe_cap64 = safe_bytes / g_stream_expert_cache_expert_bytes;
if (safe_cap64 > UINT32_MAX) safe_cap64 = UINT32_MAX;
if (safe_cap64 != 0 && safe_cap64 < cap) {
cap = (uint32_t)safe_cap64;
}
}
if (cap == 0) return;
if (g_stream_expert_cache_mlock_budget_cap == 0 ||
cap < g_stream_expert_cache_mlock_budget_cap) {
g_stream_expert_cache_mlock_budget_cap = cap;
}
}
static id<MTLBuffer> ds4_gpu_stream_expert_alloc_buffer(
uint64_t len,
NSString *label) {
if (!g_device ||
len == 0 ||
len > (uint64_t)NSUIntegerMax) {
return nil;
}
id<MTLBuffer> buffer = [g_device newBufferWithLength:(NSUInteger)len
options:MTLResourceStorageModeShared];
if (!buffer) {
fprintf(stderr,
"ds4: Metal streaming expert explicit buffer allocation failed (%.2f MiB)\n",
ds4_gpu_mib(len));
return nil;
}
buffer.label = label;
g_stream_expert_cache_buffer_allocs++;
if (g_ssd_streaming_mode) {
void *ptr = [buffer contents];
const NSUInteger n = [buffer length];
const double t0 = ds4_gpu_now_ms();
if (ptr && n != 0 && mlock(ptr, (size_t)n) == 0) {
const double dt = ds4_gpu_now_ms() - t0;
g_stream_expert_cache_mlock_ms += dt;
if (g_stream_expert_cache_mlock_bytes > UINT64_MAX - (uint64_t)n) {
g_stream_expert_cache_mlock_bytes = UINT64_MAX;
} else {
g_stream_expert_cache_mlock_bytes += (uint64_t)n;
}
} else {
const double dt = ds4_gpu_now_ms() - t0;
g_stream_expert_cache_mlock_ms += dt;
g_stream_expert_cache_mlock_failures++;
if (g_stream_expert_cache_mlock_fail_bytes > UINT64_MAX - (uint64_t)n) {
g_stream_expert_cache_mlock_fail_bytes = UINT64_MAX;
} else {
g_stream_expert_cache_mlock_fail_bytes += (uint64_t)n;
}
const int err = ptr && n != 0 ? errno : 0;
ds4_gpu_stream_expert_cache_cap_budget_to_locked();
ds4_gpu_stream_expert_cache_warn_mlock_failure((uint64_t)n, err);
if (getenv("DS4_METAL_STREAMING_EXPERT_BUFFER_MLOCK_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal streaming expert buffer mlock failed len=%.2f MiB: %s\n",
ds4_gpu_mib((uint64_t)n),
err != 0 ? strerror(err) : "mlock unavailable");
}
return nil;
}
}
return buffer;
}
static void ds4_gpu_stream_expert_unlock_explicit_buffer(id<MTLBuffer> buffer) {
if (!g_ssd_streaming_mode || !buffer) return;
void *ptr = [buffer contents];
const NSUInteger n = [buffer length];
if (!ptr || n == 0 || munlock(ptr, (size_t)n) != 0) return;
if (g_stream_expert_cache_mlock_bytes >= (uint64_t)n) {
g_stream_expert_cache_mlock_bytes -= (uint64_t)n;
} else {
g_stream_expert_cache_mlock_bytes = 0;
}
}
static int ds4_gpu_stream_expert_combined_buffer_enabled(void) {
return g_ssd_streaming_mode &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_COMBINED_BUFFER") == NULL;
}
static int ds4_gpu_stream_expert_slab_enabled(void) {
return ds4_gpu_stream_expert_combined_buffer_enabled() &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_SLABS") == NULL;
}
/*
* Large PRO caches otherwise create thousands of small shared Metal buffers.
* Slabs keep the buffer object set small while locking pages only for slots
* that actually hold a streamed expert.
*/
static uint64_t ds4_gpu_stream_expert_slab_target_bytes(void) {
const uint64_t mib = 1024ull * 1024ull;
uint64_t target = 4096ull * mib;
const char *env = getenv("DS4_METAL_STREAMING_EXPERT_SLAB_MB");
if (env && env[0]) {
char *end = NULL;
unsigned long long v = strtoull(env, &end, 10);
if (end != env && *end == '\0' && v != 0) {
target = v > UINT64_MAX / mib ? UINT64_MAX : (uint64_t)v * mib;
}
}
return target;
}
static id<MTLBuffer> ds4_gpu_stream_expert_alloc_slab_buffer(
uint64_t len,
NSString *label) {
if (!g_device ||
len == 0 ||
len > (uint64_t)NSUIntegerMax) {
return nil;
}
id<MTLBuffer> buffer = [g_device newBufferWithLength:(NSUInteger)len
options:MTLResourceStorageModeShared];
if (!buffer) {
fprintf(stderr,
"ds4: Metal streaming expert slab allocation failed (%.2f MiB)\n",
ds4_gpu_mib(len));
return nil;
}
buffer.label = label;
g_stream_expert_cache_buffer_allocs++;
return buffer;
}
static int ds4_gpu_stream_expert_slab_slot_range(
uint32_t slot,
uint32_t *slab_index,
uint64_t *slot_base) {
for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) {
const uint32_t start = g_stream_expert_cache_slab_start_slot[i];
const uint32_t count = g_stream_expert_cache_slab_slot_count[i];
if (slot < start || slot >= start + count) continue;
if (slab_index) *slab_index = i;
if (slot_base) {
*slot_base =
(uint64_t)(slot - start) * g_stream_expert_cache_slab_slot_bytes;
}
return 1;
}
return 0;
}
static int ds4_gpu_stream_expert_slab_slot_for_buffer(
id<MTLBuffer> buffer,
NSUInteger gate_inner,
uint32_t *slot_out) {
if (!buffer || g_stream_expert_cache_slab_slot_bytes == 0 || !slot_out) {
return 0;
}
for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) {
if (g_stream_expert_cache_slabs[i] != buffer) continue;
const uint64_t inner = (uint64_t)gate_inner;
const uint64_t slot_bytes = g_stream_expert_cache_slab_slot_bytes;
if (slot_bytes == 0 || inner % slot_bytes != 0) return 0;
const uint64_t local_slot = inner / slot_bytes;
if (local_slot >= g_stream_expert_cache_slab_slot_count[i]) return 0;
const uint64_t slot =
(uint64_t)g_stream_expert_cache_slab_start_slot[i] + local_slot;
if (slot > UINT32_MAX) return 0;
*slot_out = (uint32_t)slot;
return 1;
}
return 0;
}
static void ds4_gpu_stream_expert_slab_push_free_slot(uint32_t slot) {
if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES ||
g_stream_expert_cache_free_slot_count >=
DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES) {
return;
}
g_stream_expert_cache_free_slots[g_stream_expert_cache_free_slot_count++] =
slot;
}
static int ds4_gpu_stream_expert_slab_lock_slot(uint32_t slot) {
if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES ||
g_stream_expert_cache_slab_slot_locked[slot]) {
return 1;
}
uint32_t slab = UINT32_MAX;
uint64_t base = 0;
if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) ||
slab >= g_stream_expert_cache_slab_count ||
!g_stream_expert_cache_slabs[slab] ||
g_stream_expert_cache_slab_slot_bytes == 0 ||
base > (uint64_t)NSUIntegerMax) {
return 0;
}
void *contents = [g_stream_expert_cache_slabs[slab] contents];
if (!contents) return 0;
const double t0 = ds4_gpu_now_ms();
void *ptr = (uint8_t *)contents + (NSUInteger)base;
const size_t n = (size_t)g_stream_expert_cache_slab_slot_bytes;
if (mlock(ptr, n) == 0) {
const double dt = ds4_gpu_now_ms() - t0;
g_stream_expert_cache_mlock_ms += dt;
if (g_stream_expert_cache_mlock_bytes >
UINT64_MAX - g_stream_expert_cache_slab_slot_bytes) {
g_stream_expert_cache_mlock_bytes = UINT64_MAX;
} else {
g_stream_expert_cache_mlock_bytes +=
g_stream_expert_cache_slab_slot_bytes;
}
g_stream_expert_cache_slab_slot_locked[slot] = 1;
return 1;
}
const double dt = ds4_gpu_now_ms() - t0;
g_stream_expert_cache_mlock_ms += dt;
g_stream_expert_cache_mlock_failures++;
if (g_stream_expert_cache_mlock_fail_bytes >
UINT64_MAX - g_stream_expert_cache_slab_slot_bytes) {
g_stream_expert_cache_mlock_fail_bytes = UINT64_MAX;
} else {
g_stream_expert_cache_mlock_fail_bytes +=
g_stream_expert_cache_slab_slot_bytes;
}
ds4_gpu_stream_expert_cache_cap_budget_to_locked();
ds4_gpu_stream_expert_cache_warn_mlock_failure(
g_stream_expert_cache_slab_slot_bytes,
errno);
return 0;
}
static int ds4_gpu_stream_expert_slab_unlock_slot(uint32_t slot) {
if (slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES ||
!g_stream_expert_cache_slab_slot_locked[slot]) {
return 1;
}
uint32_t slab = UINT32_MAX;
uint64_t base = 0;
if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) ||
slab >= g_stream_expert_cache_slab_count ||
!g_stream_expert_cache_slabs[slab] ||
g_stream_expert_cache_slab_slot_bytes == 0 ||
base > (uint64_t)NSUIntegerMax) {
return 0;
}
void *contents = [g_stream_expert_cache_slabs[slab] contents];
if (!contents) return 0;
void *ptr = (uint8_t *)contents + (NSUInteger)base;
const size_t n = (size_t)g_stream_expert_cache_slab_slot_bytes;
if (munlock(ptr, n) != 0) return 0;
g_stream_expert_cache_slab_slot_locked[slot] = 0;
if (g_stream_expert_cache_mlock_bytes >= g_stream_expert_cache_slab_slot_bytes) {
g_stream_expert_cache_mlock_bytes -= g_stream_expert_cache_slab_slot_bytes;
} else {
g_stream_expert_cache_mlock_bytes = 0;
}
return 1;
}
static int ds4_gpu_stream_expert_slab_slot_buffers(
uint32_t slot,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
__strong id<MTLBuffer> *gate_buf,
__strong id<MTLBuffer> *up_buf,
__strong id<MTLBuffer> *down_buf,
NSUInteger *gate_inner,
NSUInteger *up_inner,
NSUInteger *down_inner) {
uint32_t slab = UINT32_MAX;
uint64_t base = 0;
if (!ds4_gpu_stream_expert_slab_slot_range(slot, &slab, &base) ||
slab >= g_stream_expert_cache_slab_count ||
!g_stream_expert_cache_slabs[slab]) {
return 0;
}
if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull ||
base > UINT64_MAX - (gate_expert_bytes * 2ull + down_expert_bytes) ||
base + gate_expert_bytes * 2ull + down_expert_bytes >
(uint64_t)NSUIntegerMax) {
return 0;
}
id<MTLBuffer> b = g_stream_expert_cache_slabs[slab];
if (!ds4_gpu_stream_expert_slab_lock_slot(slot)) return 0;
*gate_buf = b;
*up_buf = b;
*down_buf = b;
*gate_inner = (NSUInteger)base;
*up_inner = (NSUInteger)(base + gate_expert_bytes);
*down_inner = (NSUInteger)(base + gate_expert_bytes * 2ull);
return 1;
}
static int ds4_gpu_stream_expert_alloc_slab_slot(
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
__strong id<MTLBuffer> *gate_buf,
__strong id<MTLBuffer> *up_buf,
__strong id<MTLBuffer> *down_buf,
NSUInteger *gate_inner,
NSUInteger *up_inner,
NSUInteger *down_inner) {
if (!ds4_gpu_stream_expert_slab_enabled() ||
!gate_buf || !up_buf || !down_buf ||
!gate_inner || !up_inner || !down_inner ||
gate_expert_bytes == 0 || down_expert_bytes == 0 ||
gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) {
return 0;
}
uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes;
if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0;
const uint64_t page = (uint64_t)getpagesize();
if (page != 0) {
slot_bytes = round_up_u64(slot_bytes, page);
if (slot_bytes == 0 || slot_bytes > (uint64_t)NSUIntegerMax) return 0;
}
if (g_stream_expert_cache_slab_slot_bytes != 0 &&
g_stream_expert_cache_slab_slot_bytes != slot_bytes) {
return 0;
}
g_stream_expert_cache_slab_slot_bytes = slot_bytes;
if (g_stream_expert_cache_free_slot_count != 0) {
const uint32_t slot =
g_stream_expert_cache_free_slots[--g_stream_expert_cache_free_slot_count];
return ds4_gpu_stream_expert_slab_slot_buffers(slot,
gate_expert_bytes,
down_expert_bytes,
gate_buf,
up_buf,
down_buf,
gate_inner,
up_inner,
down_inner);
}
uint32_t slab = g_stream_expert_cache_slab_count;
if (slab != 0 &&
g_stream_expert_cache_slab_slots_used[slab - 1] <
g_stream_expert_cache_slab_slot_count[slab - 1]) {
slab--;
} else {
if (g_stream_expert_cache_slab_count >=
DS4_METAL_STREAM_EXPERT_CACHE_MAX_SLABS) {
return 0;
}
const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
if (budget != 0 && g_stream_expert_cache_slab_total_slots >= budget) {
return 0;
}
uint64_t target = ds4_gpu_stream_expert_slab_target_bytes();
uint64_t slots64 = target / slot_bytes;
if (slots64 == 0) slots64 = 1;
if (slots64 > UINT32_MAX) slots64 = UINT32_MAX;
uint32_t slots = (uint32_t)slots64;
if (budget != 0) {
const uint32_t remaining =
budget - g_stream_expert_cache_slab_total_slots;
if (slots > remaining) slots = remaining;
}
if (slots == 0) return 0;
id<MTLBuffer> slab_buffer = nil;
while (slots != 0) {
if ((uint64_t)slots <= UINT64_MAX / slot_bytes &&
(uint64_t)slots * slot_bytes <= (uint64_t)NSUIntegerMax) {
slab_buffer =
ds4_gpu_stream_expert_alloc_slab_buffer(
(uint64_t)slots * slot_bytes,
@"ds4_stream_expert_slab");
if (slab_buffer) break;
}
slots /= 2u;
}
if (!slab_buffer || slots == 0) return 0;
slab = g_stream_expert_cache_slab_count++;
g_stream_expert_cache_slabs[slab] = slab_buffer;
g_stream_expert_cache_slab_start_slot[slab] =
g_stream_expert_cache_slab_total_slots;
g_stream_expert_cache_slab_slot_count[slab] = slots;
g_stream_expert_cache_slab_slots_used[slab] = 0;
g_stream_expert_cache_slab_total_slots += slots;
}
const uint32_t local_slot = g_stream_expert_cache_slab_slots_used[slab]++;
const uint32_t slot =
g_stream_expert_cache_slab_start_slot[slab] + local_slot;
return ds4_gpu_stream_expert_slab_slot_buffers(slot,
gate_expert_bytes,
down_expert_bytes,
gate_buf,
up_buf,
down_buf,
gate_inner,
up_inner,
down_inner);
}
static uint64_t ds4_gpu_stream_expert_buffer_object_count(
id<MTLBuffer> gate,
id<MTLBuffer> up,
id<MTLBuffer> down) {
if (!gate || !up || !down) return 0;
if (gate == up && gate == down) return 1;
if (gate == up || gate == down || up == down) return 2;
return 3;
}
static int ds4_gpu_stream_expert_evict_dontneed_enabled(void) {
return g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_EVICT_DONTNEED") != NULL &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EVICT_DONTNEED") == NULL;
}
static void ds4_gpu_stream_expert_evict_dontneed_range(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t len) {
if (!ds4_gpu_stream_expert_evict_dontneed_enabled() ||
!model_map ||
model_size == 0 ||
offset > model_size ||
len == 0 ||
len > model_size - offset) {
return;
}
#if defined(POSIX_MADV_DONTNEED)
const uint64_t page = (uint64_t)getpagesize();
const uint64_t page_offset = offset & ~(page - 1);
const uint64_t leading = offset - page_offset;
if (len > UINT64_MAX - leading ||
leading + len > UINT64_MAX - (page - 1)) {
return;
}
uint64_t advise_bytes = round_up_u64(leading + len, page);
if (advise_bytes > model_size - page_offset) {
advise_bytes = model_size - page_offset;
}
if (advise_bytes == 0 || advise_bytes > (uint64_t)SIZE_MAX) return;
const uintptr_t base = (uintptr_t)model_map;
if (page_offset > (uint64_t)(UINTPTR_MAX - base)) return;
void *addr = (void *)(base + (uintptr_t)page_offset);
const int rc = posix_madvise(addr, (size_t)advise_bytes, POSIX_MADV_DONTNEED);
if (rc == 0) {
if (g_stream_expert_cache_evict_advise_bytes > UINT64_MAX - advise_bytes) {
g_stream_expert_cache_evict_advise_bytes = UINT64_MAX;
} else {
g_stream_expert_cache_evict_advise_bytes += advise_bytes;
}
} else if (getenv("DS4_METAL_STREAMING_EXPERT_EVICT_DONTNEED_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal streaming expert evict DONTNEED failed offset=%.2f GiB len=%.2f MiB: %s\n",
ds4_gpu_gib(offset),
ds4_gpu_mib(len),
strerror(rc));
}
#else
(void)model_map;
(void)model_size;
(void)offset;
(void)len;
#endif
}
static int ds4_gpu_stream_expert_split_requested(void) {
return g_ssd_streaming_mode;
}
static uint32_t ds4_gpu_stream_expert_split_min_decode_tokens(void) {
return 4;
}
static uint32_t ds4_gpu_stream_expert_split_min_cached(void) {
uint32_t min_cached = 1024;
const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
if (budget != 0 && budget < min_cached * 2u) {
min_cached = budget / 2u;
}
return min_cached;
}
static int ds4_gpu_stream_expert_split_ready(void) {
if (!ds4_gpu_stream_expert_split_requested()) return 0;
if (g_stream_expert_cache_decode_tokens <
ds4_gpu_stream_expert_split_min_decode_tokens()) {
return 0;
}
return g_stream_expert_cache_entry_count >=
ds4_gpu_stream_expert_split_min_cached();
}
static void ds4_gpu_stream_expert_cache_decay_route_hotness(void) {
for (uint32_t layer = 0;
layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
layer++) {
for (uint32_t expert = 0;
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
expert++) {
g_stream_expert_cache_route_hotness[layer][expert] >>= 1;
}
}
}
void ds4_gpu_stream_expert_cache_reset_route_hotness(void) {
memset(g_stream_expert_cache_route_hotness,
0,
sizeof(g_stream_expert_cache_route_hotness));
g_stream_expert_cache_hotness_decay_token =
g_stream_expert_cache_decode_tokens;
}
static void ds4_gpu_stream_expert_cache_maybe_decay_route_hotness(void) {
if (g_stream_expert_cache_decode_tokens == 0) return;
if (g_stream_expert_cache_hotness_decay_token == 0) {
g_stream_expert_cache_hotness_decay_token =
g_stream_expert_cache_decode_tokens;
return;
}
while (g_stream_expert_cache_decode_tokens -
g_stream_expert_cache_hotness_decay_token >=
DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS) {
ds4_gpu_stream_expert_cache_decay_route_hotness();
g_stream_expert_cache_hotness_decay_token +=
DS4_METAL_STREAM_EXPERT_HOTNESS_DECAY_TOKENS;
}
}
static void ds4_gpu_stream_expert_cache_note_route_hotness(
uint32_t layer,
uint32_t expert,
uint32_t amount) {
if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
amount == 0) {
return;
}
uint32_t *hotness = &g_stream_expert_cache_route_hotness[layer][expert];
if (*hotness > UINT32_MAX - amount) {
*hotness = UINT32_MAX;
} else {
*hotness += amount;
}
}
static void ds4_gpu_stream_expert_cache_note_selected_hotness(
uint32_t layer,
const int32_t *selected_ids,
uint32_t n_selected) {
if (!selected_ids ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
n_selected == 0) {
return;
}
ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
for (uint32_t i = 0; i < n_selected; i++) {
if (selected_ids[i] < 0 ||
selected_ids[i] >=
(int32_t)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
continue;
}
ds4_gpu_stream_expert_cache_note_route_hotness(
layer,
(uint32_t)selected_ids[i],
1);
}
}
static void ds4_gpu_stream_expert_cache_note_frequency_hotness(
uint32_t layer,
const uint32_t *frequency,
uint32_t n_total_expert) {
if (!frequency || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) {
return;
}
if (n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
n_total_expert = DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
}
ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
for (uint32_t expert = 0; expert < n_total_expert; expert++) {
ds4_gpu_stream_expert_cache_note_route_hotness(layer,
expert,
frequency[expert]);
}
}
static void ds4_gpu_stream_expert_cache_note_token(uint32_t layer_index) {
if (!g_ssd_streaming_mode ||
g_stream_expert_cache_decode_tokens == UINT64_MAX) {
return;
}
/* GLM starts with dense layers; a pipeline slice may start even later.
* Typed and generic routed kernels must advance the same aging clock. */
if (layer_index < g_stream_expert_cache_first_decode_layer)
g_stream_expert_cache_first_decode_layer = layer_index;
if (layer_index != g_stream_expert_cache_first_decode_layer) return;
g_stream_expert_cache_decode_tokens++;
ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
}
static int ds4_gpu_stream_compact_addr_requested(void) {
return g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_STREAMING_COMPACT_ADDR") != NULL &&
getenv("DS4_METAL_DISABLE_STREAMING_COMPACT_ADDR") == NULL &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL;
}
static int ds4_gpu_stream_expert_addr_table_requested(void) {
return g_ssd_streaming_mode &&
(getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL ||
getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL ||
getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL ||
g_stream_prefill_batch_selected_addr_building ||
g_glm_stream_expert_addr_table_building ||
(getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL &&
getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") == NULL) ||
ds4_gpu_stream_expert_split_requested()) &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL;
}
static int ds4_gpu_stream_expert_addr_table_kernel_requested(void) {
return g_ssd_streaming_mode &&
(getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL ||
getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL ||
getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL ||
ds4_gpu_stream_expert_split_ready()) &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL;
}
static int ds4_gpu_stream_expert_masked_addr_requested(void) {
return g_ssd_streaming_mode &&
(getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_MASKED_ADDR") != NULL ||
ds4_gpu_stream_expert_split_ready()) &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_MASKED_ADDR") == NULL;
}
static int ds4_gpu_stream_expert_hit_validator_requested(void) {
return g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_STREAMING_EXPERT_HIT_VALIDATOR") != NULL &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_HIT_VALIDATOR") == NULL;
}
static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_max(
uint32_t n_total_expert) {
const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MAX");
if (env && env[0]) {
char *end = NULL;
const long v = strtol(env, &end, 10);
if (end != env) {
if (v <= 0) return 0;
if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX;
return (uint32_t)v;
}
}
if (n_total_expert == 384) return 800u;
if (n_total_expert == 256) return 760u;
return 0;
}
static uint32_t ds4_gpu_stream_prefill_batch_selected_addr_auto_min(
uint32_t n_total_expert) {
const char *env = getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_MIN");
if (env && env[0]) {
char *end = NULL;
const long v = strtol(env, &end, 10);
if (end != env) {
if (v <= 0) return 0;
if ((unsigned long)v > (unsigned long)UINT32_MAX) return UINT32_MAX;
return (uint32_t)v;
}
}
if (n_total_expert == 384 || n_total_expert == 256) return 2u;
return 0;
}
static int ds4_gpu_stream_prefill_batch_selected_addr_enabled(
uint32_t n_tokens,
uint32_t n_total_expert,
uint32_t n_expert,
uint32_t gate_type,
uint32_t down_type) {
if (!g_ssd_streaming_mode ||
n_tokens <= 1 ||
n_total_expert == 0 ||
n_expert != 6 ||
gate_type != DS4_METAL_TENSOR_IQ2_XXS ||
down_type != DS4_METAL_TENSOR_Q2_K ||
g_quality_mode ||
getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL ||
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") != NULL ||
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL ||
getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") != NULL) {
return 0;
}
/* All unique experts for one layer must fit simultaneously because the
* address-table kernels consume them in one dispatch. Once the global
* cache fills, preparation reuses entries owned by other layers. */
if (ds4_gpu_stream_expert_cache_configured_count() < n_total_expert) {
return 0;
}
if (getenv("DS4_METAL_ENABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") != NULL) {
return 1;
}
const uint32_t max_tokens =
ds4_gpu_stream_prefill_batch_selected_addr_auto_max(n_total_expert);
const uint32_t min_tokens =
ds4_gpu_stream_prefill_batch_selected_addr_auto_min(n_total_expert);
return max_tokens != 0 && n_tokens >= min_tokens && n_tokens <= max_tokens;
}
static int ds4_gpu_glm_streaming_prefill_full_layer_active(void) {
return g_glm_streaming_prefill_full_layer_runtime ||
getenv("DS4_METAL_GLM_STREAMING_PREFILL_FULL_LAYER") != NULL;
}
static int ds4_gpu_stream_full_expert_addr_table_requested(void) {
return g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") != NULL &&
getenv("DS4_METAL_DISABLE_STREAMING_FULL_EXPERT_ADDR_TABLE") == NULL;
}
static uint64_t ds4_gpu_buffer_address(id<MTLBuffer> buffer, NSUInteger inner) {
if (!buffer) return 0;
#if TARGET_OS_OSX
if (@available(macOS 13.0, *)) {
return (uint64_t)[buffer gpuAddress] + (uint64_t)inner;
}
#endif
return 0;
}
static int ds4_gpu_stream_compact_addr_ensure_buffers(uint32_t layer) {
if (!g_device || layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return 0;
const NSUInteger addr_bytes = 6u * sizeof(uint64_t);
const NSUInteger ids_bytes = 6u * sizeof(int32_t);
for (uint32_t i = 0; i < 4; i++) {
id<MTLBuffer> current = nil;
NSUInteger bytes = addr_bytes;
NSString *label = @"ds4_stream_compact_gate_addresses";
switch (i) {
case 0:
current = g_stream_compact_gate_addr_buffers[layer];
label = @"ds4_stream_compact_gate_addresses";
break;
case 1:
current = g_stream_compact_up_addr_buffers[layer];
label = @"ds4_stream_compact_up_addresses";
break;
case 2:
current = g_stream_compact_down_addr_buffers[layer];
label = @"ds4_stream_compact_down_addresses";
break;
default:
current = g_stream_compact_selected_buffers[layer];
bytes = ids_bytes;
label = @"ds4_stream_compact_selected_ids";
break;
}
if (current) continue;
id<MTLBuffer> b = [g_device newBufferWithLength:bytes
options:MTLResourceStorageModeShared];
if (!b) {
fprintf(stderr, "ds4: Metal streaming compact address buffer allocation failed\n");
return 0;
}
b.label = label;
memset([b contents], 0, bytes);
switch (i) {
case 0:
g_stream_compact_gate_addr_buffers[layer] = b;
break;
case 1:
g_stream_compact_up_addr_buffers[layer] = b;
break;
case 2:
g_stream_compact_down_addr_buffers[layer] = b;
break;
default:
g_stream_compact_selected_buffers[layer] = b;
break;
}
}
return 1;
}
static int ds4_gpu_stream_compact_addr_prepare(
uint32_t layer,
ds4_gpu_stream_expert_cache_entry * const entries[6],
uint32_t n_entries,
id<MTLBuffer> *gate_addrs,
id<MTLBuffer> *up_addrs,
id<MTLBuffer> *down_addrs,
id<MTLBuffer> *selected_ids) {
if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
!entries || !gate_addrs || !up_addrs || !down_addrs || !selected_ids ||
n_entries == 0 || n_entries > 6) {
return 0;
}
if (!ds4_gpu_stream_compact_addr_ensure_buffers(layer)) return 0;
uint64_t gate_values[6] = {0, 0, 0, 0, 0, 0};
uint64_t up_values[6] = {0, 0, 0, 0, 0, 0};
uint64_t down_values[6] = {0, 0, 0, 0, 0, 0};
int32_t slot_ids[6] = {0, 1, 2, 3, 4, 5};
for (uint32_t i = 0; i < n_entries; i++) {
ds4_gpu_stream_expert_cache_entry *e = entries[i];
if (!e || !e->gate_buffer || !e->up_buffer || !e->down_buffer) {
return 0;
}
gate_values[i] = ds4_gpu_buffer_address(e->gate_buffer, e->gate_inner);
up_values[i] = ds4_gpu_buffer_address(e->up_buffer, e->up_inner);
down_values[i] = ds4_gpu_buffer_address(e->down_buffer, e->down_inner);
if (gate_values[i] == 0 || up_values[i] == 0 || down_values[i] == 0) {
fprintf(stderr, "ds4: Metal streaming compact address path requires GPU addresses\n");
return 0;
}
}
const NSUInteger addr_bytes = 6u * sizeof(uint64_t);
const NSUInteger ids_bytes = 6u * sizeof(int32_t);
id<MTLBuffer> gb = g_stream_compact_gate_addr_buffers[layer];
id<MTLBuffer> ub = g_stream_compact_up_addr_buffers[layer];
id<MTLBuffer> db = g_stream_compact_down_addr_buffers[layer];
id<MTLBuffer> ib = g_stream_compact_selected_buffers[layer];
memcpy([gb contents], gate_values, addr_bytes);
memcpy([ub contents], up_values, addr_bytes);
memcpy([db contents], down_values, addr_bytes);
memcpy([ib contents], slot_ids, ids_bytes);
*gate_addrs = gb;
*up_addrs = ub;
*down_addrs = db;
*selected_ids = ib;
return 1;
}
static int ds4_gpu_stream_selected_ids_prepare(
uint32_t layer,
const int32_t *selected_ids,
uint32_t n_selected,
id<MTLBuffer> *selected_buf,
NSUInteger *selected_off) {
if (!g_device ||
!selected_ids ||
!selected_buf ||
!selected_off ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
n_selected == 0 ||
n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED) {
return 0;
}
const NSUInteger bytes =
(NSUInteger)DS4_METAL_MAX_ROUTED_EXPERT_USED * sizeof(int32_t);
id<MTLBuffer> b = g_stream_selected_id_buffers[layer];
if (!b) {
b = [g_device newBufferWithLength:bytes
options:MTLResourceStorageModeShared];
if (!b) {
fprintf(stderr, "ds4: Metal streaming selected-id buffer allocation failed\n");
return 0;
}
b.label = @"ds4_stream_selected_ids";
g_stream_selected_id_buffers[layer] = b;
}
int32_t ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = {0};
memcpy(ids, selected_ids, (size_t)n_selected * sizeof(ids[0]));
memcpy([b contents], ids, bytes);
*selected_buf = b;
*selected_off = 0;
return 1;
}
static int ds4_gpu_stream_expert_cache_ensure_addr_buffers(uint32_t layer) {
if (!g_device ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) {
return 0;
}
const NSUInteger bytes =
(NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t);
id<MTLBuffer> buffers[3] = {
g_stream_expert_cache_gate_addr_buffers[layer],
g_stream_expert_cache_up_addr_buffers[layer],
g_stream_expert_cache_down_addr_buffers[layer],
};
for (uint32_t i = 0; i < 3; i++) {
if (buffers[i]) continue;
id<MTLBuffer> b = [g_device newBufferWithLength:bytes
options:MTLResourceStorageModeShared];
if (!b) {
fprintf(stderr, "ds4: Metal streaming expert address table allocation failed\n");
return 0;
}
b.label =
i == 0 ? @"ds4_stream_expert_gate_addresses" :
(i == 1 ? @"ds4_stream_expert_up_addresses" :
@"ds4_stream_expert_down_addresses");
memset([b contents], 0, bytes);
if (i == 0) {
g_stream_expert_cache_gate_addr_buffers[layer] = b;
} else if (i == 1) {
g_stream_expert_cache_up_addr_buffers[layer] = b;
} else {
g_stream_expert_cache_down_addr_buffers[layer] = b;
}
}
return 1;
}
static void ds4_gpu_stream_expert_cache_zero_addr_slot(uint32_t layer, uint32_t expert) {
if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
return;
}
id<MTLBuffer> buffers[3] = {
g_stream_expert_cache_gate_addr_buffers[layer],
g_stream_expert_cache_up_addr_buffers[layer],
g_stream_expert_cache_down_addr_buffers[layer],
};
const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t);
for (uint32_t i = 0; i < 3; i++) {
if (!buffers[i]) continue;
uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off);
*addr = 0;
}
}
static int ds4_gpu_stream_expert_cache_set_addr_slot_raw(
uint32_t layer,
uint32_t expert,
id<MTLBuffer> gate_buffer,
NSUInteger gate_inner,
id<MTLBuffer> up_buffer,
NSUInteger up_inner,
id<MTLBuffer> down_buffer,
NSUInteger down_inner) {
if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0;
const uint64_t values[3] = {
ds4_gpu_buffer_address(gate_buffer, gate_inner),
ds4_gpu_buffer_address(up_buffer, up_inner),
ds4_gpu_buffer_address(down_buffer, down_inner),
};
if (values[0] == 0 || values[1] == 0 || values[2] == 0) {
fprintf(stderr, "ds4: Metal streaming expert address table requires GPU addresses\n");
return 0;
}
id<MTLBuffer> buffers[3] = {
g_stream_expert_cache_gate_addr_buffers[layer],
g_stream_expert_cache_up_addr_buffers[layer],
g_stream_expert_cache_down_addr_buffers[layer],
};
const NSUInteger off = (NSUInteger)expert * sizeof(uint64_t);
for (uint32_t i = 0; i < 3; i++) {
uint64_t *addr = (uint64_t *)((uint8_t *)[buffers[i] contents] + off);
*addr = values[i];
}
return 1;
}
static int ds4_gpu_stream_expert_cache_set_addr_slot(
uint32_t layer,
uint32_t expert,
id<MTLBuffer> gate_buffer,
NSUInteger gate_inner,
id<MTLBuffer> up_buffer,
NSUInteger up_inner,
id<MTLBuffer> down_buffer,
NSUInteger down_inner) {
if (!ds4_gpu_stream_expert_addr_table_requested()) return 1;
return ds4_gpu_stream_expert_cache_set_addr_slot_raw(layer,
expert,
gate_buffer,
gate_inner,
up_buffer,
up_inner,
down_buffer,
down_inner);
}
static int ds4_gpu_stream_expert_cache_addr_buffers(
uint32_t layer,
id<MTLBuffer> *gate,
id<MTLBuffer> *up,
id<MTLBuffer> *down) {
if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) return 0;
if (gate) *gate = g_stream_expert_cache_gate_addr_buffers[layer];
if (up) *up = g_stream_expert_cache_up_addr_buffers[layer];
if (down) *down = g_stream_expert_cache_down_addr_buffers[layer];
return g_stream_expert_cache_gate_addr_buffers[layer] &&
g_stream_expert_cache_up_addr_buffers[layer] &&
g_stream_expert_cache_down_addr_buffers[layer];
}
static void ds4_gpu_stream_full_expert_addr_clear_layer(uint32_t layer) {
if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return;
ds4_gpu_stream_expert_cache_entry *e = &g_stream_full_expert_addr_entry[layer];
e->gate_buffer = nil;
e->up_buffer = nil;
e->down_buffer = nil;
e->model_map = NULL;
e->model_size = 0;
e->gate_abs_offset = 0;
e->up_abs_offset = 0;
e->down_abs_offset = 0;
e->gate_expert_bytes = 0;
e->down_expert_bytes = 0;
e->logical_bytes = 0;
e->last_used = 0;
e->use_count = 0;
e->gate_inner = 0;
e->up_inner = 0;
e->down_inner = 0;
e->slab_slot = 0;
e->valid = 0;
e->slab_backed = 0;
}
static int ds4_gpu_stream_full_expert_addr_table_prepare(
const void *model_map,
uint64_t model_size,
uint32_t layer,
uint32_t n_total_expert,
uint64_t gate_abs_offset,
uint64_t up_abs_offset,
uint64_t down_abs_offset,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
id<MTLBuffer> *gate_addrs,
id<MTLBuffer> *up_addrs,
id<MTLBuffer> *down_addrs,
ds4_gpu_stream_expert_cache_entry **entry_out) {
if (!ds4_gpu_stream_full_expert_addr_table_requested()) return 0;
if (!model_map || model_size == 0 ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
n_total_expert == 0 ||
n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
gate_expert_bytes == 0 ||
down_expert_bytes == 0 ||
n_total_expert > UINT64_MAX / gate_expert_bytes ||
n_total_expert > UINT64_MAX / down_expert_bytes) {
return 0;
}
const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes;
const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes;
if (gate_abs_offset > model_size ||
up_abs_offset > model_size ||
down_abs_offset > model_size ||
gate_tensor_bytes > model_size - gate_abs_offset ||
gate_tensor_bytes > model_size - up_abs_offset ||
down_tensor_bytes > model_size - down_abs_offset) {
return 0;
}
ds4_gpu_stream_expert_cache_entry *entry =
&g_stream_full_expert_addr_entry[layer];
if (!entry->valid ||
entry->model_map != model_map ||
entry->model_size != model_size ||
entry->gate_abs_offset != gate_abs_offset ||
entry->up_abs_offset != up_abs_offset ||
entry->down_abs_offset != down_abs_offset ||
entry->gate_expert_bytes != gate_expert_bytes ||
entry->down_expert_bytes != down_expert_bytes ||
!entry->gate_buffer ||
!entry->up_buffer ||
!entry->down_buffer) {
ds4_gpu_stream_full_expert_addr_clear_layer(layer);
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
uint64_t down_inner = 0;
id<MTLBuffer> gate_buf =
ds4_gpu_wrap_model_exact_range_owned(model_map,
model_size,
gate_abs_offset,
gate_tensor_bytes,
&gate_inner);
id<MTLBuffer> up_buf =
ds4_gpu_wrap_model_exact_range_owned(model_map,
model_size,
up_abs_offset,
gate_tensor_bytes,
&up_inner);
id<MTLBuffer> down_buf =
ds4_gpu_wrap_model_exact_range_owned(model_map,
model_size,
down_abs_offset,
down_tensor_bytes,
&down_inner);
if (!gate_buf || !up_buf || !down_buf) return 0;
entry->gate_buffer = gate_buf;
entry->up_buffer = up_buf;
entry->down_buffer = down_buf;
entry->model_map = model_map;
entry->model_size = model_size;
entry->gate_abs_offset = gate_abs_offset;
entry->up_abs_offset = up_abs_offset;
entry->down_abs_offset = down_abs_offset;
entry->gate_expert_bytes = gate_expert_bytes;
entry->down_expert_bytes = down_expert_bytes;
entry->logical_bytes = gate_tensor_bytes * 2ull + down_tensor_bytes;
entry->gate_inner = (NSUInteger)gate_inner;
entry->up_inner = (NSUInteger)up_inner;
entry->down_inner = (NSUInteger)down_inner;
entry->valid = 1;
if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) {
ds4_gpu_stream_full_expert_addr_clear_layer(layer);
return 0;
}
const uint64_t gate_base = ds4_gpu_buffer_address(entry->gate_buffer,
entry->gate_inner);
const uint64_t up_base = ds4_gpu_buffer_address(entry->up_buffer,
entry->up_inner);
const uint64_t down_base = ds4_gpu_buffer_address(entry->down_buffer,
entry->down_inner);
if (gate_base == 0 || up_base == 0 || down_base == 0) {
fprintf(stderr, "ds4: Metal full streaming expert address table requires GPU addresses\n");
ds4_gpu_stream_full_expert_addr_clear_layer(layer);
return 0;
}
id<MTLBuffer> buffers[3] = {
g_stream_expert_cache_gate_addr_buffers[layer],
g_stream_expert_cache_up_addr_buffers[layer],
g_stream_expert_cache_down_addr_buffers[layer],
};
uint64_t *gate_addr = (uint64_t *)[buffers[0] contents];
uint64_t *up_addr = (uint64_t *)[buffers[1] contents];
uint64_t *down_addr = (uint64_t *)[buffers[2] contents];
for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) {
if (expert < n_total_expert) {
gate_addr[expert] = gate_base + (uint64_t)expert * gate_expert_bytes;
up_addr[expert] = up_base + (uint64_t)expert * gate_expert_bytes;
down_addr[expert] = down_base + (uint64_t)expert * down_expert_bytes;
} else {
gate_addr[expert] = 0;
up_addr[expert] = 0;
down_addr[expert] = 0;
}
}
}
if (gate_addrs) *gate_addrs = g_stream_expert_cache_gate_addr_buffers[layer];
if (up_addrs) *up_addrs = g_stream_expert_cache_up_addr_buffers[layer];
if (down_addrs) *down_addrs = g_stream_expert_cache_down_addr_buffers[layer];
if (entry_out) *entry_out = entry;
return entry->valid &&
g_stream_expert_cache_gate_addr_buffers[layer] &&
g_stream_expert_cache_up_addr_buffers[layer] &&
g_stream_expert_cache_down_addr_buffers[layer];
}
static id<MTLBuffer> ds4_gpu_stream_expert_validate_status_buffer(void) {
if (!g_device) return nil;
if (g_stream_expert_validate_status_buffer) {
return g_stream_expert_validate_status_buffer;
}
const NSUInteger bytes =
(NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t);
id<MTLBuffer> b = [g_device newBufferWithLength:bytes
options:MTLResourceStorageModeShared];
if (!b) {
fprintf(stderr, "ds4: Metal streaming expert validator allocation failed\n");
return nil;
}
b.label = @"ds4_stream_expert_validate_status";
memset([b contents], 0, bytes);
g_stream_expert_validate_status_buffer = b;
return b;
}
static int ds4_gpu_encode_stream_expert_cache_validate(
id<MTLCommandBuffer> cb,
const ds4_gpu_stream_expert_validate_args *args,
id<MTLBuffer> selected,
NSUInteger selected_off,
id<MTLBuffer> gate_addrs,
id<MTLBuffer> up_addrs,
id<MTLBuffer> down_addrs,
id<MTLBuffer> status) {
if (!cb || !args || !selected || !gate_addrs || !up_addrs || !down_addrs ||
!status || !g_moe_stream_expert_cache_validate_pipeline ||
args->n_total_expert == 0 || args->n_total_expert > 384 ||
args->n_expert == 0 || args->n_expert > 6) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_moe_stream_expert_cache_validate_pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:selected offset:selected_off atIndex:1];
[enc setBuffer:gate_addrs offset:0 atIndex:2];
[enc setBuffer:up_addrs offset:0 atIndex:3];
[enc setBuffer:down_addrs offset:0 atIndex:4];
[enc setBuffer:status offset:0 atIndex:5];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_stream_expert_cache_validate_selected(
const ds4_gpu_tensor *selected,
id<MTLBuffer> gate_addrs,
id<MTLBuffer> up_addrs,
id<MTLBuffer> down_addrs,
uint32_t n_total_expert,
uint32_t n_expert,
int32_t selected_ids[6],
uint32_t *all_cached,
uint32_t *miss_mask,
uint32_t *invalid_mask) {
if (!selected || !selected_ids || !all_cached || !miss_mask || !invalid_mask) {
return 0;
}
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> status = ds4_gpu_stream_expert_validate_status_buffer();
if (!selectedbuf || !status) return 0;
const NSUInteger status_bytes =
(NSUInteger)DS4_METAL_STREAM_EXPERT_VALIDATE_WORDS * sizeof(uint32_t);
memset([status contents], 0, status_bytes);
ds4_gpu_stream_expert_validate_args args = {
.n_total_expert = n_total_expert,
.n_expert = n_expert,
};
const int had_batch = g_batch_cb != nil;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_stream_expert_cache_validate(cb,
&args,
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_addrs,
up_addrs,
down_addrs,
status)) {
return 0;
}
if (had_batch) {
if (ds4_gpu_end_commands() == 0) return 0;
} else if (!ds4_gpu_finish_command_buffer(cb, owned,
"streaming expert cache validator")) {
return 0;
}
const uint32_t *words = (const uint32_t *)[status contents];
*all_cached = words[0];
*miss_mask = words[1];
*invalid_mask = words[2];
for (uint32_t i = 0; i < 6; i++) {
selected_ids[i] = (int32_t)words[4 + i];
}
if (had_batch && ds4_gpu_begin_commands() == 0) return 0;
return 1;
}
static void ds4_gpu_stream_expert_cache_clear_entry_internal(
uint32_t layer,
uint32_t expert,
int count_eviction,
int recycle_slab_slot,
ds4_gpu_stream_expert_reusable_buffers *reuse) {
if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
return;
}
if (reuse) {
reuse->gate_buffer = nil;
reuse->up_buffer = nil;
reuse->down_buffer = nil;
reuse->gate_inner = 0;
reuse->up_inner = 0;
reuse->down_inner = 0;
}
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (!e->valid) return;
if (ds4_gpu_stream_expert_cache_entry_inflight(e)) {
return;
}
const uint64_t bytes = e->logical_bytes;
ds4_gpu_stream_expert_evict_dontneed_range(e->model_map,
e->model_size,
e->gate_abs_offset,
e->gate_expert_bytes);
ds4_gpu_stream_expert_evict_dontneed_range(e->model_map,
e->model_size,
e->up_abs_offset,
e->gate_expert_bytes);
ds4_gpu_stream_expert_evict_dontneed_range(e->model_map,
e->model_size,
e->down_abs_offset,
e->down_expert_bytes);
ds4_gpu_stream_expert_cache_zero_addr_slot(layer, expert);
if (reuse) {
reuse->gate_buffer = e->gate_buffer;
reuse->up_buffer = e->up_buffer;
reuse->down_buffer = e->down_buffer;
reuse->gate_inner = e->gate_inner;
reuse->up_inner = e->up_inner;
reuse->down_inner = e->down_inner;
} else if (e->slab_backed && recycle_slab_slot) {
ds4_gpu_stream_expert_slab_push_free_slot(e->slab_slot);
} else if (!e->slab_backed) {
ds4_gpu_stream_expert_unlock_explicit_buffer(e->gate_buffer);
if (e->up_buffer != e->gate_buffer) {
ds4_gpu_stream_expert_unlock_explicit_buffer(e->up_buffer);
}
if (e->down_buffer != e->gate_buffer &&
e->down_buffer != e->up_buffer) {
ds4_gpu_stream_expert_unlock_explicit_buffer(e->down_buffer);
}
}
e->gate_buffer = nil;
e->up_buffer = nil;
e->down_buffer = nil;
e->model_map = NULL;
e->model_size = 0;
e->gate_abs_offset = 0;
e->up_abs_offset = 0;
e->down_abs_offset = 0;
e->gate_expert_bytes = 0;
e->down_expert_bytes = 0;
e->logical_bytes = 0;
e->last_used = 0;
e->use_count = 0;
e->gate_inner = 0;
e->up_inner = 0;
e->down_inner = 0;
e->inflight_seq = 0;
e->slab_slot = 0;
e->valid = 0;
e->slab_backed = 0;
if (g_stream_expert_cache_layer_count[layer] > 0) {
g_stream_expert_cache_layer_count[layer]--;
}
if (g_stream_expert_cache_entry_count > 0) {
g_stream_expert_cache_entry_count--;
}
if (g_stream_expert_cache_bytes >= bytes) {
g_stream_expert_cache_bytes -= bytes;
} else {
g_stream_expert_cache_bytes = 0;
}
if (count_eviction) {
g_stream_expert_cache_evictions++;
g_stream_expert_cache_layer_evictions[layer]++;
}
}
static void ds4_gpu_stream_expert_cache_clear_entry(
uint32_t layer,
uint32_t expert,
int count_eviction) {
ds4_gpu_stream_expert_cache_clear_entry_internal(layer,
expert,
count_eviction,
1,
NULL);
}
static void ds4_gpu_stream_expert_cache_clear_all(int reset_stats) {
ds4_gpu_stream_expert_pending_load_clear();
g_stream_expert_cache_done_seq = g_stream_expert_cache_cb_seq;
g_stream_expert_cache_batch_seq = 0;
g_stream_expert_cache_owned_seq = 0;
g_stream_expert_cache_pending_max_seq = 0;
for (uint32_t layer = 0; layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER; layer++) {
for (uint32_t expert = 0; expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; expert++) {
g_stream_expert_cache[layer][expert].inflight_seq = 0;
ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0);
}
ds4_gpu_stream_full_expert_addr_clear_layer(layer);
g_stream_expert_cache_layer_count[layer] = 0;
id<MTLBuffer> buffers[3] = {
g_stream_expert_cache_gate_addr_buffers[layer],
g_stream_expert_cache_up_addr_buffers[layer],
g_stream_expert_cache_down_addr_buffers[layer],
};
for (uint32_t i = 0; i < 3; i++) {
if (!buffers[i]) continue;
const NSUInteger bytes =
(NSUInteger)DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT * sizeof(uint64_t);
memset([buffers[i] contents], 0, bytes);
}
}
g_stream_expert_cache_bytes = 0;
g_stream_expert_cache_entry_count = 0;
for (uint32_t i = 0; i < g_stream_expert_cache_slab_count; i++) {
g_stream_expert_cache_slabs[i] = nil;
g_stream_expert_cache_slab_start_slot[i] = 0;
g_stream_expert_cache_slab_slot_count[i] = 0;
g_stream_expert_cache_slab_slots_used[i] = 0;
}
g_stream_expert_cache_slab_count = 0;
g_stream_expert_cache_slab_total_slots = 0;
g_stream_expert_cache_free_slot_count = 0;
g_stream_expert_cache_slab_slot_bytes = 0;
memset(g_stream_expert_cache_slab_slot_locked,
0,
sizeof(g_stream_expert_cache_slab_slot_locked));
if (reset_stats) {
g_stream_expert_cache_hits = 0;
g_stream_expert_cache_misses = 0;
g_stream_expert_cache_evictions = 0;
g_stream_expert_cache_wraps = 0;
g_stream_expert_cache_clock = 0;
g_stream_expert_cache_evict_advise_bytes = 0;
g_stream_expert_cache_willneed_advise_bytes = 0;
g_stream_expert_cache_pread_bytes = 0;
g_stream_expert_cache_pread_ms = 0.0;
g_stream_expert_cache_mlock_bytes = 0;
g_stream_expert_cache_mlock_fail_bytes = 0;
g_stream_expert_cache_mlock_failures = 0;
g_stream_expert_cache_mlock_ms = 0.0;
g_stream_expert_cache_mlock_warned = 0;
g_stream_expert_cache_mlock_budget_cap = 0;
g_stream_expert_cache_mlock_relief_applied = 0;
g_stream_expert_cache_buffer_allocs = 0;
g_stream_expert_cache_buffer_reuses = 0;
g_stream_expert_cache_decode_tokens = 0;
g_stream_expert_cache_first_decode_layer = UINT32_MAX;
g_stream_expert_cache_hotness_decay_token = 0;
memset(g_stream_expert_cache_route_hotness,
0,
sizeof(g_stream_expert_cache_route_hotness));
g_stream_expert_timing_selected_calls = 0;
g_stream_expert_timing_selected_read_ms = 0.0;
g_stream_expert_timing_selected_sync_ms = 0.0;
g_stream_expert_timing_selected_copy_ms = 0.0;
g_stream_expert_timing_selected_bind_ms = 0.0;
g_stream_expert_timing_split_layers = 0;
g_stream_expert_timing_split_resident_experts = 0;
g_stream_expert_timing_split_missing_experts = 0;
g_stream_expert_timing_split_resident_ms = 0.0;
g_stream_expert_timing_split_missing_ms = 0.0;
g_stream_expert_timing_split_missing_load_ms = 0.0;
g_stream_expert_timing_split_missing_slot_ms = 0.0;
g_stream_expert_timing_split_missing_prune_ms = 0.0;
g_stream_expert_timing_split_missing_addr_ms = 0.0;
g_stream_expert_timing_split_missing_wait_ms = 0.0;
g_stream_expert_timing_load_calls = 0;
g_stream_expert_timing_load_prepare_ms = 0.0;
g_stream_expert_timing_load_pread_ms = 0.0;
g_stream_expert_timing_load_install_ms = 0.0;
g_stream_expert_timing_prepare_batch_reuse_calls = 0;
g_stream_expert_timing_prepare_batch_reuse_ms = 0.0;
g_stream_expert_timing_prepare_buffer_calls = 0;
g_stream_expert_timing_prepare_buffer_ms = 0.0;
g_stream_expert_timing_prepare_task_experts = 0;
g_stream_expert_timing_prepare_task_ms = 0.0;
g_stream_expert_timing_reuse_scan_calls = 0;
g_stream_expert_timing_reuse_scan_entries = 0;
g_stream_expert_timing_reuse_scan_ms = 0.0;
g_stream_expert_timing_reuse_clear_ms = 0.0;
g_stream_expert_timing_readahead_calls = 0;
g_stream_expert_timing_readahead_bytes = 0;
g_stream_expert_timing_readahead_ms = 0.0;
g_stream_expert_timing_cache_all_resident_layers = 0;
g_stream_expert_timing_cache_all_missing_layers = 0;
g_stream_expert_timing_cache_mixed_layers = 0;
g_stream_expert_timing_cache_resident_experts = 0;
g_stream_expert_timing_cache_missing_experts = 0;
g_stream_expert_timing_last_report =
(ds4_gpu_stream_expert_timing_snapshot){0};
memset(g_stream_expert_cache_layer_hits,
0,
sizeof(g_stream_expert_cache_layer_hits));
memset(g_stream_expert_cache_layer_misses,
0,
sizeof(g_stream_expert_cache_layer_misses));
memset(g_stream_expert_cache_layer_evictions,
0,
sizeof(g_stream_expert_cache_layer_evictions));
memset(g_stream_expert_cache_layer_pread_bytes,
0,
sizeof(g_stream_expert_cache_layer_pread_bytes));
memset(g_stream_expert_cache_layer_pread_ms,
0,
sizeof(g_stream_expert_cache_layer_pread_ms));
}
}
static int ds4_gpu_stream_expert_cache_is_protected(
uint32_t expert,
const int32_t *protect_ids,
uint32_t n_protect) {
if (!protect_ids) return 0;
for (uint32_t i = 0; i < n_protect; i++) {
if (protect_ids[i] >= 0 && (uint32_t)protect_ids[i] == expert) {
return 1;
}
}
return 0;
}
static void ds4_gpu_stream_expert_cache_prune_layer(
uint32_t layer,
uint32_t n_total_expert,
uint32_t n_selected,
const int32_t *protect_ids,
uint32_t n_protect) {
if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return;
uint32_t cap = ds4_gpu_stream_expert_cache_effective_cap(layer,
n_total_expert,
n_selected);
if (cap == 0) return;
/*
* Route hotness counts selected experts even when they miss. Hit-count
* LFU penalizes experts that are repeatedly selected but evicted before a
* second hit, which keeps too many decode layers in the mixed-cache path.
*/
while (g_stream_expert_cache_layer_count[layer] > cap) {
uint32_t victim = UINT32_MAX;
uint32_t lowest_hotness = UINT32_MAX;
uint64_t oldest = UINT64_MAX;
for (uint32_t expert = 0; expert < n_total_expert; expert++) {
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (!e->valid ||
ds4_gpu_stream_expert_cache_entry_inflight(e) ||
ds4_gpu_stream_expert_cache_is_protected(expert, protect_ids, n_protect)) {
continue;
}
const uint32_t hotness =
g_stream_expert_cache_route_hotness[layer][expert];
if (hotness < lowest_hotness ||
(hotness == lowest_hotness && e->last_used < oldest)) {
lowest_hotness = hotness;
oldest = e->last_used;
victim = expert;
}
}
if (victim == UINT32_MAX) break;
ds4_gpu_stream_expert_cache_clear_entry(layer, victim, 1);
}
}
static int ds4_gpu_stream_expert_cache_entry_protected(
uint32_t layer,
uint32_t expert,
uint32_t protect_layer,
const int32_t *protect_ids,
uint32_t n_protect) {
if (layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER &&
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
ds4_gpu_stream_expert_cache_entry_inflight(
&g_stream_expert_cache[layer][expert])) {
return 1;
}
return layer == protect_layer &&
ds4_gpu_stream_expert_cache_is_protected(expert,
protect_ids,
n_protect);
}
static int ds4_gpu_stream_expert_cache_entry_reusable(
const ds4_gpu_stream_expert_cache_entry *e,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
return e &&
e->valid &&
e->gate_buffer &&
e->up_buffer &&
e->down_buffer &&
e->gate_expert_bytes == gate_expert_bytes &&
e->down_expert_bytes == down_expert_bytes;
}
static int ds4_gpu_stream_expert_cache_take_reusable(
int force_reuse,
uint32_t protect_layer,
const int32_t *protect_ids,
uint32_t n_protect,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
ds4_gpu_stream_expert_reusable_buffers *reuse) {
if (!reuse) return 0;
reuse->gate_buffer = nil;
reuse->up_buffer = nil;
reuse->down_buffer = nil;
reuse->gate_inner = 0;
reuse->up_inner = 0;
reuse->down_inner = 0;
const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
if (budget == 0 ||
(!force_reuse && g_stream_expert_cache_entry_count < budget)) {
return 0;
}
int waited_inflight = 0;
retry:
;
uint32_t victim_layer = UINT32_MAX;
uint32_t victim_expert = UINT32_MAX;
uint32_t lowest_hotness = UINT32_MAX;
uint64_t oldest = UINT64_MAX;
int skipped_inflight = 0;
const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0;
uint64_t scan_entries = 0;
for (uint32_t layer = 0;
layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
layer++) {
for (uint32_t expert = 0;
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
expert++) {
scan_entries++;
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (!ds4_gpu_stream_expert_cache_entry_reusable(e,
gate_expert_bytes,
down_expert_bytes)) {
continue;
}
if (ds4_gpu_stream_expert_cache_entry_inflight(e)) {
skipped_inflight = 1;
continue;
}
if (ds4_gpu_stream_expert_cache_entry_protected(layer,
expert,
protect_layer,
protect_ids,
n_protect)) {
continue;
}
const uint32_t hotness =
g_stream_expert_cache_route_hotness[layer][expert];
if (hotness < lowest_hotness ||
(hotness == lowest_hotness && e->last_used < oldest)) {
lowest_hotness = hotness;
oldest = e->last_used;
victim_layer = layer;
victim_expert = expert;
}
}
}
if (timing) {
ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries,
ds4_gpu_now_ms() - scan_t0);
}
if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) {
if (skipped_inflight && !waited_inflight &&
!ds4_gpu_stream_expert_cache_on_service_thread()) {
waited_inflight = 1;
if (!ds4_gpu_stream_expert_cache_wait_inflight(
"streaming expert cache reuse")) {
return 0;
}
goto retry;
}
return 0;
}
const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0;
ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer,
victim_expert,
1,
1,
reuse);
if (timing) {
ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() -
clear_t0);
}
if (!reuse->gate_buffer || !reuse->up_buffer || !reuse->down_buffer) {
reuse->gate_buffer = nil;
reuse->up_buffer = nil;
reuse->down_buffer = nil;
reuse->gate_inner = 0;
reuse->up_inner = 0;
reuse->down_inner = 0;
return 0;
}
g_stream_expert_cache_buffer_reuses +=
ds4_gpu_stream_expert_buffer_object_count(reuse->gate_buffer,
reuse->up_buffer,
reuse->down_buffer);
return 1;
}
static uint32_t ds4_gpu_stream_expert_cache_take_reusable_batch(
uint32_t n_needed,
uint32_t protect_layer,
const int32_t *protect_ids,
uint32_t n_protect,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
ds4_gpu_stream_expert_reusable_buffers *reuses) {
if (!reuses || n_needed == 0) return 0;
if (n_needed > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
n_needed = DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED;
}
for (uint32_t i = 0; i < n_needed; i++) {
reuses[i].gate_buffer = nil;
reuses[i].up_buffer = nil;
reuses[i].down_buffer = nil;
reuses[i].gate_inner = 0;
reuses[i].up_inner = 0;
reuses[i].down_inner = 0;
}
const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
if (budget == 0 || g_stream_expert_cache_entry_count < budget) {
return 0;
}
int waited_inflight = 0;
retry:
;
uint32_t victim_layers[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint32_t victim_experts[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint32_t victim_hotness[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint64_t victim_last_used[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint32_t victim_count = 0;
int skipped_inflight = 0;
const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
const double scan_t0 = timing ? ds4_gpu_now_ms() : 0.0;
uint64_t scan_entries = 0;
for (uint32_t i = 0; i < n_needed; i++) {
victim_layers[i] = UINT32_MAX;
victim_experts[i] = UINT32_MAX;
victim_hotness[i] = UINT32_MAX;
victim_last_used[i] = UINT64_MAX;
}
for (uint32_t layer = 0;
layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
layer++) {
for (uint32_t expert = 0;
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
expert++) {
scan_entries++;
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (!ds4_gpu_stream_expert_cache_entry_reusable(e,
gate_expert_bytes,
down_expert_bytes)) {
continue;
}
if (ds4_gpu_stream_expert_cache_entry_inflight(e)) {
skipped_inflight = 1;
continue;
}
if (ds4_gpu_stream_expert_cache_entry_protected(layer,
expert,
protect_layer,
protect_ids,
n_protect)) {
continue;
}
const uint32_t hotness =
g_stream_expert_cache_route_hotness[layer][expert];
const uint64_t last_used = e->last_used;
if (victim_count < n_needed) {
victim_layers[victim_count] = layer;
victim_experts[victim_count] = expert;
victim_hotness[victim_count] = hotness;
victim_last_used[victim_count] = last_used;
victim_count++;
continue;
}
uint32_t worst = 0;
for (uint32_t i = 1; i < victim_count; i++) {
if (victim_hotness[i] > victim_hotness[worst] ||
(victim_hotness[i] == victim_hotness[worst] &&
victim_last_used[i] > victim_last_used[worst])) {
worst = i;
}
}
if (hotness < victim_hotness[worst] ||
(hotness == victim_hotness[worst] &&
last_used < victim_last_used[worst])) {
victim_layers[worst] = layer;
victim_experts[worst] = expert;
victim_hotness[worst] = hotness;
victim_last_used[worst] = last_used;
}
}
}
if (timing) {
ds4_gpu_stream_expert_timing_note_reuse_scan(scan_entries,
ds4_gpu_now_ms() - scan_t0);
}
if (victim_count == 0) {
if (skipped_inflight && !waited_inflight) {
waited_inflight = 1;
if (!ds4_gpu_stream_expert_cache_wait_inflight(
"streaming expert cache batch reuse")) {
return 0;
}
goto retry;
}
return 0;
}
uint32_t reuse_count = 0;
const double clear_t0 = timing ? ds4_gpu_now_ms() : 0.0;
for (uint32_t i = 0; i < victim_count; i++) {
if (victim_layers[i] == UINT32_MAX ||
victim_experts[i] == UINT32_MAX) {
continue;
}
ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layers[i],
victim_experts[i],
1,
1,
&reuses[reuse_count]);
if (!reuses[reuse_count].gate_buffer ||
!reuses[reuse_count].up_buffer ||
!reuses[reuse_count].down_buffer) {
reuses[reuse_count].gate_buffer = nil;
reuses[reuse_count].up_buffer = nil;
reuses[reuse_count].down_buffer = nil;
reuses[reuse_count].gate_inner = 0;
reuses[reuse_count].up_inner = 0;
reuses[reuse_count].down_inner = 0;
continue;
}
g_stream_expert_cache_buffer_reuses +=
ds4_gpu_stream_expert_buffer_object_count(
reuses[reuse_count].gate_buffer,
reuses[reuse_count].up_buffer,
reuses[reuse_count].down_buffer);
reuse_count++;
}
if (timing) {
ds4_gpu_stream_expert_timing_note_reuse_clear(ds4_gpu_now_ms() -
clear_t0);
}
return reuse_count;
}
static int ds4_gpu_stream_expert_batch_reuse_enabled(
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) {
return 0;
}
const uint64_t slot_bytes = gate_expert_bytes * 2ull + down_expert_bytes;
/*
* One global victim scan per selected miss is measurable for GLM Q2-size
* slots, but batching larger Q4-size slots regressed short decode on M5.
* Keep larger slots on the older single-victim path until profiling says
* otherwise.
*/
return slot_bytes <= 16ull * 1024ull * 1024ull;
}
static uint32_t ds4_gpu_stream_expert_cache_release_mlock_margin(
uint32_t protect_layer,
const int32_t *protect_ids,
uint32_t n_protect) {
if (g_stream_expert_cache_mlock_relief_applied ||
g_stream_expert_cache_slab_slot_bytes == 0) {
return 0;
}
const uint32_t locked_before =
ds4_gpu_stream_expert_slab_locked_slot_count();
if (locked_before == 0) return 0;
const uint32_t target_release = (locked_before + 9u) / 10u;
uint32_t released = 0;
while (released < target_release) {
uint32_t victim_layer = UINT32_MAX;
uint32_t victim_expert = UINT32_MAX;
uint32_t victim_slot = UINT32_MAX;
uint32_t lowest_hotness = UINT32_MAX;
uint64_t oldest = UINT64_MAX;
for (uint32_t layer = 0;
layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
layer++) {
for (uint32_t expert = 0;
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
expert++) {
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (!e->valid ||
!e->slab_backed ||
e->slab_slot >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_ENTRIES ||
!g_stream_expert_cache_slab_slot_locked[e->slab_slot] ||
ds4_gpu_stream_expert_cache_entry_inflight(e) ||
ds4_gpu_stream_expert_cache_entry_protected(layer,
expert,
protect_layer,
protect_ids,
n_protect)) {
continue;
}
const uint32_t hotness =
g_stream_expert_cache_route_hotness[layer][expert];
if (hotness < lowest_hotness ||
(hotness == lowest_hotness && e->last_used < oldest)) {
lowest_hotness = hotness;
oldest = e->last_used;
victim_layer = layer;
victim_expert = expert;
victim_slot = e->slab_slot;
}
}
}
if (victim_layer == UINT32_MAX ||
victim_expert == UINT32_MAX ||
victim_slot == UINT32_MAX) {
break;
}
ds4_gpu_stream_expert_cache_clear_entry_internal(victim_layer,
victim_expert,
1,
0,
NULL);
if (g_stream_expert_cache[victim_layer][victim_expert].valid) {
break;
}
if (ds4_gpu_stream_expert_slab_unlock_slot(victim_slot)) {
released++;
}
}
if (released == 0) return 0;
g_stream_expert_cache_mlock_relief_applied = 1;
uint32_t cap = g_stream_expert_cache_entry_count;
const uint32_t locked_after =
ds4_gpu_stream_expert_slab_locked_slot_count();
if (locked_after != 0 && locked_after < cap) cap = locked_after;
if (cap != 0 &&
(g_stream_expert_cache_mlock_budget_cap == 0 ||
cap < g_stream_expert_cache_mlock_budget_cap)) {
g_stream_expert_cache_mlock_budget_cap = cap;
}
const uint64_t released_bytes =
(uint64_t)released * g_stream_expert_cache_slab_slot_bytes;
fprintf(stderr,
"ds4: released locked cache margin: %u experts / %.2f GiB; "
"runtime cache cap now %u experts\n",
released,
ds4_gpu_gib(released_bytes),
g_stream_expert_cache_mlock_budget_cap);
return released;
}
static int ds4_gpu_stream_expert_cache_take_capped_reusable(
uint32_t protect_layer,
const int32_t *protect_ids,
uint32_t n_protect,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
__strong id<MTLBuffer> *gate_buf,
__strong id<MTLBuffer> *up_buf,
__strong id<MTLBuffer> *down_buf,
NSUInteger *gate_inner,
NSUInteger *up_inner,
NSUInteger *down_inner) {
if (g_stream_expert_cache_mlock_budget_cap == 0) return 0;
ds4_gpu_stream_expert_cache_release_mlock_margin(protect_layer,
protect_ids,
n_protect);
ds4_gpu_stream_expert_reusable_buffers reuse = { nil, nil, nil, 0, 0, 0 };
if (!ds4_gpu_stream_expert_cache_take_reusable(1,
protect_layer,
protect_ids,
n_protect,
gate_expert_bytes,
down_expert_bytes,
&reuse)) {
return 0;
}
*gate_buf = reuse.gate_buffer;
*up_buf = reuse.up_buffer;
*down_buf = reuse.down_buffer;
*gate_inner = reuse.gate_inner;
*up_inner = reuse.up_inner;
*down_inner = reuse.down_inner;
return *gate_buf && *up_buf && *down_buf;
}
static int ds4_gpu_stream_expert_cache_prepare_load_buffers(
uint32_t layer,
uint32_t expert,
uint32_t protect_layer,
const int32_t *protect_ids,
uint32_t n_protect,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
int force_reuse,
__strong id<MTLBuffer> *gate_buf,
__strong id<MTLBuffer> *up_buf,
__strong id<MTLBuffer> *down_buf,
NSUInteger *gate_inner,
NSUInteger *up_inner,
NSUInteger *down_inner) {
if (!gate_buf || !up_buf || !down_buf ||
!gate_inner || !up_inner || !down_inner ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
return 0;
}
*gate_buf = nil;
*up_buf = nil;
*down_buf = nil;
*gate_inner = 0;
*up_inner = 0;
*down_inner = 0;
ds4_gpu_stream_expert_reusable_buffers reuse = { nil, nil, nil, 0, 0, 0 };
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (e->valid && ds4_gpu_stream_expert_cache_entry_inflight(e)) {
if (ds4_gpu_stream_expert_cache_on_service_thread()) return 0;
if (!ds4_gpu_stream_expert_cache_wait_inflight(
"streaming expert cache replacement")) {
return 0;
}
e = &g_stream_expert_cache[layer][expert];
}
if (ds4_gpu_stream_expert_cache_entry_reusable(e,
gate_expert_bytes,
down_expert_bytes)) {
ds4_gpu_stream_expert_cache_clear_entry_internal(layer,
expert,
0,
1,
&reuse);
if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) {
g_stream_expert_cache_buffer_reuses +=
ds4_gpu_stream_expert_buffer_object_count(reuse.gate_buffer,
reuse.up_buffer,
reuse.down_buffer);
}
} else if (e->valid) {
ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0);
}
if (!reuse.gate_buffer || !reuse.up_buffer || !reuse.down_buffer) {
if (!ds4_gpu_stream_expert_cache_take_reusable(force_reuse,
protect_layer,
protect_ids,
n_protect,
gate_expert_bytes,
down_expert_bytes,
&reuse)) {
reuse.gate_buffer = nil;
reuse.up_buffer = nil;
reuse.down_buffer = nil;
}
}
if (reuse.gate_buffer && reuse.up_buffer && reuse.down_buffer) {
*gate_buf = reuse.gate_buffer;
*up_buf = reuse.up_buffer;
*down_buf = reuse.down_buffer;
*gate_inner = reuse.gate_inner;
*up_inner = reuse.up_inner;
*down_inner = reuse.down_inner;
return 1;
}
if (ds4_gpu_stream_expert_combined_buffer_enabled()) {
if (gate_expert_bytes > UINT64_MAX - gate_expert_bytes ||
gate_expert_bytes * 2ull > UINT64_MAX - down_expert_bytes ||
gate_expert_bytes > (uint64_t)NSUIntegerMax ||
gate_expert_bytes * 2ull > (uint64_t)NSUIntegerMax ||
gate_expert_bytes * 2ull + down_expert_bytes >
(uint64_t)NSUIntegerMax) {
return 0;
}
if (g_stream_expert_cache_mlock_budget_cap != 0) {
return ds4_gpu_stream_expert_cache_take_capped_reusable(
protect_layer, protect_ids, n_protect,
gate_expert_bytes, down_expert_bytes,
gate_buf, up_buf, down_buf,
gate_inner, up_inner, down_inner);
}
if (ds4_gpu_stream_expert_alloc_slab_slot(gate_expert_bytes,
down_expert_bytes,
gate_buf,
up_buf,
down_buf,
gate_inner,
up_inner,
down_inner)) {
return 1;
}
if (g_stream_expert_cache_mlock_budget_cap != 0) {
return ds4_gpu_stream_expert_cache_take_capped_reusable(
protect_layer, protect_ids, n_protect,
gate_expert_bytes, down_expert_bytes,
gate_buf, up_buf, down_buf,
gate_inner, up_inner, down_inner);
}
const uint64_t up_off = gate_expert_bytes;
const uint64_t down_off = gate_expert_bytes * 2ull;
const uint64_t combined_bytes = down_off + down_expert_bytes;
id<MTLBuffer> combined =
ds4_gpu_stream_expert_alloc_buffer(combined_bytes,
@"ds4_stream_expert_combined");
if (!combined) {
return ds4_gpu_stream_expert_cache_take_capped_reusable(
protect_layer, protect_ids, n_protect,
gate_expert_bytes, down_expert_bytes,
gate_buf, up_buf, down_buf,
gate_inner, up_inner, down_inner);
}
*gate_buf = combined;
*up_buf = combined;
*down_buf = combined;
*gate_inner = 0;
*up_inner = (NSUInteger)up_off;
*down_inner = (NSUInteger)down_off;
return 1;
}
id<MTLBuffer> explicit_gate =
ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes,
@"ds4_stream_expert_gate");
id<MTLBuffer> explicit_up =
ds4_gpu_stream_expert_alloc_buffer(gate_expert_bytes,
@"ds4_stream_expert_up");
id<MTLBuffer> explicit_down =
ds4_gpu_stream_expert_alloc_buffer(down_expert_bytes,
@"ds4_stream_expert_down");
if (!explicit_gate || !explicit_up || !explicit_down) {
ds4_gpu_stream_expert_unlock_explicit_buffer(explicit_gate);
ds4_gpu_stream_expert_unlock_explicit_buffer(explicit_up);
ds4_gpu_stream_expert_unlock_explicit_buffer(explicit_down);
return ds4_gpu_stream_expert_cache_take_capped_reusable(
protect_layer, protect_ids, n_protect,
gate_expert_bytes, down_expert_bytes,
gate_buf, up_buf, down_buf,
gate_inner, up_inner, down_inner);
}
*gate_buf = explicit_gate;
*up_buf = explicit_up;
*down_buf = explicit_down;
return 1;
}
static void ds4_gpu_stream_expert_cache_prune_global(
uint32_t protect_layer,
const int32_t *protect_ids,
uint32_t n_protect) {
const uint32_t budget = ds4_gpu_stream_expert_cache_configured_budget();
if (budget == 0 || g_stream_expert_cache_entry_count <= budget) return;
while (g_stream_expert_cache_entry_count > budget) {
uint32_t victim_layer = UINT32_MAX;
uint32_t victim_expert = UINT32_MAX;
uint32_t lowest_hotness = UINT32_MAX;
uint64_t oldest = UINT64_MAX;
for (uint32_t layer = 0;
layer < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER;
layer++) {
for (uint32_t expert = 0;
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
expert++) {
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (!e->valid ||
ds4_gpu_stream_expert_cache_entry_protected(layer,
expert,
protect_layer,
protect_ids,
n_protect)) {
continue;
}
const uint32_t hotness =
g_stream_expert_cache_route_hotness[layer][expert];
if (hotness < lowest_hotness ||
(hotness == lowest_hotness && e->last_used < oldest)) {
lowest_hotness = hotness;
oldest = e->last_used;
victim_layer = layer;
victim_expert = expert;
}
}
}
if (victim_layer == UINT32_MAX || victim_expert == UINT32_MAX) break;
ds4_gpu_stream_expert_cache_clear_entry(victim_layer, victim_expert, 1);
}
}
static int ds4_gpu_stream_expert_cache_entry_matches(
const ds4_gpu_stream_expert_cache_entry *e,
const void *model_map,
uint64_t model_size,
uint64_t gate_abs_offset,
uint64_t up_abs_offset,
uint64_t down_abs_offset,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
return e &&
e->valid &&
e->model_map == model_map &&
e->model_size == model_size &&
e->gate_abs_offset == gate_abs_offset &&
e->up_abs_offset == up_abs_offset &&
e->down_abs_offset == down_abs_offset &&
e->gate_expert_bytes == gate_expert_bytes &&
e->down_expert_bytes == down_expert_bytes &&
e->gate_buffer && e->up_buffer && e->down_buffer;
}
static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_peek(
const void *model_map,
uint64_t model_size,
uint32_t layer,
uint32_t expert,
uint32_t n_total_expert,
uint32_t n_selected,
uint64_t gate_abs_offset,
uint64_t up_abs_offset,
uint64_t down_abs_offset,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
if (!g_ssd_streaming_mode ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
expert >= n_total_expert ||
!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) ||
ds4_gpu_stream_expert_cache_effective_cap(layer,
n_total_expert,
n_selected) == 0) {
return NULL;
}
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (!ds4_gpu_stream_expert_cache_entry_matches(e,
model_map,
model_size,
gate_abs_offset,
up_abs_offset,
down_abs_offset,
gate_expert_bytes,
down_expert_bytes)) {
return NULL;
}
e->last_used = ++g_stream_expert_cache_clock;
e->use_count++;
g_stream_expert_cache_hits++;
g_stream_expert_cache_layer_hits[layer]++;
return e;
}
static ds4_gpu_stream_expert_cache_entry *
ds4_gpu_stream_expert_cache_install_loaded(
const void *model_map,
uint64_t model_size,
uint32_t layer,
uint32_t expert,
uint64_t gate_abs_offset,
uint64_t up_abs_offset,
uint64_t down_abs_offset,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
id<MTLBuffer> gate_buf,
id<MTLBuffer> up_buf,
id<MTLBuffer> down_buf,
NSUInteger gate_inner,
NSUInteger up_inner,
NSUInteger down_inner) {
if (!gate_buf || !up_buf || !down_buf ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT) {
return NULL;
}
if (gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) {
fprintf(stderr, "ds4: Metal streaming expert cache byte size overflow\n");
return NULL;
}
const uint64_t logical_bytes = gate_expert_bytes * 2ull + down_expert_bytes;
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (e->valid) {
if (ds4_gpu_stream_expert_cache_entry_inflight(e) &&
ds4_gpu_stream_expert_cache_on_service_thread()) {
return NULL;
}
if (ds4_gpu_stream_expert_cache_entry_inflight(e) &&
!ds4_gpu_stream_expert_cache_wait_inflight(
"streaming expert cache install")) {
return NULL;
}
if (ds4_gpu_stream_expert_cache_entry_inflight(e)) return NULL;
ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0);
if (e->valid) return NULL;
}
if (!ds4_gpu_stream_expert_cache_set_addr_slot(layer,
expert,
gate_buf,
gate_inner,
up_buf,
up_inner,
down_buf,
down_inner)) {
return NULL;
}
e->gate_buffer = gate_buf;
e->up_buffer = up_buf;
e->down_buffer = down_buf;
e->model_map = model_map;
e->model_size = model_size;
e->gate_abs_offset = gate_abs_offset;
e->up_abs_offset = up_abs_offset;
e->down_abs_offset = down_abs_offset;
e->gate_expert_bytes = gate_expert_bytes;
e->down_expert_bytes = down_expert_bytes;
e->logical_bytes = logical_bytes;
e->last_used = ++g_stream_expert_cache_clock;
e->use_count = 1;
e->gate_inner = gate_inner;
e->up_inner = up_inner;
e->down_inner = down_inner;
e->inflight_seq = 0;
uint32_t slab_slot = 0;
if (gate_buf == up_buf &&
gate_buf == down_buf &&
ds4_gpu_stream_expert_slab_slot_for_buffer(gate_buf,
gate_inner,
&slab_slot)) {
e->slab_backed = 1;
e->slab_slot = slab_slot;
} else {
e->slab_backed = 0;
e->slab_slot = 0;
}
e->valid = 1;
g_stream_expert_cache_layer_count[layer]++;
if (g_stream_expert_cache_entry_count < UINT32_MAX) {
g_stream_expert_cache_entry_count++;
}
if (g_stream_expert_cache_bytes > UINT64_MAX - logical_bytes) {
g_stream_expert_cache_bytes = UINT64_MAX;
} else {
g_stream_expert_cache_bytes += logical_bytes;
}
g_stream_expert_cache_misses++;
g_stream_expert_cache_layer_misses[layer]++;
g_stream_expert_cache_wraps += 3;
return e;
}
static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get_protected(
const void *model_map,
uint64_t model_size,
uint32_t layer,
uint32_t expert,
uint32_t n_total_expert,
uint32_t n_selected,
uint64_t gate_abs_offset,
uint64_t up_abs_offset,
uint64_t down_abs_offset,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
const int32_t *protect_ids,
uint32_t n_protect) {
if (!g_ssd_streaming_mode ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
expert >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
expert >= n_total_expert ||
!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) ||
ds4_gpu_stream_expert_cache_effective_cap(layer,
n_total_expert,
n_selected) == 0) {
return NULL;
}
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][expert];
if (ds4_gpu_stream_expert_cache_entry_matches(e,
model_map,
model_size,
gate_abs_offset,
up_abs_offset,
down_abs_offset,
gate_expert_bytes,
down_expert_bytes)) {
e->last_used = ++g_stream_expert_cache_clock;
e->use_count++;
g_stream_expert_cache_hits++;
g_stream_expert_cache_layer_hits[layer]++;
return e;
}
ds4_gpu_stream_expert_readahead_range(gate_abs_offset, gate_expert_bytes);
ds4_gpu_stream_expert_readahead_range(up_abs_offset, gate_expert_bytes);
ds4_gpu_stream_expert_readahead_range(down_abs_offset, down_expert_bytes);
id<MTLBuffer> gate_buf = nil;
id<MTLBuffer> up_buf = nil;
id<MTLBuffer> down_buf = nil;
NSUInteger gate_inner = 0;
NSUInteger up_inner = 0;
NSUInteger down_inner = 0;
const int32_t protect_one = (int32_t)expert;
if (!protect_ids || n_protect == 0) {
protect_ids = &protect_one;
n_protect = 1;
}
const uint32_t cache_budget =
ds4_gpu_stream_expert_cache_configured_budget();
const int force_reuse =
cache_budget != 0 && g_stream_expert_cache_entry_count >= cache_budget;
if (!ds4_gpu_stream_expert_cache_prepare_load_buffers(layer,
expert,
layer,
protect_ids,
n_protect,
gate_expert_bytes,
down_expert_bytes,
force_reuse,
&gate_buf,
&up_buf,
&down_buf,
&gate_inner,
&up_inner,
&down_inner)) {
return NULL;
}
if (!gate_buf || !up_buf || !down_buf) return NULL;
uint8_t *gate_dst = (uint8_t *)[gate_buf contents] + gate_inner;
uint8_t *up_dst = (uint8_t *)[up_buf contents] + up_inner;
uint8_t *down_dst = (uint8_t *)[down_buf contents] + down_inner;
if (!gate_dst || !up_dst || !down_dst) return NULL;
ds4_gpu_stream_expert_pread_task tasks[3] = {
{
.offset = gate_abs_offset,
.len = gate_expert_bytes,
.dst = gate_dst,
},
{
.offset = up_abs_offset,
.len = gate_expert_bytes,
.dst = up_dst,
},
{
.offset = down_abs_offset,
.len = down_expert_bytes,
.dst = down_dst,
},
};
uint64_t read_bytes = 0;
double read_ms = 0.0;
if (!ds4_gpu_stream_expert_pread_tasks(tasks, 3, &read_bytes, &read_ms)) {
return NULL;
}
ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms);
if (getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal streaming expert parallel pread layer=%u experts=1 tensors=3 "
"threads=%u bytes=%.2f GiB wall=%.3f ms\n",
layer,
ds4_gpu_stream_expert_pread_thread_count(3),
ds4_gpu_gib(read_bytes),
read_ms);
}
return ds4_gpu_stream_expert_cache_install_loaded(model_map,
model_size,
layer,
expert,
gate_abs_offset,
up_abs_offset,
down_abs_offset,
gate_expert_bytes,
down_expert_bytes,
gate_buf,
up_buf,
down_buf,
gate_inner,
up_inner,
down_inner);
}
static ds4_gpu_stream_expert_cache_entry *ds4_gpu_stream_expert_cache_get(
const void *model_map,
uint64_t model_size,
uint32_t layer,
uint32_t expert,
uint32_t n_total_expert,
uint32_t n_selected,
uint64_t gate_abs_offset,
uint64_t up_abs_offset,
uint64_t down_abs_offset,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
return ds4_gpu_stream_expert_cache_get_protected(model_map,
model_size,
layer,
expert,
n_total_expert,
n_selected,
gate_abs_offset,
up_abs_offset,
down_abs_offset,
gate_expert_bytes,
down_expert_bytes,
NULL,
0);
}
static int ds4_gpu_stream_expert_pending_load_profile_enabled(void) {
return getenv("DS4_METAL_STREAMING_EXPERT_EARLY_LOAD_PROFILE") != NULL;
}
static void ds4_gpu_stream_expert_pending_load_release_buffers(
ds4_gpu_stream_expert_pending_load *p) {
if (!p) return;
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
p->gate_bufs[i] = nil;
p->up_bufs[i] = nil;
p->down_bufs[i] = nil;
p->gate_inners[i] = 0;
p->up_inners[i] = 0;
p->down_inners[i] = 0;
}
}
static int ds4_gpu_stream_expert_pending_load_install(
ds4_gpu_stream_expert_pending_load *p,
ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED],
double elapsed_ms) {
if (!p || p->n_loads == 0) return 1;
uint64_t read_bytes = 0;
int ok = 1;
for (uint32_t i = 0; i < p->n_tasks; i++) {
if (!p->tasks[i].ok) ok = 0;
if (read_bytes > UINT64_MAX - p->tasks[i].read_bytes) {
read_bytes = UINT64_MAX;
} else {
read_bytes += p->tasks[i].read_bytes;
}
}
if (!ok) return 0;
ds4_gpu_stream_expert_cache_note_pread(p->layer, read_bytes, elapsed_ms);
const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled();
double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
double load_install_ms = 0.0;
ds4_gpu_stream_expert_cache_entry
*loaded_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
loaded_entries[i] = NULL;
}
for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) {
const uint32_t slot = p->load_slots[load_i];
const uint32_t expert = (uint32_t)p->selected_ids[slot];
ds4_gpu_stream_expert_cache_entry *entry =
ds4_gpu_stream_expert_cache_install_loaded(p->model_map,
p->model_size,
p->layer,
expert,
p->gate_abs_offsets[slot],
p->up_abs_offsets[slot],
p->down_abs_offsets[slot],
p->gate_expert_bytes,
p->down_expert_bytes,
p->gate_bufs[load_i],
p->up_bufs[load_i],
p->down_bufs[load_i],
p->gate_inners[load_i],
p->up_inners[load_i],
p->down_inners[load_i]);
if (!entry) return 0;
loaded_entries[slot] = entry;
if (entries) entries[slot] = entry;
}
for (uint32_t i = 0; i < p->n_selected; i++) {
if ((p->missing_mask & (1u << i)) == 0) continue;
if (entries && entries[i]) continue;
const uint32_t source = p->source_slots[i];
if (source >= p->n_selected) return 0;
ds4_gpu_stream_expert_cache_entry *entry = entries && entries[source] ?
entries[source] : loaded_entries[source];
if (!entry) return 0;
entry->use_count++;
if (entries) entries[i] = entry;
}
if (load_timing) {
load_install_ms = ds4_gpu_now_ms() - load_t0;
ds4_gpu_stream_expert_timing_note_load_detail(p->prepare_ms,
elapsed_ms,
load_install_ms);
}
if (ds4_gpu_stream_expert_pending_load_profile_enabled()) {
fprintf(stderr,
"ds4: Metal streaming expert early-load finish layer=%u experts=%u tensors=%u bytes=%.2f GiB wall=%.3f ms\n",
p->layer,
p->n_loads,
p->n_tasks,
ds4_gpu_gib(read_bytes),
elapsed_ms);
}
return 1;
}
static int ds4_gpu_stream_expert_pending_load_finish(
ds4_gpu_stream_expert_cache_entry *entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED]) {
ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load;
if (!p->active) return 1;
const double start_ms = p->start_ms;
if (!ds4_gpu_stream_expert_pread_pool_wait()) {
ds4_gpu_stream_expert_pending_load_release_buffers(p);
p->active = 0;
return 0;
}
const double elapsed_ms = ds4_gpu_now_ms() - start_ms;
p->active = 0;
const int ok = ds4_gpu_stream_expert_pending_load_install(p,
entries,
elapsed_ms);
ds4_gpu_stream_expert_pending_load_release_buffers(p);
p->n_tasks = 0;
p->n_loads = 0;
p->prepare_ms = 0.0;
return ok;
}
static void ds4_gpu_stream_expert_pending_load_clear(void) {
if (!g_stream_expert_pending_load.active) {
ds4_gpu_stream_expert_pending_load_release_buffers(
&g_stream_expert_pending_load);
return;
}
(void)ds4_gpu_stream_expert_pending_load_finish(NULL);
}
static int ds4_gpu_stream_expert_pending_load_matches(
const void *model_map,
uint64_t model_size,
uint32_t layer,
const int32_t *selected_ids,
uint32_t n_total_expert,
uint32_t n_selected,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes) {
ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load;
if (!p->active ||
p->model_map != model_map ||
p->model_size != model_size ||
p->layer != layer ||
p->n_total_expert != n_total_expert ||
p->n_selected != n_selected ||
p->gate_expert_bytes != gate_expert_bytes ||
p->down_expert_bytes != down_expert_bytes) {
return 0;
}
for (uint32_t i = 0; i < n_selected; i++) {
if (p->selected_ids[i] != selected_ids[i]) return 0;
}
return 1;
}
int ds4_gpu_stream_expert_cache_begin_selected_load(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_selected) {
if (!g_ssd_streaming_mode ||
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_EARLY_LOAD") != NULL) {
return 1;
}
if (!table) return 0;
const void *model_map = table->model_map;
const uint64_t model_size = table->model_size;
const uint32_t layer = table->layer;
const uint32_t n_total_expert = table->n_total_expert;
const uint64_t gate_offset = table->gate_offset;
const uint64_t up_offset = table->up_offset;
const uint64_t down_offset = table->down_offset;
const uint64_t gate_expert_bytes = table->gate_expert_bytes;
const uint64_t down_expert_bytes = table->down_expert_bytes;
if (!model_map || !selected_ids ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
n_selected == 0 ||
n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED ||
n_total_expert == 0 ||
n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) ||
ds4_gpu_stream_expert_cache_effective_cap(layer,
n_total_expert,
n_selected) == 0) {
return 1;
}
if (!g_initialized && !ds4_gpu_init()) return 0;
if (ds4_gpu_stream_expert_pending_load_matches(model_map,
model_size,
layer,
selected_ids,
n_total_expert,
n_selected,
gate_expert_bytes,
down_expert_bytes)) {
return 1;
}
ds4_gpu_stream_expert_pending_load_clear();
ds4_gpu_stream_expert_pending_load *p = &g_stream_expert_pending_load;
p->active = 0;
p->model_map = model_map;
p->model_size = model_size;
p->layer = layer;
p->n_total_expert = n_total_expert;
p->n_selected = n_selected;
p->missing_mask = 0;
p->n_loads = 0;
p->n_tasks = 0;
p->gate_expert_bytes = gate_expert_bytes;
p->down_expert_bytes = down_expert_bytes;
p->prepare_ms = 0.0;
const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled();
const double load_prepare_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
p->selected_ids[i] = -1;
p->load_slots[i] = 0;
p->source_slots[i] = UINT32_MAX;
p->gate_abs_offsets[i] = 0;
p->up_abs_offsets[i] = 0;
p->down_abs_offsets[i] = 0;
p->gate_bufs[i] = nil;
p->up_bufs[i] = nil;
p->down_bufs[i] = nil;
p->gate_inners[i] = 0;
p->up_inners[i] = 0;
p->down_inners[i] = 0;
}
memset(p->tasks, 0, sizeof(p->tasks));
for (uint32_t i = 0; i < n_selected; i++) {
if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) {
fprintf(stderr,
"ds4: Metal streaming early-load expert id %d is outside 0..%u\n",
selected_ids[i],
n_total_expert);
return 0;
}
p->selected_ids[i] = selected_ids[i];
const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i];
if (expert_id > UINT64_MAX / gate_expert_bytes ||
expert_id > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n");
return 0;
}
const uint64_t gate_rel = expert_id * gate_expert_bytes;
const uint64_t down_rel = expert_id * down_expert_bytes;
if (gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal streaming early-load offset overflow\n");
return 0;
}
p->gate_abs_offsets[i] = gate_offset + gate_rel;
p->up_abs_offsets[i] = up_offset + gate_rel;
p->down_abs_offsets[i] = down_offset + down_rel;
ds4_gpu_stream_expert_cache_entry *e =
&g_stream_expert_cache[layer][(uint32_t)selected_ids[i]];
if (ds4_gpu_stream_expert_cache_entry_matches(e,
model_map,
model_size,
p->gate_abs_offsets[i],
p->up_abs_offsets[i],
p->down_abs_offsets[i],
gate_expert_bytes,
down_expert_bytes)) {
continue;
}
uint32_t source = UINT32_MAX;
for (uint32_t prev = 0; prev < i; prev++) {
if (selected_ids[prev] == selected_ids[i] &&
p->gate_abs_offsets[prev] == p->gate_abs_offsets[i] &&
p->up_abs_offsets[prev] == p->up_abs_offsets[i] &&
p->down_abs_offsets[prev] == p->down_abs_offsets[i] &&
(p->missing_mask & (1u << prev)) != 0) {
source = p->source_slots[prev] != UINT32_MAX ?
p->source_slots[prev] : prev;
break;
}
}
if (source != UINT32_MAX) {
p->source_slots[i] = source;
p->missing_mask |= 1u << i;
continue;
}
p->source_slots[i] = i;
p->missing_mask |= 1u << i;
p->load_slots[p->n_loads++] = i;
}
if (p->n_loads == 0) return 1;
const uint32_t cache_budget =
ds4_gpu_stream_expert_cache_configured_budget();
uint32_t reserved_entries = g_stream_expert_cache_entry_count;
ds4_gpu_stream_expert_reusable_buffers
batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
batch_reuse[i] =
(ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 };
}
uint32_t batch_reuse_count = 0;
if (cache_budget != 0 &&
reserved_entries >= cache_budget &&
p->n_loads > 1 &&
ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes,
down_expert_bytes)) {
const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
batch_reuse_count =
ds4_gpu_stream_expert_cache_take_reusable_batch(
p->n_loads,
layer,
selected_ids,
n_selected,
gate_expert_bytes,
down_expert_bytes,
batch_reuse);
if (load_timing) {
ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(
ds4_gpu_now_ms() - reuse_t0);
}
}
for (uint32_t load_i = 0; load_i < p->n_loads; load_i++) {
const uint32_t slot = p->load_slots[load_i];
const uint32_t expert = (uint32_t)p->selected_ids[slot];
const int force_reuse =
cache_budget != 0 && reserved_entries >= cache_budget;
ds4_gpu_stream_expert_readahead_range(p->gate_abs_offsets[slot],
gate_expert_bytes);
ds4_gpu_stream_expert_readahead_range(p->up_abs_offsets[slot],
gate_expert_bytes);
ds4_gpu_stream_expert_readahead_range(p->down_abs_offsets[slot],
down_expert_bytes);
if (load_i < batch_reuse_count &&
batch_reuse[load_i].gate_buffer &&
batch_reuse[load_i].up_buffer &&
batch_reuse[load_i].down_buffer) {
p->gate_bufs[load_i] = batch_reuse[load_i].gate_buffer;
p->up_bufs[load_i] = batch_reuse[load_i].up_buffer;
p->down_bufs[load_i] = batch_reuse[load_i].down_buffer;
p->gate_inners[load_i] = batch_reuse[load_i].gate_inner;
p->up_inners[load_i] = batch_reuse[load_i].up_inner;
p->down_inners[load_i] = batch_reuse[load_i].down_inner;
} else {
const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
const int prepared =
ds4_gpu_stream_expert_cache_prepare_load_buffers(layer,
expert,
layer,
selected_ids,
n_selected,
gate_expert_bytes,
down_expert_bytes,
force_reuse,
&p->gate_bufs[load_i],
&p->up_bufs[load_i],
&p->down_bufs[load_i],
&p->gate_inners[load_i],
&p->up_inners[load_i],
&p->down_inners[load_i]);
if (load_timing) {
ds4_gpu_stream_expert_timing_note_prepare_buffer(
ds4_gpu_now_ms() - buffer_t0);
}
if (!prepared) {
ds4_gpu_stream_expert_pending_load_release_buffers(p);
return 0;
}
}
if (!force_reuse && reserved_entries < UINT32_MAX) {
reserved_entries++;
}
uint8_t *gate_dst = (uint8_t *)[p->gate_bufs[load_i] contents] +
p->gate_inners[load_i];
uint8_t *up_dst = (uint8_t *)[p->up_bufs[load_i] contents] +
p->up_inners[load_i];
uint8_t *down_dst = (uint8_t *)[p->down_bufs[load_i] contents] +
p->down_inners[load_i];
if (!gate_dst || !up_dst || !down_dst) {
ds4_gpu_stream_expert_pending_load_release_buffers(p);
return 0;
}
const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = p->gate_abs_offsets[slot],
.len = gate_expert_bytes,
.dst = gate_dst,
};
p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = p->up_abs_offsets[slot],
.len = gate_expert_bytes,
.dst = up_dst,
};
p->tasks[p->n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = p->down_abs_offsets[slot],
.len = down_expert_bytes,
.dst = down_dst,
};
if (load_timing) {
ds4_gpu_stream_expert_timing_note_prepare_task(
1,
ds4_gpu_now_ms() - task_t0);
}
}
const uint32_t n_workers =
ds4_gpu_stream_expert_pread_thread_count(p->n_tasks);
p->start_ms = ds4_gpu_now_ms();
if (load_timing) {
p->prepare_ms = p->start_ms - load_prepare_t0;
}
if (ds4_gpu_stream_expert_pread_pool_begin(p->tasks,
p->n_tasks,
n_workers)) {
p->active = 1;
if (ds4_gpu_stream_expert_pending_load_profile_enabled()) {
fprintf(stderr,
"ds4: Metal streaming expert early-load begin layer=%u experts=%u tensors=%u threads=%u\n",
layer,
p->n_loads,
p->n_tasks,
n_workers);
}
return 1;
}
uint64_t read_bytes = 0;
double read_ms = 0.0;
if (!ds4_gpu_stream_expert_pread_tasks(p->tasks,
p->n_tasks,
&read_bytes,
&read_ms)) {
ds4_gpu_stream_expert_pending_load_release_buffers(p);
return 0;
}
(void)read_bytes;
if (!ds4_gpu_stream_expert_pending_load_install(p, NULL, read_ms)) {
ds4_gpu_stream_expert_pending_load_release_buffers(p);
return 0;
}
ds4_gpu_stream_expert_pending_load_release_buffers(p);
p->n_tasks = 0;
p->n_loads = 0;
p->prepare_ms = 0.0;
return 1;
}
static int ds4_gpu_stream_expert_cache_load_selected_missing_with_source(
const void *model_map,
uint64_t model_size,
uint32_t layer,
const int32_t *selected_ids,
uint32_t n_total_expert,
uint32_t n_selected,
const uint64_t *gate_abs_offsets,
const uint64_t *up_abs_offsets,
const uint64_t *down_abs_offsets,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
uint32_t missing_mask,
int gpu_copy_source,
ds4_gpu_stream_expert_cache_entry **entries) {
if (!g_ssd_streaming_mode ||
!model_map ||
!selected_ids ||
!gate_abs_offsets ||
!up_abs_offsets ||
!down_abs_offsets ||
!entries ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
n_selected == 0 ||
n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED ||
n_total_expert == 0 ||
n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) ||
ds4_gpu_stream_expert_cache_effective_cap(layer,
n_total_expert,
n_selected) == 0) {
return 0;
}
missing_mask &= (1u << n_selected) - 1u;
if (missing_mask == 0) return 1;
if (g_stream_expert_pending_load.active &&
n_selected <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
if (ds4_gpu_stream_expert_pending_load_matches(model_map,
model_size,
layer,
selected_ids,
n_total_expert,
n_selected,
gate_expert_bytes,
down_expert_bytes)) {
if (!ds4_gpu_stream_expert_pending_load_finish(entries)) return 0;
for (uint32_t i = 0; i < n_selected; i++) {
if ((missing_mask & (1u << i)) != 0 && entries[i]) {
missing_mask &= ~(1u << i);
}
}
if (missing_mask == 0) return 1;
} else {
ds4_gpu_stream_expert_pending_load_clear();
}
} else if (g_stream_expert_pending_load.active) {
ds4_gpu_stream_expert_pending_load_clear();
}
const int load_timing = ds4_gpu_stream_expert_timing_summary_enabled();
double load_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
double load_prepare_ms = 0.0;
double load_install_ms = 0.0;
uint32_t load_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
uint32_t source_slots[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
load_slots[i] = 0;
source_slots[i] = UINT32_MAX;
}
uint32_t n_loads = 0;
for (uint32_t i = 0; i < n_selected; i++) {
if ((missing_mask & (1u << i)) == 0) continue;
if (entries[i]) {
source_slots[i] = i;
continue;
}
if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) {
fprintf(stderr,
"ds4: Metal streaming selected missing expert id %d is outside 0..%u\n",
selected_ids[i],
n_total_expert);
return 0;
}
uint32_t source = UINT32_MAX;
for (uint32_t prev = 0; prev < i; prev++) {
if (selected_ids[prev] == selected_ids[i] &&
gate_abs_offsets[prev] == gate_abs_offsets[i] &&
up_abs_offsets[prev] == up_abs_offsets[i] &&
down_abs_offsets[prev] == down_abs_offsets[i] &&
(entries[prev] || (missing_mask & (1u << prev)) != 0)) {
source = source_slots[prev] != UINT32_MAX ?
source_slots[prev] : prev;
break;
}
}
if (source != UINT32_MAX) {
source_slots[i] = source;
continue;
}
source_slots[i] = i;
load_slots[n_loads++] = i;
}
if (n_loads == 0) {
for (uint32_t i = 0; i < n_selected; i++) {
if ((missing_mask & (1u << i)) == 0 || entries[i]) continue;
const uint32_t source = source_slots[i];
if (source >= n_selected || !entries[source]) return 0;
entries[i] = entries[source];
entries[i]->use_count++;
}
return 1;
}
__strong id<MTLBuffer> gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
__strong id<MTLBuffer> up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
__strong id<MTLBuffer> down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
gate_bufs[i] = nil;
up_bufs[i] = nil;
down_bufs[i] = nil;
gate_inners[i] = 0;
up_inners[i] = 0;
down_inners[i] = 0;
}
ds4_gpu_stream_expert_pread_task tasks[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED * 3u];
memset(tasks, 0, sizeof(tasks));
uint32_t n_tasks = 0;
const uint32_t cache_budget =
ds4_gpu_stream_expert_cache_configured_budget();
uint32_t reserved_entries = g_stream_expert_cache_entry_count;
ds4_gpu_stream_expert_reusable_buffers
batch_reuse[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
batch_reuse[i] =
(ds4_gpu_stream_expert_reusable_buffers){ nil, nil, nil, 0, 0, 0 };
}
uint32_t batch_reuse_count = 0;
if (cache_budget != 0 &&
reserved_entries >= cache_budget &&
n_loads > 1 &&
ds4_gpu_stream_expert_batch_reuse_enabled(gate_expert_bytes,
down_expert_bytes)) {
const double reuse_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
batch_reuse_count =
ds4_gpu_stream_expert_cache_take_reusable_batch(
n_loads,
layer,
selected_ids,
n_selected,
gate_expert_bytes,
down_expert_bytes,
batch_reuse);
if (load_timing) {
ds4_gpu_stream_expert_timing_note_prepare_batch_reuse(
ds4_gpu_now_ms() - reuse_t0);
}
}
for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
const uint32_t slot = load_slots[load_i];
const uint32_t expert = (uint32_t)selected_ids[slot];
const int force_reuse =
cache_budget != 0 && reserved_entries >= cache_budget;
if (!gpu_copy_source) {
ds4_gpu_stream_expert_readahead_range(gate_abs_offsets[slot],
gate_expert_bytes);
ds4_gpu_stream_expert_readahead_range(up_abs_offsets[slot],
gate_expert_bytes);
ds4_gpu_stream_expert_readahead_range(down_abs_offsets[slot],
down_expert_bytes);
}
if (load_i < batch_reuse_count &&
batch_reuse[load_i].gate_buffer &&
batch_reuse[load_i].up_buffer &&
batch_reuse[load_i].down_buffer) {
gate_bufs[load_i] = batch_reuse[load_i].gate_buffer;
up_bufs[load_i] = batch_reuse[load_i].up_buffer;
down_bufs[load_i] = batch_reuse[load_i].down_buffer;
gate_inners[load_i] = batch_reuse[load_i].gate_inner;
up_inners[load_i] = batch_reuse[load_i].up_inner;
down_inners[load_i] = batch_reuse[load_i].down_inner;
} else {
const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
const int prepared =
ds4_gpu_stream_expert_cache_prepare_load_buffers(layer,
expert,
layer,
selected_ids,
n_selected,
gate_expert_bytes,
down_expert_bytes,
force_reuse,
&gate_bufs[load_i],
&up_bufs[load_i],
&down_bufs[load_i],
&gate_inners[load_i],
&up_inners[load_i],
&down_inners[load_i]);
if (load_timing) {
ds4_gpu_stream_expert_timing_note_prepare_buffer(
ds4_gpu_now_ms() - buffer_t0);
}
if (!prepared) {
return 0;
}
}
if (!force_reuse && reserved_entries < UINT32_MAX) {
reserved_entries++;
}
if (!gate_bufs[load_i] || !up_bufs[load_i] || !down_bufs[load_i]) {
return 0;
}
uint8_t *gate_dst = (uint8_t *)[gate_bufs[load_i] contents] +
gate_inners[load_i];
uint8_t *up_dst = (uint8_t *)[up_bufs[load_i] contents] +
up_inners[load_i];
uint8_t *down_dst = (uint8_t *)[down_bufs[load_i] contents] +
down_inners[load_i];
if (!gate_dst || !up_dst || !down_dst) return 0;
if (gpu_copy_source &&
(gate_abs_offsets[slot] > model_size ||
gate_expert_bytes > model_size - gate_abs_offsets[slot] ||
up_abs_offsets[slot] > model_size ||
gate_expert_bytes > model_size - up_abs_offsets[slot] ||
down_abs_offsets[slot] > model_size ||
down_expert_bytes > model_size - down_abs_offsets[slot])) {
fprintf(stderr,
"ds4: Metal streaming mapped expert source is outside the model mapping\n");
return 0;
}
const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = gate_abs_offsets[slot],
.len = gate_expert_bytes,
.dst = gate_dst,
};
tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = up_abs_offsets[slot],
.len = gate_expert_bytes,
.dst = up_dst,
};
tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = down_abs_offsets[slot],
.len = down_expert_bytes,
.dst = down_dst,
};
if (load_timing) {
ds4_gpu_stream_expert_timing_note_prepare_task(
1,
ds4_gpu_now_ms() - task_t0);
}
}
if (load_timing) {
const double now_ms = ds4_gpu_now_ms();
load_prepare_ms = now_ms - load_t0;
load_t0 = now_ms;
}
uint64_t read_bytes = 0;
double read_ms = 0.0;
int ok = 1;
if (gpu_copy_source) {
if (!g_batch_cb ||
gate_expert_bytes > (uint64_t)NSUIntegerMax ||
down_expert_bytes > (uint64_t)NSUIntegerMax) {
return 0;
}
__strong id<MTLBuffer> gate_srcs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
__strong id<MTLBuffer> up_srcs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
__strong id<MTLBuffer> down_srcs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger gate_src_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger up_src_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
NSUInteger down_src_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
const uint32_t slot = load_slots[load_i];
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
uint64_t down_inner = 0;
gate_srcs[load_i] = ds4_gpu_wrap_model_range(model_map,
model_size,
gate_abs_offsets[slot],
gate_expert_bytes,
&gate_inner);
up_srcs[load_i] = ds4_gpu_wrap_model_range(model_map,
model_size,
up_abs_offsets[slot],
gate_expert_bytes,
&up_inner);
down_srcs[load_i] = ds4_gpu_wrap_model_range(model_map,
model_size,
down_abs_offsets[slot],
down_expert_bytes,
&down_inner);
if (!gate_srcs[load_i] || !up_srcs[load_i] || !down_srcs[load_i] ||
gate_inner > (uint64_t)NSUIntegerMax ||
up_inner > (uint64_t)NSUIntegerMax ||
down_inner > (uint64_t)NSUIntegerMax) {
return 0;
}
gate_src_inners[load_i] = (NSUInteger)gate_inner;
up_src_inners[load_i] = (NSUInteger)up_inner;
down_src_inners[load_i] = (NSUInteger)down_inner;
}
ds4_gpu_close_batch_encoder();
id<MTLBlitCommandEncoder> blit = ds4_gpu_blit_encoder(g_batch_cb, "expert_cache_fill", n_loads * 3u);
if (!blit) return 0;
for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
[blit copyFromBuffer:gate_srcs[load_i]
sourceOffset:gate_src_inners[load_i]
toBuffer:gate_bufs[load_i]
destinationOffset:gate_inners[load_i]
size:(NSUInteger)gate_expert_bytes];
[blit copyFromBuffer:up_srcs[load_i]
sourceOffset:up_src_inners[load_i]
toBuffer:up_bufs[load_i]
destinationOffset:up_inners[load_i]
size:(NSUInteger)gate_expert_bytes];
[blit copyFromBuffer:down_srcs[load_i]
sourceOffset:down_src_inners[load_i]
toBuffer:down_bufs[load_i]
destinationOffset:down_inners[load_i]
size:(NSUInteger)down_expert_bytes];
const uint64_t bytes =
gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2u ?
UINT64_MAX : gate_expert_bytes * 2u + down_expert_bytes;
read_bytes = read_bytes > UINT64_MAX - bytes ?
UINT64_MAX : read_bytes + bytes;
}
[blit endEncoding];
g_batch_has_work = YES;
} else {
ok = ds4_gpu_stream_expert_pread_tasks(tasks,
n_tasks,
&read_bytes,
&read_ms);
}
if (!ok) return 0;
if (load_timing) {
load_t0 = ds4_gpu_now_ms();
}
if (!gpu_copy_source) {
ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms);
}
if (gpu_copy_source &&
getenv("DS4_METAL_STREAMING_PREFILL_CACHE_SEED_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal streaming expert GPU cache copy layer=%u "
"experts=%u bytes=%.2f GiB\n",
layer,
n_loads,
ds4_gpu_gib(read_bytes));
} else if (!gpu_copy_source &&
getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal streaming expert parallel pread layer=%u experts=%u tensors=%u "
"threads=%u bytes=%.2f GiB wall=%.3f ms\n",
layer,
n_loads,
n_tasks,
ds4_gpu_stream_expert_pread_thread_count(n_tasks),
ds4_gpu_gib(read_bytes),
read_ms);
}
for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
const uint32_t slot = load_slots[load_i];
const uint32_t expert = (uint32_t)selected_ids[slot];
ds4_gpu_stream_expert_cache_entry *entry =
ds4_gpu_stream_expert_cache_install_loaded(model_map,
model_size,
layer,
expert,
gate_abs_offsets[slot],
up_abs_offsets[slot],
down_abs_offsets[slot],
gate_expert_bytes,
down_expert_bytes,
gate_bufs[load_i],
up_bufs[load_i],
down_bufs[load_i],
gate_inners[load_i],
up_inners[load_i],
down_inners[load_i]);
if (!entry) return 0;
entries[slot] = entry;
}
for (uint32_t i = 0; i < n_selected; i++) {
if ((missing_mask & (1u << i)) == 0 || entries[i]) continue;
const uint32_t source = source_slots[i];
if (source >= n_selected || !entries[source]) return 0;
entries[i] = entries[source];
entries[i]->use_count++;
}
for (uint32_t i = 0; i < n_selected; i++) {
if ((missing_mask & (1u << i)) != 0 && !entries[i]) return 0;
}
if (load_timing && !gpu_copy_source) {
load_install_ms = ds4_gpu_now_ms() - load_t0;
ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms,
read_ms,
load_install_ms);
}
return 1;
}
static int ds4_gpu_stream_expert_cache_load_selected_missing(
const void *model_map,
uint64_t model_size,
uint32_t layer,
const int32_t *selected_ids,
uint32_t n_total_expert,
uint32_t n_selected,
const uint64_t *gate_abs_offsets,
const uint64_t *up_abs_offsets,
const uint64_t *down_abs_offsets,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
uint32_t missing_mask,
ds4_gpu_stream_expert_cache_entry **entries) {
return ds4_gpu_stream_expert_cache_load_selected_missing_with_source(
model_map,
model_size,
layer,
selected_ids,
n_total_expert,
n_selected,
gate_abs_offsets,
up_abs_offsets,
down_abs_offsets,
gate_expert_bytes,
down_expert_bytes,
missing_mask,
0,
entries);
}
static void ds4_gpu_glm_stream_selected_prefetch_set(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_selected) {
g_glm_stream_selected_prefetch.active = 0;
if (!table || !selected_ids ||
n_selected == 0 ||
n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
return;
}
g_glm_stream_selected_prefetch.model_map = table->model_map;
g_glm_stream_selected_prefetch.model_size = table->model_size;
g_glm_stream_selected_prefetch.layer = table->layer;
g_glm_stream_selected_prefetch.n_total_expert = table->n_total_expert;
g_glm_stream_selected_prefetch.n_selected = n_selected;
g_glm_stream_selected_prefetch.gate_offset = table->gate_offset;
g_glm_stream_selected_prefetch.up_offset = table->up_offset;
g_glm_stream_selected_prefetch.down_offset = table->down_offset;
g_glm_stream_selected_prefetch.gate_expert_bytes = table->gate_expert_bytes;
g_glm_stream_selected_prefetch.down_expert_bytes = table->down_expert_bytes;
for (uint32_t i = 0; i < n_selected; i++) {
g_glm_stream_selected_prefetch.selected_ids[i] = selected_ids[i];
}
g_glm_stream_selected_prefetch.active = 1;
}
static int ds4_gpu_glm_stream_selected_prefetch_take(
const void *model_map,
uint64_t model_size,
uint32_t layer,
uint32_t n_total_expert,
uint32_t n_selected,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
int32_t *selected_ids_out) {
if (!selected_ids_out || !g_glm_stream_selected_prefetch.active) return 0;
ds4_gpu_glm_stream_selected_prefetch *p = &g_glm_stream_selected_prefetch;
if (p->model_map != model_map ||
p->model_size != model_size ||
p->layer != layer ||
p->n_total_expert != n_total_expert ||
p->n_selected != n_selected ||
p->gate_offset != gate_offset ||
p->up_offset != up_offset ||
p->down_offset != down_offset ||
p->gate_expert_bytes != gate_expert_bytes ||
p->down_expert_bytes != down_expert_bytes) {
return 0;
}
for (uint32_t i = 0; i < n_selected; i++) {
selected_ids_out[i] = p->selected_ids[i];
}
p->active = 0;
return 1;
}
int ds4_gpu_glm_stream_expert_cache_begin_selected_load_tensor(
const ds4_gpu_stream_expert_table *table,
const ds4_gpu_tensor *selected,
uint32_t n_selected) {
g_glm_stream_selected_prefetch.active = 0;
if (!g_ssd_streaming_mode ||
getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_EARLY_LOAD") != NULL) {
return 1;
}
if (!table || !selected ||
n_selected == 0 ||
n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED) {
return 1;
}
const int timing = ds4_gpu_stream_expert_timing_summary_enabled();
double t0 = timing ? ds4_gpu_now_ms() : 0.0;
double selected_sync_ms = 0.0;
double selected_copy_ms = 0.0;
double selected_bind_ms = 0.0;
const int had_batch = g_batch_cb != nil;
if (had_batch) {
if (ds4_gpu_end_commands() == 0) return 0;
if (timing) {
selected_sync_ms = ds4_gpu_now_ms() - t0;
t0 = ds4_gpu_now_ms();
}
}
int32_t selected_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED];
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED; i++) {
selected_ids[i] = -1;
}
int ok = ds4_gpu_tensor_read(selected,
0,
selected_ids,
(uint64_t)n_selected * sizeof(selected_ids[0])) != 0;
if (timing) {
selected_copy_ms = ds4_gpu_now_ms() - t0;
t0 = ds4_gpu_now_ms();
}
if (ok) {
ok = ds4_gpu_stream_expert_cache_begin_selected_load(table,
selected_ids,
n_selected) != 0;
}
if (timing) {
selected_bind_ms = ds4_gpu_now_ms() - t0;
ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms,
selected_copy_ms,
selected_bind_ms);
}
if (ok) {
ds4_gpu_glm_stream_selected_prefetch_set(table,
selected_ids,
n_selected);
}
if (had_batch && ds4_gpu_begin_commands() == 0) ok = 0;
return ok;
}
static void ds4_gpu_stream_expert_cache_clear_layer(uint32_t layer) {
if (layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER) return;
for (uint32_t expert = 0;
expert < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT;
expert++) {
ds4_gpu_stream_expert_cache_clear_entry(layer, expert, 0);
}
g_stream_expert_cache_layer_count[layer] = 0;
}
static int ds4_gpu_stream_expert_cache_prepare_selected_batch(
const void *model_map,
uint64_t model_size,
uint32_t layer,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_total_expert,
uint32_t n_selected,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint64_t gate_expert_bytes,
uint64_t down_expert_bytes,
id<MTLBuffer> *gate_addrs,
id<MTLBuffer> *up_addrs,
id<MTLBuffer> *down_addrs,
ds4_gpu_stream_expert_cache_entry **resources,
uint32_t *n_resources,
uint32_t *unique_out,
id<MTLBuffer> *overflow_gate,
id<MTLBuffer> *overflow_up,
id<MTLBuffer> *overflow_down) {
if (overflow_gate) *overflow_gate = nil;
if (overflow_up) *overflow_up = nil;
if (overflow_down) *overflow_down = nil;
if (!g_ssd_streaming_mode ||
!model_map ||
!selected ||
!gate_addrs ||
!up_addrs ||
!down_addrs ||
!resources ||
!n_resources ||
!unique_out ||
!overflow_gate ||
!overflow_up ||
!overflow_down ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
n_tokens == 0 ||
n_selected == 0 ||
n_selected > DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED ||
n_total_expert == 0 ||
n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
gate_expert_bytes == 0 ||
down_expert_bytes == 0) {
return 0;
}
if (n_tokens > UINT32_MAX / n_selected) return 0;
const uint64_t n_ids = (uint64_t)n_tokens * n_selected;
if (n_ids > SIZE_MAX / sizeof(int32_t)) return 0;
int32_t *ids = malloc((size_t)n_ids * sizeof(ids[0]));
if (!ids) return 0;
const bool profile =
getenv("DS4_METAL_STREAMING_PREFILL_BATCH_SELECTED_ADDR_PROFILE") != NULL;
const double t0 = profile ? ds4_gpu_now_ms() : 0.0;
int ok = ds4_gpu_tensor_read(selected,
0,
ids,
n_ids * sizeof(ids[0]));
const double t_read = profile ? ds4_gpu_now_ms() : 0.0;
bool seen[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { false };
uint32_t frequency[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { 0 };
int32_t unique_ids[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
uint32_t unique_count = 0;
*n_resources = 0;
*unique_out = 0;
if (ok) {
if (!ds4_gpu_stream_expert_cache_ensure_addr_buffers(layer)) {
ok = 0;
}
}
if (ok) {
for (uint64_t i = 0; i < n_ids; i++) {
const int32_t selected_id = ids[i];
if (selected_id < 0 || (uint32_t)selected_id >= n_total_expert) {
fprintf(stderr,
"ds4: Metal streaming batch selected expert id %d is outside 0..%u at layer %u\n",
selected_id,
n_total_expert,
layer);
ok = 0;
break;
}
frequency[(uint32_t)selected_id]++;
if (!seen[(uint32_t)selected_id]) {
seen[(uint32_t)selected_id] = true;
unique_ids[unique_count++] = selected_id;
}
}
}
if (ok) {
ds4_gpu_stream_expert_cache_note_frequency_hotness(layer,
frequency,
n_total_expert);
}
/*
* When the layer's unique selected set does not fit the cache budget, the
* extra experts are addressed straight into whole-tensor mapped model views
* instead of falling back to a different MoE kernel path. The address-table
* kernels read identical expert bytes either way, so the cache/view split
* does not change the computed logits.
*/
uint32_t view_served = 0;
uint64_t overflow_gate_inner = 0;
uint64_t overflow_up_inner = 0;
uint64_t overflow_down_inner = 0;
uint64_t unique_gate_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
uint64_t unique_up_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
uint64_t unique_down_offsets[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
ds4_gpu_stream_expert_cache_entry
*unique_entries[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
uint32_t load_unique[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
__strong id<MTLBuffer>
gate_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
__strong id<MTLBuffer>
up_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
__strong id<MTLBuffer>
down_bufs[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
NSUInteger gate_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
NSUInteger up_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
NSUInteger down_inners[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT];
for (uint32_t i = 0; i < DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT; i++) {
unique_gate_offsets[i] = 0;
unique_up_offsets[i] = 0;
unique_down_offsets[i] = 0;
unique_entries[i] = NULL;
load_unique[i] = 0;
gate_bufs[i] = nil;
up_bufs[i] = nil;
down_bufs[i] = nil;
gate_inners[i] = 0;
up_inners[i] = 0;
down_inners[i] = 0;
}
ds4_gpu_stream_expert_pread_task *tasks = NULL;
uint32_t n_loads = 0;
uint32_t n_tasks = 0;
double load_prepare_ms = 0.0;
double load_install_ms = 0.0;
double load_timing_t0 = ds4_gpu_stream_expert_timing_summary_enabled() ?
ds4_gpu_now_ms() : 0.0;
const int load_timing = load_timing_t0 != 0.0;
if (ok && unique_count != 0) {
tasks = calloc((size_t)unique_count * 3u, sizeof(tasks[0]));
if (!tasks) ok = 0;
}
if (ok) {
const uint32_t cache_budget =
ds4_gpu_stream_expert_cache_configured_budget();
uint32_t reserved_entries = g_stream_expert_cache_entry_count;
for (uint32_t u = 0; u < unique_count; u++) {
const uint32_t expert = (uint32_t)unique_ids[u];
if ((uint64_t)expert > UINT64_MAX / gate_expert_bytes ||
(uint64_t)expert > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n");
ok = 0;
break;
}
const uint64_t gate_rel = (uint64_t)expert * gate_expert_bytes;
const uint64_t down_rel = (uint64_t)expert * down_expert_bytes;
if (gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal streaming batch selected expert offset overflow\n");
ok = 0;
break;
}
unique_gate_offsets[u] = gate_offset + gate_rel;
unique_up_offsets[u] = up_offset + gate_rel;
unique_down_offsets[u] = down_offset + down_rel;
ds4_gpu_stream_expert_cache_entry *entry =
ds4_gpu_stream_expert_cache_peek(model_map,
model_size,
layer,
expert,
n_total_expert,
n_selected,
unique_gate_offsets[u],
unique_up_offsets[u],
unique_down_offsets[u],
gate_expert_bytes,
down_expert_bytes);
if (entry) {
unique_entries[u] = entry;
continue;
}
const int force_reuse =
cache_budget != 0 && reserved_entries >= cache_budget;
/* The worker pool reads this entire batch below. Serial read-ahead
* here waits for the same pages before parallel I/O can begin. */
const double buffer_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
const int prepared =
ds4_gpu_stream_expert_cache_prepare_load_buffers(layer,
expert,
layer,
unique_ids,
unique_count,
gate_expert_bytes,
down_expert_bytes,
force_reuse,
&gate_bufs[n_loads],
&up_bufs[n_loads],
&down_bufs[n_loads],
&gate_inners[n_loads],
&up_inners[n_loads],
&down_inners[n_loads]);
if (load_timing) {
ds4_gpu_stream_expert_timing_note_prepare_buffer(
ds4_gpu_now_ms() - buffer_t0);
}
if (!prepared) {
ok = 0;
break;
}
if (!force_reuse && reserved_entries < UINT32_MAX) {
reserved_entries++;
}
if (!gate_bufs[n_loads] ||
!up_bufs[n_loads] ||
!down_bufs[n_loads]) {
ok = 0;
break;
}
uint8_t *gate_dst = (uint8_t *)[gate_bufs[n_loads] contents] +
gate_inners[n_loads];
uint8_t *up_dst = (uint8_t *)[up_bufs[n_loads] contents] +
up_inners[n_loads];
uint8_t *down_dst = (uint8_t *)[down_bufs[n_loads] contents] +
down_inners[n_loads];
if (!gate_dst || !up_dst || !down_dst) {
ok = 0;
break;
}
load_unique[n_loads] = u;
const double task_t0 = load_timing ? ds4_gpu_now_ms() : 0.0;
tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = unique_gate_offsets[u],
.len = gate_expert_bytes,
.dst = gate_dst,
};
tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = unique_up_offsets[u],
.len = gate_expert_bytes,
.dst = up_dst,
};
tasks[n_tasks++] = (ds4_gpu_stream_expert_pread_task) {
.offset = unique_down_offsets[u],
.len = down_expert_bytes,
.dst = down_dst,
};
if (load_timing) {
ds4_gpu_stream_expert_timing_note_prepare_task(
1,
ds4_gpu_now_ms() - task_t0);
}
n_loads++;
}
}
if (ok && n_loads != 0) {
if (load_timing_t0 != 0.0) {
const double now_ms = ds4_gpu_now_ms();
load_prepare_ms = now_ms - load_timing_t0;
load_timing_t0 = now_ms;
}
uint64_t read_bytes = 0;
double read_ms = 0.0;
ok = ds4_gpu_stream_expert_pread_tasks(tasks,
n_tasks,
&read_bytes,
&read_ms);
if (ok) {
ds4_gpu_stream_expert_cache_note_pread(layer, read_bytes, read_ms);
}
if (load_timing_t0 != 0.0) {
load_timing_t0 = ds4_gpu_now_ms();
}
if (ok && getenv("DS4_METAL_STREAMING_EXPERT_PREAD_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal streaming batch expert parallel pread layer=%u experts=%u tensors=%u "
"threads=%u bytes=%.2f GiB wall=%.3f ms\n",
layer,
n_loads,
n_tasks,
ds4_gpu_stream_expert_pread_thread_count(n_tasks),
ds4_gpu_gib(read_bytes),
read_ms);
}
if (ok) {
for (uint32_t load_i = 0; load_i < n_loads; load_i++) {
const uint32_t u = load_unique[load_i];
const uint32_t expert = (uint32_t)unique_ids[u];
ds4_gpu_stream_expert_cache_entry *entry =
ds4_gpu_stream_expert_cache_install_loaded(model_map,
model_size,
layer,
expert,
unique_gate_offsets[u],
unique_up_offsets[u],
unique_down_offsets[u],
gate_expert_bytes,
down_expert_bytes,
gate_bufs[load_i],
up_bufs[load_i],
down_bufs[load_i],
gate_inners[load_i],
up_inners[load_i],
down_inners[load_i]);
if (!entry) {
ok = 0;
break;
}
unique_entries[u] = entry;
}
}
if (load_timing_t0 != 0.0) {
load_install_ms = ds4_gpu_now_ms() - load_timing_t0;
ds4_gpu_stream_expert_timing_note_load_detail(load_prepare_ms,
read_ms,
load_install_ms);
}
}
if (tasks) free(tasks);
if (ok) {
for (uint32_t u = 0; u < unique_count; u++) {
ds4_gpu_stream_expert_cache_entry *entry = unique_entries[u];
const uint32_t expert = (uint32_t)unique_ids[u];
if (!entry) {
const uint64_t gate_rel = unique_gate_offsets[u] - gate_offset;
const uint64_t down_rel = unique_down_offsets[u] - down_offset;
if (!*overflow_gate) {
const uint64_t gate_tensor_bytes =
(uint64_t)n_total_expert * gate_expert_bytes;
const uint64_t down_tensor_bytes =
(uint64_t)n_total_expert * down_expert_bytes;
uint64_t gate_view_inner = 0;
uint64_t up_view_inner = 0;
uint64_t down_view_inner = 0;
id<MTLBuffer> gv = ds4_gpu_wrap_model_range(model_map,
model_size,
gate_offset,
gate_tensor_bytes,
&gate_view_inner);
id<MTLBuffer> uv = ds4_gpu_wrap_model_range(model_map,
model_size,
up_offset,
gate_tensor_bytes,
&up_view_inner);
id<MTLBuffer> dv = ds4_gpu_wrap_model_range(model_map,
model_size,
down_offset,
down_tensor_bytes,
&down_view_inner);
if (!gv || !uv || !dv) {
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr "
"failed to map overflow expert views at layer %u\n",
layer);
ok = 0;
break;
}
*overflow_gate = gv;
*overflow_up = uv;
*overflow_down = dv;
overflow_gate_inner = gate_view_inner;
overflow_up_inner = up_view_inner;
overflow_down_inner = down_view_inner;
}
if (!ds4_gpu_stream_expert_cache_set_addr_slot(
layer,
expert,
*overflow_gate,
(NSUInteger)(overflow_gate_inner + gate_rel),
*overflow_up,
(NSUInteger)(overflow_up_inner + gate_rel),
*overflow_down,
(NSUInteger)(overflow_down_inner + down_rel))) {
ok = 0;
break;
}
view_served++;
continue;
}
const uint32_t extra_uses =
frequency[expert] > 0 ? frequency[expert] - 1u : 0;
if (extra_uses != 0) {
if (entry->use_count > UINT64_MAX - extra_uses) {
entry->use_count = UINT64_MAX;
} else {
entry->use_count += extra_uses;
}
if (g_stream_expert_cache_hits > UINT64_MAX - extra_uses) {
g_stream_expert_cache_hits = UINT64_MAX;
} else {
g_stream_expert_cache_hits += extra_uses;
}
if (g_stream_expert_cache_layer_hits[layer] >
UINT64_MAX - extra_uses) {
g_stream_expert_cache_layer_hits[layer] = UINT64_MAX;
} else {
g_stream_expert_cache_layer_hits[layer] += extra_uses;
}
}
resources[*n_resources] = entry;
(*n_resources)++;
}
}
free(ids);
if (!ok || (*n_resources == 0 && view_served == 0)) {
ds4_gpu_stream_expert_cache_clear_layer(layer);
return 0;
}
if (!ds4_gpu_stream_expert_cache_addr_buffers(layer,
gate_addrs,
up_addrs,
down_addrs)) {
ds4_gpu_stream_expert_cache_clear_layer(layer);
return 0;
}
*unique_out = *n_resources + view_served;
if (view_served != 0 &&
ds4_gpu_stream_expert_timing_summary_enabled()) {
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr layer=%u "
"served %u/%u unique experts via mapped views (cache budget %u)\n",
layer,
view_served,
unique_count,
ds4_gpu_stream_expert_cache_configured_budget());
}
if (profile) {
const double t_done = ds4_gpu_now_ms();
uint64_t per_expert_bytes = UINT64_MAX;
uint64_t logical_bytes = UINT64_MAX;
if (gate_expert_bytes <= (UINT64_MAX - down_expert_bytes) / 2ull) {
per_expert_bytes = gate_expert_bytes * 2ull + down_expert_bytes;
if (per_expert_bytes == 0 ||
*n_resources <= UINT64_MAX / per_expert_bytes) {
logical_bytes = (uint64_t)(*n_resources) * per_expert_bytes;
}
}
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr layer=%u "
"tokens=%u unique=%u read=%.3f ms wrap=%.3f ms bytes=%.2f GiB\n",
layer,
n_tokens,
*n_resources,
t_read - t0,
t_done - t_read,
ds4_gpu_gib(logical_bytes));
}
return 1;
}
int ds4_gpu_stream_expert_cache_seed_selected(
const ds4_gpu_stream_expert_table *table,
const int32_t *selected_ids,
uint32_t n_selected) {
if (!g_ssd_streaming_mode) return 1;
if (!table) return 0;
const void *model_map = table->model_map;
const uint64_t model_size = table->model_size;
const uint32_t layer = table->layer;
const uint32_t n_total_expert = table->n_total_expert;
const uint64_t gate_offset = table->gate_offset;
const uint64_t up_offset = table->up_offset;
const uint64_t down_offset = table->down_offset;
const uint64_t gate_expert_bytes = table->gate_expert_bytes;
const uint64_t down_expert_bytes = table->down_expert_bytes;
if (!model_map || !selected_ids ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
n_selected == 0 ||
n_selected > 6 ||
n_total_expert == 0 ||
n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) ||
ds4_gpu_stream_expert_cache_effective_cap(layer,
n_total_expert,
n_selected) == 0) {
return 1;
}
if (!g_initialized && !ds4_gpu_init()) return 0;
ds4_gpu_stream_expert_cache_note_selected_hotness(layer,
selected_ids,
n_selected);
for (uint32_t i = 0; i < n_selected; i++) {
if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) {
fprintf(stderr,
"ds4: Metal prefill expert-cache seed selected expert id %d is outside 0..%u\n",
selected_ids[i],
n_total_expert);
return 0;
}
const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i];
if (expert_id > UINT64_MAX / gate_expert_bytes ||
expert_id > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n");
return 0;
}
const uint64_t gate_rel = expert_id * gate_expert_bytes;
const uint64_t down_rel = expert_id * down_expert_bytes;
if (gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal prefill expert-cache seed offset overflow\n");
return 0;
}
if (!ds4_gpu_stream_expert_cache_get(model_map,
model_size,
layer,
(uint32_t)selected_ids[i],
n_total_expert,
n_selected,
gate_offset + gate_rel,
up_offset + gate_rel,
down_offset + down_rel,
gate_expert_bytes,
down_expert_bytes)) {
return 0;
}
}
ds4_gpu_stream_expert_cache_prune_layer(layer,
n_total_expert,
n_selected,
selected_ids,
n_selected);
ds4_gpu_stream_expert_cache_prune_global(layer,
selected_ids,
n_selected);
return 1;
}
static int ds4_gpu_stream_expert_cache_seed_experts_impl(
const ds4_gpu_stream_expert_table *table,
const int32_t *expert_ids,
const uint32_t *expert_priorities,
uint32_t n_experts,
int gpu_copy_source) {
if (!g_ssd_streaming_mode) return 1;
if (!table) return 0;
const void *model_map = table->model_map;
const uint64_t model_size = table->model_size;
const uint32_t layer = table->layer;
const uint32_t n_total_expert = table->n_total_expert;
const uint64_t gate_offset = table->gate_offset;
const uint64_t up_offset = table->up_offset;
const uint64_t down_offset = table->down_offset;
const uint64_t gate_expert_bytes = table->gate_expert_bytes;
const uint64_t down_expert_bytes = table->down_expert_bytes;
if (!model_map || !expert_ids ||
layer >= DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER ||
n_experts == 0 ||
n_total_expert == 0 ||
n_total_expert > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT ||
!ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) ||
ds4_gpu_stream_expert_cache_effective_cap(layer,
n_total_expert,
1) == 0) {
return 1;
}
if (!g_initialized && !ds4_gpu_init()) return 0;
ds4_gpu_stream_expert_cache_maybe_decay_route_hotness();
uint32_t remaining = n_experts;
while (remaining != 0) {
const uint32_t batch =
remaining > 6u ? 6u : remaining;
remaining -= batch;
int32_t selected_ids[6] = { -1, -1, -1, -1, -1, -1 };
uint64_t gate_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 };
uint64_t up_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 };
uint64_t down_abs_offsets[6] = { 0, 0, 0, 0, 0, 0 };
ds4_gpu_stream_expert_cache_entry *entries[6] = {
NULL, NULL, NULL, NULL, NULL, NULL
};
uint32_t missing_mask = 0;
for (uint32_t i = 0; i < batch; i++) {
const int32_t expert = expert_ids[remaining + i];
const uint32_t priority =
expert_priorities ? expert_priorities[remaining + i] : 0;
if (expert < 0 || (uint32_t)expert >= n_total_expert) {
fprintf(stderr,
"ds4: Metal streaming hotlist seed expert id %d is outside 0..%u\n",
expert,
n_total_expert);
return 0;
}
const uint64_t expert_id = (uint64_t)(uint32_t)expert;
if (expert_id > UINT64_MAX / gate_expert_bytes ||
expert_id > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n");
return 0;
}
const uint64_t gate_rel = expert_id * gate_expert_bytes;
const uint64_t down_rel = expert_id * down_expert_bytes;
if (gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal streaming hotlist seed offset overflow\n");
return 0;
}
ds4_gpu_stream_expert_cache_note_route_hotness(
layer,
(uint32_t)expert,
priority != 0 ? priority : 1u);
selected_ids[i] = expert;
gate_abs_offsets[i] = gate_offset + gate_rel;
up_abs_offsets[i] = up_offset + gate_rel;
down_abs_offsets[i] = down_offset + down_rel;
entries[i] = ds4_gpu_stream_expert_cache_peek(model_map,
model_size,
layer,
(uint32_t)expert,
n_total_expert,
batch,
gate_abs_offsets[i],
up_abs_offsets[i],
down_abs_offsets[i],
gate_expert_bytes,
down_expert_bytes);
if (!entries[i]) missing_mask |= 1u << i;
if (entries[i] && priority != 0 &&
entries[i]->use_count < (uint64_t)priority) {
entries[i]->use_count = (uint64_t)priority;
}
}
if (missing_mask != 0 &&
!ds4_gpu_stream_expert_cache_load_selected_missing_with_source(
model_map,
model_size,
layer,
selected_ids,
n_total_expert,
batch,
gate_abs_offsets,
up_abs_offsets,
down_abs_offsets,
gate_expert_bytes,
down_expert_bytes,
missing_mask,
gpu_copy_source,
entries)) {
return 0;
}
if (expert_priorities) {
for (uint32_t i = 0; i < batch; i++) {
const uint32_t priority = expert_priorities[remaining + i];
if (entries[i] && priority != 0 &&
entries[i]->use_count < (uint64_t)priority) {
entries[i]->use_count = (uint64_t)priority;
}
}
}
}
const uint32_t protect_n = n_experts < 6u ? n_experts : 6u;
ds4_gpu_stream_expert_cache_prune_layer(layer,
n_total_expert,
1,
expert_ids,
protect_n);
ds4_gpu_stream_expert_cache_prune_global(layer,
expert_ids,
protect_n);
return 1;
}
int ds4_gpu_stream_expert_cache_seed_experts(
const ds4_gpu_stream_expert_table *table,
const int32_t *expert_ids,
const uint32_t *expert_priorities,
uint32_t n_experts) {
@autoreleasepool {
return ds4_gpu_stream_expert_cache_seed_experts_impl(table,
expert_ids,
expert_priorities,
n_experts,
0);
}
}
int ds4_gpu_stream_expert_cache_seed_experts_gpu_copy(
const ds4_gpu_stream_expert_table *table,
const int32_t *expert_ids,
const uint32_t *expert_priorities,
uint32_t n_experts) {
/* The C prefill caller has no pool. In particular, blit encoders retain
* destination slabs even after the cache drops its last references. */
@autoreleasepool {
return ds4_gpu_stream_expert_cache_seed_experts_impl(table,
expert_ids,
expert_priorities,
n_experts,
1);
}
}
static uint32_t ds4_gpu_q4_expert_table_group_size(uint32_t n_total_expert) {
const char *env = getenv("DS4_METAL_Q4_EXPERT_TABLE_GROUP_SIZE");
if (!env || !env[0]) return 1;
char *end = NULL;
unsigned long v = strtoul(env, &end, 10);
if (end == env || *end != '\0' || v < 2 || v > n_total_expert) {
return 1;
}
return (uint32_t)v;
}
static bool ds4_gpu_q4_table_queue_residency_requested(void) {
return getenv("DS4_METAL_Q4_TABLE_QUEUE_RESIDENCY_SET") != NULL;
}
static bool ds4_gpu_q4_table_queue_residency_available(void) {
#if TARGET_OS_OSX
if (@available(macOS 15.0, *)) {
return g_queue && [g_queue respondsToSelector:@selector(addResidencySet:)];
}
#endif
return false;
}
static bool ds4_gpu_q4_non_streaming_opt_in_enabled(void) {
return getenv("DS4_METAL_ENABLE_Q4_SELECTED_EXPERT_VIEWS") != NULL ||
getenv("DS4_METAL_ENABLE_PRO_Q4_SELECTED_EXPERT_VIEWS") != NULL ||
getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL ||
getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL ||
getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL ||
getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL;
}
static bool ds4_gpu_q4_selected_paths_allowed(void) {
if (g_ssd_streaming_mode) return true;
if (g_glm_model_mode) return false;
return ds4_gpu_q4_non_streaming_opt_in_enabled();
}
int ds4_gpu_pro_q4_expert_table_auto_available(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
/* Indirect tables retain every expert and its residency set. They cannot
* coexist with a bounded SSD cache, even when only six experts are read. */
return !g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_TABLE_AUTO") != NULL &&
getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL &&
getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL &&
ds4_gpu_q4_table_queue_residency_available();
}
static bool ds4_gpu_q4_table_queue_residency_enabled(bool auto_queue_residency) {
return (auto_queue_residency || ds4_gpu_q4_table_queue_residency_requested()) &&
ds4_gpu_q4_table_queue_residency_available();
}
static bool ds4_gpu_q4_table_model_residency_enabled(void) {
return getenv("DS4_METAL_Q4_TABLE_MODEL_RESIDENCY_SET") != NULL;
}
static bool ds4_gpu_pro_q4_expert_indirect_shape_supported(
uint32_t n_total_expert,
uint32_t n_expert,
uint64_t gate_tensor_bytes,
uint64_t down_tensor_bytes) {
const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull;
return n_total_expert == 384 &&
n_expert == 6 &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes;
}
static bool ds4_gpu_pro_q4_expert_table_auto_enabled(
uint32_t n_total_expert,
uint32_t n_expert,
uint64_t gate_tensor_bytes,
uint64_t down_tensor_bytes) {
/*
* This path lets the shader choose among exact per-expert resources.
* It is only automatic when a Metal queue residency set can make every
* indirect expert resource visible up front; otherwise the selected
* exact-slice path remains the fallback.
*/
return ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert,
n_expert,
gate_tensor_bytes,
down_tensor_bytes) &&
ds4_gpu_pro_q4_expert_table_auto_available();
}
static bool ds4_gpu_pro_q4_expert_address_auto_enabled(
uint32_t n_total_expert,
uint32_t n_expert,
uint64_t gate_tensor_bytes,
uint64_t down_tensor_bytes) {
/*
* GPU-address expert tables are useful for experiments, but they are not
* safe as an automatic path unless every indirect expert resource is made
* visible to Metal. Keep this behind an explicit opt-in while the selected
* active-slice path remains the correctness/performance baseline.
*/
return !g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") != NULL &&
ds4_gpu_q4_selected_paths_allowed() &&
ds4_gpu_pro_q4_expert_indirect_shape_supported(n_total_expert,
n_expert,
gate_tensor_bytes,
down_tensor_bytes) &&
g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil &&
getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_TABLE_AUTO") == NULL &&
getenv("DS4_METAL_DISABLE_PRO_Q4_EXPERT_ADDRESS_AUTO") == NULL &&
getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL &&
ds4_gpu_q4_table_queue_residency_available();
}
static bool ds4_gpu_q4_table_bind_anchors_enabled(void) {
return getenv("DS4_METAL_Q4_TABLE_BIND_ANCHORS") != NULL;
}
static int ds4_gpu_use_model_residency_set(id<MTLCommandBuffer> cb) {
if (!ds4_gpu_q4_table_model_residency_enabled()) return 1;
#if TARGET_OS_OSX
if (@available(macOS 15.0, *)) {
if (cb && g_model_residency_set && [cb respondsToSelector:@selector(useResidencySet:)]) {
[cb useResidencySet:g_model_residency_set];
return 1;
}
}
#endif
fprintf(stderr, "ds4: Metal Q4 table model residency set is not available\n");
return 0;
}
static int ds4_gpu_bind_q4_expert_table_anchors(
id<MTLComputeCommandEncoder> enc,
DS4MetalQ4ExpertTable *table,
NSUInteger first_index,
NSUInteger max_count) {
if (!ds4_gpu_q4_table_bind_anchors_enabled()) return 1;
if (!enc || !table || !table.expertBuffers) return 0;
const NSUInteger count = [table.expertBuffers count];
if (count > max_count) {
fprintf(stderr,
"ds4: Metal Q4 table anchor count %lu exceeds available slots %lu\n",
(unsigned long)count,
(unsigned long)max_count);
return 0;
}
for (NSUInteger i = 0; i < count; i++) {
[enc setBuffer:[table.expertBuffers objectAtIndex:i]
offset:0
atIndex:first_index + i];
}
return 1;
}
static id ds4_gpu_q4_expert_table_residency_set(NSMutableArray<id<MTLBuffer>> *buffers) {
if (!buffers || [buffers count] == 0 ||
getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL ||
getenv("DS4_METAL_Q4_TABLE_PER_TENSOR_RESIDENCY_SET") == NULL) {
return nil;
}
#if TARGET_OS_OSX
if (@available(macOS 15.0, *)) {
MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init];
desc.label = @"ds4_q4_expert_table";
desc.initialCapacity = [buffers count];
NSError *error = nil;
id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error];
if (!residency_set) {
fprintf(stderr, "ds4: Metal Q4 expert table residency set creation failed: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
for (id<MTLBuffer> buffer in buffers) {
[residency_set addAllocation:buffer];
}
[residency_set commit];
[residency_set requestResidency];
return residency_set;
}
#endif
return nil;
}
static void ds4_gpu_q4_residency_add_table(id residency_set,
DS4MetalQ4ExpertTable *table) {
#if TARGET_OS_OSX
if (!residency_set || !table) return;
if (@available(macOS 15.0, *)) {
if (table.argumentBuffer) {
[residency_set addAllocation:table.argumentBuffer];
}
if (table.addressBuffer) {
[residency_set addAllocation:table.addressBuffer];
}
for (id<MTLBuffer> buffer in table.expertBuffers) {
[residency_set addAllocation:buffer];
}
}
#else
(void)residency_set;
(void)table;
#endif
}
static id ds4_gpu_q4_expert_layer_residency_set(DS4MetalQ4ExpertTable *gate_table,
DS4MetalQ4ExpertTable *up_table,
DS4MetalQ4ExpertTable *down_table,
bool auto_queue_residency) {
const bool queue_residency =
ds4_gpu_q4_table_queue_residency_enabled(auto_queue_residency);
if (!gate_table || !up_table || !down_table ||
!g_device || !g_q4_expert_layer_residency_cache ||
(getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") == NULL &&
!queue_residency)) {
return nil;
}
#if TARGET_OS_OSX
if (@available(macOS 15.0, *)) {
NSString *key = [NSString stringWithFormat:@"%p:%p:%p",
gate_table, up_table, down_table];
DS4MetalQ4LayerResidency *cached =
[g_q4_expert_layer_residency_cache objectForKey:key];
if (cached) return cached.residencySet;
const NSUInteger capacity =
[gate_table.expertBuffers count] +
[up_table.expertBuffers count] +
[down_table.expertBuffers count] + 3u;
MTLResidencySetDescriptor *desc = [[MTLResidencySetDescriptor alloc] init];
desc.label = @"ds4_q4_expert_layer";
desc.initialCapacity = capacity;
NSError *error = nil;
id residency_set = [g_device newResidencySetWithDescriptor:desc error:&error];
if (!residency_set) {
fprintf(stderr,
"ds4: Metal Q4 expert layer residency set creation failed: %s\n",
[[error localizedDescription] UTF8String]);
return nil;
}
ds4_gpu_q4_residency_add_table(residency_set, gate_table);
ds4_gpu_q4_residency_add_table(residency_set, up_table);
ds4_gpu_q4_residency_add_table(residency_set, down_table);
[residency_set commit];
[residency_set requestResidency];
DS4MetalQ4LayerResidency *entry = [DS4MetalQ4LayerResidency new];
entry.residencySet = residency_set;
if (queue_residency) {
[g_queue addResidencySet:residency_set];
entry.addedToQueue = YES;
}
[g_q4_expert_layer_residency_cache setObject:entry forKey:key];
return residency_set;
}
#endif
return nil;
}
static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table(
const void *model_map,
uint64_t model_size,
uint64_t tensor_offset,
uint64_t expert_bytes,
uint32_t n_total_expert,
id<MTLArgumentEncoder> encoder);
static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table(
const void *model_map,
uint64_t model_size,
uint64_t tensor_offset,
uint64_t expert_bytes,
uint32_t n_total_expert);
int ds4_gpu_preload_q4_expert_tables(const void *model_map, uint64_t model_size,
uint64_t gate_offset, uint64_t up_offset, uint64_t down_offset,
uint64_t gate_expert_bytes, uint64_t down_expert_bytes,
uint32_t n_total_expert) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!model_map || model_size == 0 || gate_expert_bytes == 0 ||
down_expert_bytes == 0 || n_total_expert == 0) {
return 0;
}
if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal Q4 expert table preload byte size overflow\n");
return 0;
}
const uint64_t gate_tensor_bytes = (uint64_t)n_total_expert * gate_expert_bytes;
const uint64_t down_tensor_bytes = (uint64_t)n_total_expert * down_expert_bytes;
const bool address_auto =
ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert,
6,
gate_tensor_bytes,
down_tensor_bytes);
const bool table_auto =
ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert,
6,
gate_tensor_bytes,
down_tensor_bytes);
if (!address_auto && !table_auto) {
return 1;
}
if (address_auto) {
@autoreleasepool {
DS4MetalQ4ExpertTable *gate_table =
ds4_gpu_q4_expert_address_table(model_map,
model_size,
gate_offset,
gate_expert_bytes,
n_total_expert);
DS4MetalQ4ExpertTable *up_table =
ds4_gpu_q4_expert_address_table(model_map,
model_size,
up_offset,
gate_expert_bytes,
n_total_expert);
DS4MetalQ4ExpertTable *down_table =
ds4_gpu_q4_expert_address_table(model_map,
model_size,
down_offset,
down_expert_bytes,
n_total_expert);
if (!gate_table || !up_table || !down_table) {
return 0;
}
id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table,
up_table,
down_table,
true);
if (!residency) {
fprintf(stderr, "ds4: Metal Q4 expert address table preload failed to create queue residency set\n");
return 0;
}
}
return 1;
}
if (!g_moe_table_q4_pair_gate_encoder || !g_moe_table_q4_pair_up_encoder ||
!g_moe_table_q4_sum_down_encoder) {
fprintf(stderr, "ds4: Metal Q4 expert table preload missing argument encoders\n");
return 0;
}
@autoreleasepool {
DS4MetalQ4ExpertTable *gate_table =
ds4_gpu_q4_expert_table(model_map,
model_size,
gate_offset,
gate_expert_bytes,
n_total_expert,
g_moe_table_q4_pair_gate_encoder);
DS4MetalQ4ExpertTable *up_table =
ds4_gpu_q4_expert_table(model_map,
model_size,
up_offset,
gate_expert_bytes,
n_total_expert,
g_moe_table_q4_pair_up_encoder);
DS4MetalQ4ExpertTable *down_table =
ds4_gpu_q4_expert_table(model_map,
model_size,
down_offset,
down_expert_bytes,
n_total_expert,
g_moe_table_q4_sum_down_encoder);
if (!gate_table || !up_table || !down_table) {
return 0;
}
id residency = ds4_gpu_q4_expert_layer_residency_set(gate_table,
up_table,
down_table,
true);
if (!residency) {
fprintf(stderr, "ds4: Metal Q4 expert table preload failed to create queue residency set\n");
return 0;
}
}
return 1;
}
static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_table(
const void *model_map,
uint64_t model_size,
uint64_t tensor_offset,
uint64_t expert_bytes,
uint32_t n_total_expert,
id<MTLArgumentEncoder> encoder) {
if (g_ssd_streaming_mode || !model_map || !g_device || !g_q4_expert_table_cache || !encoder ||
model_size == 0 || expert_bytes == 0 ||
n_total_expert == 0 || n_total_expert > 384) {
return nil;
}
if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) {
fprintf(stderr, "ds4: Metal Q4 expert table byte size overflow\n");
return nil;
}
const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes;
if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) {
fprintf(stderr, "ds4: Metal Q4 expert table is outside the mapped model\n");
return nil;
}
const uint32_t table_group_size =
ds4_gpu_q4_expert_table_group_size(n_total_expert);
NSString *key = [NSString stringWithFormat:@"%p:%llu:%llu:%llu:%u:%llu:%u",
model_map,
(unsigned long long)model_size,
(unsigned long long)tensor_offset,
(unsigned long long)expert_bytes,
n_total_expert,
(unsigned long long)[encoder encodedLength],
table_group_size];
DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key];
if (cached) return cached;
id<MTLBuffer> arg_buffer =
[g_device newBufferWithLength:[encoder encodedLength]
options:MTLResourceStorageModeShared];
if (!arg_buffer) {
fprintf(stderr, "ds4: Metal Q4 expert table argument buffer allocation failed\n");
return nil;
}
arg_buffer.label = @"ds4_q4_expert_table";
[encoder setArgumentBuffer:arg_buffer offset:0];
NSMutableArray<id<MTLBuffer>> *expert_buffers =
[NSMutableArray arrayWithCapacity:table_group_size > 1 ?
(n_total_expert + table_group_size - 1u) / table_group_size :
n_total_expert];
if (!expert_buffers) return nil;
if (table_group_size > 1) {
for (uint32_t first = 0; first < n_total_expert; first += table_group_size) {
const uint32_t remaining = n_total_expert - first;
const uint32_t group_n =
remaining < table_group_size ? remaining : table_group_size;
if ((uint64_t)first > UINT64_MAX / expert_bytes ||
(uint64_t)group_n > UINT64_MAX / expert_bytes) {
fprintf(stderr, "ds4: Metal Q4 expert table group byte overflow\n");
return nil;
}
const uint64_t rel = (uint64_t)first * expert_bytes;
const uint64_t group_bytes = (uint64_t)group_n * expert_bytes;
if (rel > UINT64_MAX - tensor_offset) {
fprintf(stderr, "ds4: Metal Q4 expert table group offset overflow\n");
return nil;
}
uint64_t inner = 0;
id<MTLBuffer> group_buf =
ds4_gpu_wrap_model_range(model_map,
model_size,
tensor_offset + rel,
group_bytes,
&inner);
if (!group_buf) return nil;
for (uint32_t j = 0; j < group_n; j++) {
const uint64_t expert_inner = inner + (uint64_t)j * expert_bytes;
[encoder setBuffer:group_buf offset:(NSUInteger)expert_inner atIndex:first + j];
}
[expert_buffers addObject:group_buf];
}
} else {
for (uint32_t i = 0; i < n_total_expert; i++) {
const uint64_t rel = (uint64_t)i * expert_bytes;
if (rel > UINT64_MAX - tensor_offset) {
fprintf(stderr, "ds4: Metal Q4 expert table offset overflow\n");
return nil;
}
uint64_t inner = 0;
id<MTLBuffer> expert_buf =
ds4_gpu_wrap_model_exact_range_owned(model_map,
model_size,
tensor_offset + rel,
expert_bytes,
&inner);
if (!expert_buf) return nil;
[encoder setBuffer:expert_buf offset:(NSUInteger)inner atIndex:i];
[expert_buffers addObject:expert_buf];
}
}
if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal Q4 expert table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n",
n_total_expert,
table_group_size,
(unsigned long)[expert_buffers count],
ds4_gpu_mib(expert_bytes));
}
DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new];
table.argumentBuffer = arg_buffer;
table.expertBuffers = expert_buffers;
table.residencySet = ds4_gpu_q4_expert_table_residency_set(expert_buffers);
table.nExpert = n_total_expert;
table.expertBytes = expert_bytes;
[g_q4_expert_table_cache setObject:table forKey:key];
return table;
}
static DS4MetalQ4ExpertTable *ds4_gpu_q4_expert_address_table(
const void *model_map,
uint64_t model_size,
uint64_t tensor_offset,
uint64_t expert_bytes,
uint32_t n_total_expert) {
if (g_ssd_streaming_mode || !model_map || !g_device || !g_q4_expert_table_cache ||
model_size == 0 || expert_bytes == 0 ||
n_total_expert == 0 || n_total_expert > 384) {
return nil;
}
if ((uint64_t)n_total_expert > UINT64_MAX / expert_bytes) {
fprintf(stderr, "ds4: Metal Q4 expert address table byte size overflow\n");
return nil;
}
const uint64_t tensor_bytes = (uint64_t)n_total_expert * expert_bytes;
if (tensor_offset > model_size || tensor_bytes > model_size - tensor_offset) {
fprintf(stderr, "ds4: Metal Q4 expert address table is outside the mapped model\n");
return nil;
}
const uint32_t table_group_size =
ds4_gpu_q4_expert_table_group_size(n_total_expert);
NSString *key = [NSString stringWithFormat:@"addr:%p:%llu:%llu:%llu:%u:%u",
model_map,
(unsigned long long)model_size,
(unsigned long long)tensor_offset,
(unsigned long long)expert_bytes,
n_total_expert,
table_group_size];
DS4MetalQ4ExpertTable *cached = [g_q4_expert_table_cache objectForKey:key];
if (cached) return cached;
id<MTLBuffer> address_buffer =
[g_device newBufferWithLength:(NSUInteger)n_total_expert * sizeof(uint64_t)
options:MTLResourceStorageModeShared];
if (!address_buffer) {
fprintf(stderr, "ds4: Metal Q4 expert address table allocation failed\n");
return nil;
}
address_buffer.label = @"ds4_q4_expert_address_table";
uint64_t *addresses = (uint64_t *)[address_buffer contents];
if (!addresses) return nil;
NSMutableArray<id<MTLBuffer>> *expert_buffers =
[NSMutableArray arrayWithCapacity:table_group_size > 1 ?
(n_total_expert + table_group_size - 1u) / table_group_size :
n_total_expert];
if (!expert_buffers) return nil;
#if TARGET_OS_OSX
if (@available(macOS 13.0, *)) {
for (uint32_t first = 0; first < n_total_expert; first += table_group_size) {
const uint32_t remaining = n_total_expert - first;
const uint32_t group_n =
remaining < table_group_size ? remaining : table_group_size;
if ((uint64_t)first > UINT64_MAX / expert_bytes ||
(uint64_t)group_n > UINT64_MAX / expert_bytes) {
fprintf(stderr, "ds4: Metal Q4 expert address table group byte overflow\n");
return nil;
}
const uint64_t rel = (uint64_t)first * expert_bytes;
const uint64_t group_bytes = (uint64_t)group_n * expert_bytes;
if (rel > UINT64_MAX - tensor_offset) {
fprintf(stderr, "ds4: Metal Q4 expert address table group offset overflow\n");
return nil;
}
uint64_t inner = 0;
id<MTLBuffer> group_buf = nil;
if (table_group_size > 1) {
group_buf = ds4_gpu_wrap_model_range(model_map,
model_size,
tensor_offset + rel,
group_bytes,
&inner);
} else {
group_buf = ds4_gpu_wrap_model_exact_range_owned(model_map,
model_size,
tensor_offset + rel,
expert_bytes,
&inner);
}
if (!group_buf) return nil;
const uint64_t base_address = (uint64_t)[group_buf gpuAddress] + inner;
for (uint32_t j = 0; j < group_n; j++) {
addresses[first + j] = base_address + (uint64_t)j * expert_bytes;
}
[expert_buffers addObject:group_buf];
}
} else
#endif
{
fprintf(stderr, "ds4: Metal GPU addresses require macOS 13 or newer\n");
return nil;
}
if (getenv("DS4_METAL_Q4_EXPERT_TABLE_PROFILE") != NULL) {
fprintf(stderr,
"ds4: Metal Q4 expert address table: experts=%u group=%u buffers=%lu expert_bytes=%.2f MiB\n",
n_total_expert,
table_group_size,
(unsigned long)[expert_buffers count],
ds4_gpu_mib(expert_bytes));
}
DS4MetalQ4ExpertTable *table = [DS4MetalQ4ExpertTable new];
table.addressBuffer = address_buffer;
table.expertBuffers = expert_buffers;
table.nExpert = n_total_expert;
table.expertBytes = expert_bytes;
[g_q4_expert_table_cache setObject:table forKey:key];
return table;
}
static void ds4_gpu_use_q4_expert_table_resources(
id<MTLCommandBuffer> cb,
id<MTLComputeCommandEncoder> enc,
DS4MetalQ4ExpertTable *table,
bool queue_residency) {
if (!enc || !table) return;
if (table.argumentBuffer) {
[enc useResource:table.argumentBuffer usage:MTLResourceUsageRead];
}
if (table.addressBuffer) {
[enc useResource:table.addressBuffer usage:MTLResourceUsageRead];
}
if (!queue_residency &&
getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL &&
table.residencySet &&
cb &&
[cb respondsToSelector:@selector(useResidencySet:)]) {
[cb useResidencySet:table.residencySet];
}
if (getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") == NULL &&
getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") == NULL) {
return;
}
const NSUInteger count = [table.expertBuffers count];
for (NSUInteger base = 0; base < count; base += 64u) {
const NSUInteger n = count - base < 64u ? count - base : 64u;
__unsafe_unretained id<MTLResource> resources[64];
for (NSUInteger i = 0; i < n; i++) {
resources[i] = [table.expertBuffers objectAtIndex:base + i];
}
[enc useResources:resources count:n usage:MTLResourceUsageRead];
}
}
int ds4_gpu_indexer_score_one_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_head,
uint32_t head_dim,
float scale) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!scores || !q || !weights || !index_comp ||
n_comp == 0 || n_head == 0 || head_dim == 0) {
return 0;
}
@autoreleasepool {
const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
const uint64_t weight_bytes = (uint64_t)n_head * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
const uint64_t score_bytes = (uint64_t)n_comp * sizeof(float);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> wbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(index_comp);
id<MTLBuffer> scorebuf = ds4_gpu_tensor_buffer(scores);
if (!qbuf || !wbuf || !compbuf || !scorebuf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(weights) < weight_bytes ||
ds4_gpu_tensor_bytes(index_comp) < comp_bytes ||
ds4_gpu_tensor_bytes(scores) < score_bytes) {
fprintf(stderr, "ds4: Metal graph indexer score received undersized buffers\n");
return 0;
}
if (n_head == 64 && head_dim == 128) {
id<MTLComputePipelineState> direct_pipeline =
ds4_gpu_hot_pipeline(g_dsv4_indexer_score_one_direct_pipeline,
"kernel_dsv4_indexer_score_one_direct");
if (!direct_pipeline) return 0;
ds4_gpu_dsv4_indexer_scores_fused_args args = {
.n_comp = n_comp,
.n_tokens = 1,
.n_head = n_head,
.head_dim = head_dim,
.pos0 = 0,
.ratio = 4,
.q_token_stride = (uint64_t)n_head * head_dim * sizeof(float),
.q_head_stride = (uint64_t)head_dim * sizeof(float),
.weights_token_stride = (uint64_t)n_head * sizeof(float),
.index_row_stride = (uint64_t)head_dim * sizeof(float),
.score_token_stride = (uint64_t)n_comp * sizeof(float),
.scale = scale,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:direct_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
[enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3];
[enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
[enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_comp, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer direct score")) return 0;
return 1;
}
const uint64_t head_score_bytes = (uint64_t)n_comp * n_head * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_head_scores_buffer,
&g_indexer_head_scores_bytes,
(NSUInteger)head_score_bytes,
"ds4_indexer_head_scores")) {
return 0;
}
ds4_gpu_q8_0_matvec_args dot_args =
ds4_gpu_make_f32_mv_args(head_dim, n_comp, n_head);
ds4_gpu_mv_dispatch dot_dispatch =
ds4_gpu_make_plain_mv_dispatch(head_dim, 1);
dot_args.nr0 = dot_dispatch.nr0;
id<MTLComputePipelineState> dot_pipeline =
ds4_gpu_get_mul_mv_pipeline(dot_dispatch.function_name, dot_dispatch.nsg);
if (!dot_pipeline) return 0;
ds4_gpu_dsv4_indexer_weighted_sum_args sum_args = {
.ne00 = (int64_t)n_comp,
.ne01 = 1,
.ne02 = (int64_t)n_head,
.nb00 = sizeof(float),
.nb01 = (uint64_t)n_comp * sizeof(float),
.nb02 = (uint64_t)n_comp * sizeof(float),
.ne10 = (int64_t)n_head,
.ne11 = 1,
.nb10 = sizeof(float),
.nb11 = (uint64_t)n_head * sizeof(float),
.ne0 = (int64_t)n_comp,
.ne1 = 1,
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_comp * sizeof(float),
.scale = scale,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:dot_pipeline];
[enc setBytes:&dot_args length:sizeof(dot_args) atIndex:0];
[enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:1];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2];
[enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:3];
if (dot_dispatch.smem) {
[enc setThreadgroupMemoryLength:dot_dispatch.smem atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + (NSUInteger)dot_dispatch.nr0 - 1u) / (NSUInteger)dot_dispatch.nr0,
n_head,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dot_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_indexer_weighted_sum_pipeline];
[enc setBytes:&sum_args length:sizeof(sum_args) atIndex:0];
[enc setBuffer:g_indexer_head_scores_buffer offset:0 atIndex:1];
[enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
[enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer score")) return 0;
}
return 1;
}
static int ds4_gpu_indexer_scores_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!scores || !q || !weights || !index_comp ||
n_comp == 0 || n_tokens == 0 || n_head == 0 || head_dim == 0 || ratio == 0) {
return 0;
}
@autoreleasepool {
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
const uint64_t weight_bytes = (uint64_t)n_tokens * n_head * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> wbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(index_comp);
id<MTLBuffer> scorebuf = ds4_gpu_tensor_buffer(scores);
if (!qbuf || !wbuf || !compbuf || !scorebuf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(weights) < weight_bytes ||
ds4_gpu_tensor_bytes(index_comp) < comp_bytes ||
ds4_gpu_tensor_bytes(scores) < score_bytes) {
fprintf(stderr, "ds4: Metal graph indexer prefill scores received undersized buffers\n");
return 0;
}
if (head_dim != 128) {
fprintf(stderr, "ds4: Metal fused DS4 indexer scores expect 128-wide rows\n");
return 0;
}
/*
* The NAX/TensorOps score builder is a prefill-only win. At small
* batches and in one-token decode the setup cost is not amortized, so
* those paths keep the older direct/tiled score kernels.
*/
const bool use_nax = ds4_gpu_mpp_available() && n_tokens >= 16u;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
use_nax ? "kernel_dsv4_indexer_scores_nax" :
(g_quality_mode ? "kernel_dsv4_indexer_scores_tiled_f32"
: "kernel_dsv4_indexer_scores_tiled"));
if (!pipeline) return 0;
ds4_gpu_dsv4_indexer_scores_fused_args args = {
.n_comp = n_comp,
.n_tokens = n_tokens,
.n_head = n_head,
.head_dim = head_dim,
.pos0 = pos0,
.ratio = ratio,
.q_token_stride = (uint64_t)n_head * head_dim * sizeof(float),
.q_head_stride = (uint64_t)head_dim * sizeof(float),
.weights_token_stride = (uint64_t)n_head * sizeof(float),
.index_row_stride = (uint64_t)head_dim * sizeof(float),
.score_token_stride = (uint64_t)n_comp * sizeof(float),
.scale = scale,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
[enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(index_comp) atIndex:3];
[enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
if (use_nax) {
const NSUInteger q_shared = 2u * 32u * 32u;
const NSUInteger k_shared = 32u * 128u;
const NSUInteger dot_shared = 32u * 32u;
[enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) +
dot_shared * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u,
((NSUInteger)n_tokens + 15u) / 16u,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
} else if (g_quality_mode) {
const NSUInteger q_shared = 8u * 128u;
const NSUInteger k_shared = 32u * 128u;
const NSUInteger dot_shared = 8u * 32u;
[enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u,
((NSUInteger)n_tokens + 7u) / 8u,
1)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
} else {
const NSUInteger q_shared = 8u * 128u;
const NSUInteger k_shared = 32u * 128u;
const NSUInteger dot_shared = 8u * 32u;
[enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) +
dot_shared * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_comp + 31u) / 32u,
((NSUInteger)n_tokens + 7u) / 8u,
1)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer prefill scores")) return 0;
}
return 1;
}
int ds4_gpu_indexer_scores_prefill_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale) {
return ds4_gpu_indexer_scores_batch_tensor(scores,
q,
weights,
index_comp,
n_comp,
n_tokens,
0,
n_head,
head_dim,
ratio,
scale);
}
int ds4_gpu_indexer_scores_decode_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *index_comp,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
uint32_t ratio,
float scale) {
return ds4_gpu_indexer_scores_batch_tensor(scores,
q,
weights,
index_comp,
n_comp,
n_tokens,
pos0,
n_head,
head_dim,
ratio,
scale);
}
/* Radix select, deterministic compaction and one in-place sort of the winners
* (kernel_topk_select_*): every dispatch spans the GPU, where the merge ladder
* ends on two or three threadgroups. Returns 0 without encoding anything when
* the pipelines are unavailable, so the caller keeps the ladder. */
static int ds4_gpu_indexer_topk_select(id<MTLBuffer> selbuf, NSUInteger sel_off,
id<MTLBuffer> scorebuf, NSUInteger score_off,
uint32_t n, uint32_t rows, uint32_t top_k) {
id<MTLComputePipelineState> hist = ds4_gpu_get_pipeline("kernel_topk_select_hist");
id<MTLComputePipelineState> count = ds4_gpu_get_pipeline("kernel_topk_select_count");
id<MTLComputePipelineState> write = ds4_gpu_get_pipeline("kernel_topk_select_write");
id<MTLComputePipelineState> sort = ds4_gpu_get_pipeline("kernel_topk_select_sort");
id<MTLComputePipelineState> zero = ds4_gpu_get_pipeline("kernel_topk_select_zero");
if (!hist || !count || !write || !sort || !zero ||
hist.maxTotalThreadsPerThreadgroup < 1024u || count.maxTotalThreadsPerThreadgroup < 1024u ||
write.maxTotalThreadsPerThreadgroup < 1024u || sort.maxTotalThreadsPerThreadgroup < 1024u) return 0;
const uint32_t n_tg = (n + 2047u) / 2048u;
/* sized for the widest row (2^23 keys), so the scratch is allocated once per row count */
const NSUInteger state_off = (NSUInteger)rows * 4u * 2048u * sizeof(uint32_t);
const NSUInteger counts_off = state_off + (NSUInteger)rows * 4u * sizeof(uint32_t);
const NSUInteger bytes = counts_off + (NSUInteger)rows * 4096u * 2u * sizeof(uint32_t);
id<MTLBuffer> before = g_topk_select_buffer;
if (!ds4_gpu_ensure_scratch_buffer(&g_topk_select_buffer, &g_topk_select_bytes, bytes,
"ds4_topk_select")) return 0;
ds4_gpu_kargs_topk_select args = {n, rows, top_k, 0, g_topk_select_parity, n_tg};
g_topk_select_parity ^= 1u;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc;
if (g_topk_select_buffer != before) { /* a fresh scratch starts with empty histograms */
const uint32_t words = (uint32_t)(g_topk_select_bytes / sizeof(uint32_t));
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:zero];
[enc setBuffer:g_topk_select_buffer offset:0 atIndex:0];
[enc setBytes:&words length:sizeof(words) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake((words + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
for (args.pass = 0; args.pass < 3u; args.pass++) {
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:hist];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:scorebuf offset:score_off atIndex:1];
[enc setBuffer:g_topk_select_buffer offset:0 atIndex:2];
[enc setThreadgroupMemoryLength:2048u * sizeof(uint32_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, rows, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:count];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:scorebuf offset:score_off atIndex:1];
[enc setBuffer:g_topk_select_buffer offset:0 atIndex:2];
[enc setBuffer:g_topk_select_buffer offset:state_off atIndex:3];
[enc setBuffer:g_topk_select_buffer offset:counts_off atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, rows, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:write];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:scorebuf offset:score_off atIndex:1];
[enc setBuffer:g_topk_select_buffer offset:state_off atIndex:2];
[enc setBuffer:g_topk_select_buffer offset:counts_off atIndex:3];
[enc setBuffer:selbuf offset:sel_off atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, rows, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:sort];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:scorebuf offset:score_off atIndex:1];
[enc setBuffer:selbuf offset:sel_off atIndex:2];
[enc setThreadgroupMemoryLength:2u * 2048u * sizeof(uint32_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1) threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "indexer top-k select") ? 1 : 0;
}
static int ds4_gpu_indexer_topk_tensor_impl(
ds4_gpu_tensor *selected,
const ds4_gpu_tensor *scores,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k,
uint32_t causal_start,
uint32_t causal_ratio) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!selected || !scores || n_comp == 0 || n_tokens == 0 || top_k == 0 || top_k > n_comp) return 0;
@autoreleasepool {
id<MTLComputePipelineState> sort_pipeline = causal_ratio ?
ds4_gpu_get_pipeline(getenv("DS4_METAL_DISABLE_V41_TOPK_SHUFFLE") ?
"kernel_argsort_f32_i32_desc_causal" : "kernel_argsort_f32_i32_desc_causal_shuffle") :
g_argsort_f32_i32_desc_pipeline;
id<MTLComputePipelineState> merge_pipeline = causal_ratio ?
ds4_gpu_get_pipeline(top_k == 512u && !getenv("DS4_METAL_DISABLE_V41_TOPK_PREFIX") ?
"kernel_argsort_merge_f32_i32_desc_causal_prefix" :
"kernel_argsort_merge_f32_i32_desc_causal") : g_argsort_merge_f32_i32_desc_pipeline;
if (!sort_pipeline || !merge_pipeline) return 0;
const uint64_t score_bytes = (uint64_t)n_comp * n_tokens * sizeof(float);
const uint64_t selected_bytes = (uint64_t)top_k * n_tokens * sizeof(uint32_t);
id<MTLBuffer> scorebuf = ds4_gpu_tensor_buffer(scores);
id<MTLBuffer> selbuf = ds4_gpu_tensor_buffer(selected);
if (!scorebuf || !selbuf ||
ds4_gpu_tensor_bytes(scores) < score_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes) {
fprintf(stderr, "ds4: Metal graph indexer top-k received undersized buffers\n");
return 0;
}
/* wide rows: the merge ladder's last rounds run on two or three threadgroups */
if (!causal_ratio && top_k <= 2048u && n_comp >= 16384u && n_comp <= (1u << 23) &&
!getenv("DS4_METAL_DISABLE_V41_TOPK_SELECT") &&
ds4_gpu_indexer_topk_select(selbuf, ds4_gpu_tensor_offset(selected), scorebuf,
ds4_gpu_tensor_offset(scores), n_comp, n_tokens, top_k)) return 1;
NSUInteger max_threads = g_argsort_f32_i32_desc_pipeline.maxTotalThreadsPerThreadgroup;
if (max_threads == 0) max_threads = 256;
int32_t nth = 1;
while ((uint32_t)nth < n_comp && (uint64_t)2u * (uint64_t)nth <= (uint64_t)max_threads) {
nth *= 2;
}
if (causal_ratio && (sort_pipeline.maxTotalThreadsPerThreadgroup < (NSUInteger)nth ||
(causal_start + 1u) / causal_ratio <= (uint32_t)nth / 2u)) return 0;
const int32_t npr = (int32_t)((n_comp + (uint32_t)nth - 1u) / (uint32_t)nth);
const int32_t block_top_k = (int32_t)(top_k < (uint32_t)nth ? top_k : (uint32_t)nth);
int32_t work_width = (int32_t)top_k;
if (npr > 1) {
const int32_t last_block = (int32_t)n_comp - (npr - 1) * nth;
work_width = (npr - 1) * block_top_k + (last_block < block_top_k ? last_block : block_top_k);
}
const uint64_t scratch_row_bytes = (uint64_t)work_width * sizeof(uint32_t);
const bool one_pass = npr <= 1;
const uint64_t scratch_bytes = one_pass ? scratch_row_bytes * n_tokens :
2u * scratch_row_bytes * n_tokens;
if (!ds4_gpu_ensure_scratch_buffer(&g_indexer_topk_buffer,
&g_indexer_topk_bytes,
(NSUInteger)scratch_bytes,
"ds4_indexer_topk")) {
return 0;
}
ds4_gpu_kargs_argsort args = {
.ne00 = (int32_t)n_comp,
.ne01 = (int32_t)n_tokens,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = (uint64_t)n_comp * sizeof(float),
.nb02 = (uint64_t)n_comp * n_tokens * sizeof(float),
.nb03 = (uint64_t)n_comp * n_tokens * sizeof(float),
.ne0 = work_width,
.ne1 = (int32_t)n_tokens,
.ne2 = 1,
.ne3 = 1,
.top_k = block_top_k,
.causal_start = causal_start,
.causal_ratio = causal_ratio,
};
// kernel_argsort_f32_i32_desc stages the block's scores behind the
// index array: nth int32 indices + nth float scores.
const NSUInteger smem = (((NSUInteger)nth * (sizeof(int32_t) + sizeof(float))) + 15u) & ~(NSUInteger)15u;
NSUInteger cur_off = 0;
NSUInteger next_off = (NSUInteger)scratch_row_bytes * n_tokens;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:sort_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1];
[enc setBuffer:one_pass ? selbuf : g_indexer_topk_buffer
offset:one_pass ? ds4_gpu_tensor_offset(selected) : cur_off
atIndex:2];
[enc setThreadgroupMemoryLength:smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)npr * n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake((NSUInteger)nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
int32_t len = block_top_k;
while (len < work_width) {
const int32_t nm = (work_width + 2 * len - 1) / (2 * len);
const bool final_merge = nm == 1;
NSUInteger merge_threads = merge_pipeline.maxTotalThreadsPerThreadgroup;
if (merge_threads == 0 || merge_threads > 512u) merge_threads = 512u;
if (merge_threads > (NSUInteger)len) merge_threads = (NSUInteger)len;
if (merge_threads == 0) merge_threads = 1;
ds4_gpu_kargs_argsort_merge merge_args = {
.ne00 = (int64_t)n_comp,
.ne01 = (int64_t)n_tokens,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = (uint64_t)n_comp * sizeof(float),
.nb02 = (uint64_t)n_comp * n_tokens * sizeof(float),
.nb03 = (uint64_t)n_comp * n_tokens * sizeof(float),
.ne0 = work_width,
.ne1 = (int32_t)n_tokens,
.ne2 = 1,
.ne3 = 1,
.top_k = nm == 1 ? (int32_t)top_k : work_width,
.len = len,
.causal_start = causal_start,
.causal_ratio = causal_ratio,
.block_width = (uint32_t)nth,
.block_top_k = (uint32_t)block_top_k,
};
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:merge_pipeline];
[enc setBytes:&merge_args length:sizeof(merge_args) atIndex:0];
[enc setBuffer:scorebuf offset:ds4_gpu_tensor_offset(scores) atIndex:1];
[enc setBuffer:g_indexer_topk_buffer offset:cur_off atIndex:2];
[enc setBuffer:final_merge ? selbuf : g_indexer_topk_buffer
offset:final_merge ? ds4_gpu_tensor_offset(selected) : next_off
atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nm * n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(merge_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
const NSUInteger tmp = cur_off;
cur_off = next_off;
next_off = tmp;
len <<= 1;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "indexer top-k")) return 0;
}
return 1;
}
int ds4_gpu_indexer_topk_tensor(ds4_gpu_tensor *selected, const ds4_gpu_tensor *scores,
uint32_t n_comp, uint32_t n_tokens, uint32_t top_k) {
return ds4_gpu_indexer_topk_tensor_impl(selected, scores, n_comp, n_tokens, top_k, 0, 0);
}
int ds4_gpu_dsv41_indexer_topk_batch(ds4_gpu_tensor *selected, const ds4_gpu_tensor *scores,
uint32_t width, uint32_t rows, uint32_t start, uint32_t ratio) {
if ((ratio != 1u && ratio != 2u) || !rows || rows > UINT32_MAX - start ||
width > INT32_MAX || rows > INT32_MAX || (start + rows) / ratio > width ||
(start + 1u) / ratio <= 512u) return 0;
return ds4_gpu_indexer_topk_tensor_impl(selected, scores, width, rows, 512u, start, ratio);
}
int ds4_gpu_dsv41_indexer_all_batch(ds4_gpu_tensor *selected, uint32_t rows,
uint32_t start, uint32_t ratio) {
if (!ratio || !rows || rows > UINT32_MAX - start || (start + rows) / ratio > 512u ||
!selected || ds4_gpu_tensor_bytes(selected) < (uint64_t)rows * 512u * sizeof(int32_t)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_indexer_all");
if (!pipeline) return 0;
const uint32_t args[] = {rows, start, ratio, 512u};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(selected) offset:ds4_gpu_tensor_offset(selected) atIndex:1];
[enc dispatchThreads:MTLSizeMake(512u, rows, 1)
threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 indexer select all");
}
}
int ds4_gpu_argmax_tensor(
ds4_gpu_tensor *out_idx,
const ds4_gpu_tensor *logits,
uint32_t n_vocab) {
if (!out_idx || !logits || n_vocab == 0) return 0;
if (ds4_gpu_tensor_bytes(out_idx) < sizeof(int32_t) ||
ds4_gpu_tensor_bytes(logits) < (uint64_t)n_vocab * sizeof(float)) {
fprintf(stderr, "ds4: Metal graph argmax received undersized buffers\n");
return 0;
}
return ds4_gpu_indexer_topk_tensor(out_idx, logits, n_vocab, 1, 1);
}
static int ds4_gpu_dsv4_topk_mask_tensor_impl(
ds4_gpu_tensor *mask,
const ds4_gpu_tensor *topk,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k,
uint32_t mask_stride) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!mask || !topk || n_comp == 0 || n_tokens == 0 || top_k == 0 || mask_stride < n_comp) return 0;
@autoreleasepool {
const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t);
const uint64_t mask_bytes = ((uint64_t)mask_stride * (n_tokens - 1u) + n_comp) * sizeof(float);
id<MTLBuffer> topkbuf = ds4_gpu_tensor_buffer(topk);
id<MTLBuffer> maskbuf = ds4_gpu_tensor_buffer(mask);
if (!topkbuf || !maskbuf ||
ds4_gpu_tensor_bytes(topk) < topk_bytes ||
ds4_gpu_tensor_bytes(mask) < mask_bytes) {
fprintf(stderr, "ds4: Metal dsv4 top-k mask received undersized buffers\n");
return 0;
}
ds4_gpu_dsv4_topk_mask_args args = {
.ne00 = (int64_t)top_k,
.ne01 = (int64_t)n_tokens,
.nb00 = sizeof(int32_t),
.nb01 = (uint64_t)top_k * sizeof(int32_t),
.ne0 = (int64_t)n_comp,
.ne1 = (int64_t)n_tokens,
.nb0 = sizeof(float),
.nb1 = (uint64_t)mask_stride * sizeof(float),
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_topk_mask_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1];
[enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)n_comp * n_tokens) + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_topk_mask_scatter_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1];
[enc setBuffer:maskbuf offset:ds4_gpu_tensor_offset(mask) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake((((NSUInteger)top_k * n_tokens) + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "dsv4 top-k mask")) return 0;
}
return 1;
}
int ds4_gpu_dsv4_topk_mask_tensor(
ds4_gpu_tensor *mask,
const ds4_gpu_tensor *topk,
uint32_t n_comp,
uint32_t n_tokens,
uint32_t top_k) {
return ds4_gpu_dsv4_topk_mask_tensor_impl(mask, topk, n_comp, n_tokens, top_k, n_comp);
}
static int ds4_gpu_matmul_q8_0_legacy_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
bool prefer_decode_mpp,
bool prefill_default,
int round_bf16,
int *fused_round) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if ((in_dim & 31u) != 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal Q8_0 tensor matmul received undersized activation buffers\n");
return 0;
}
const uint64_t blocks = in_dim / 32;
const uint64_t row_bytes = blocks * 34;
const uint64_t weight_bytes = out_dim * row_bytes;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal Q8_0 tensor matmul range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) {
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (n_tok == 1) {
if (ds4_gpu_mpp_available() &&
prefer_decode_mpp &&
(in_dim % 64u) == 0) {
const char *nax_fn = "kernel_mul_mm_q8_0_f32_nax_direct_rhs";
id<MTLComputePipelineState> mpp_pipeline =
ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false);
if (mpp_pipeline) {
ds4_gpu_mul_mm_args args =
ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:mpp_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1u,
((NSUInteger)out_dim + 63u) / 64u,
1u)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 decode MPP matmul")) {
return 0;
}
return 1;
}
ds4_gpu_warn_mpp_fallback();
}
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
if (out_dim > 65536u) mv_dispatch.nsg = 8;
mv_args.nr0 = mv_dispatch.nr0;
const char *mv_fn = mv_dispatch.function_name;
if (round_bf16 && !strcmp(mv_fn, "kernel_mul_mv_q8_0_f32")) {
mv_fn = round_bf16 == 2 ? "kernel_mul_mv_q8_0_f32_bf16io" : "kernel_mul_mv_q8_0_f32_bf16";
if (fused_round) *fused_round = 1;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(mv_fn, mv_dispatch.nsg);
if (!pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
1,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matvec")) {
return 0;
}
return 1;
}
const uint64_t mv_ext_max_tokens =
ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u);
if (n_tok <= mv_ext_max_tokens && (in_dim % 128u) == 0) {
const int16_t nsg = ds4_gpu_mv_ext_nsg();
const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok);
const char *fn_name = ds4_gpu_mv_ext_name(1, r1ptg);
id<MTLComputePipelineState> pipeline =
fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
if (!pipeline) return 0;
const int16_t nypsg = 32 / nxpsg;
const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
ds4_gpu_mul_mv_ext_args args =
ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg,
((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor mul_mv_ext")) {
return 0;
}
return 1;
}
/*
* Dense Q8_0 prefill is the cleanest DS4 TensorOps shape: M/N/K are
* aligned and the RHS activation matrix is already dense. The retained
* kernel dequantizes each 64x32 weight tile to half in threadgroup
* memory, then uses direct-RHS MPP for the activation tile. This avoids
* staging RHS into threadgroup memory and was the direct replacement for
* the slower generic MPP prototype.
*/
/*
* An unaligned token count used to send the entire projection through
* the generic kernel. Run its aligned prefix through TensorOps and
* leave only the final partial tile to the boundary-safe kernel.
* Tiny prompts do not amortize the second dispatch, while --quality
* deliberately retains the single-kernel arithmetic schedule.
*/
const bool split_nax_prefix =
!g_quality_mode && n_tok >= 192u && (n_tok % 32u) != 0u;
const uint64_t nax_rows =
(n_tok % 32u) == 0u ? n_tok :
(split_nax_prefix ? n_tok - (n_tok % 32u) : 0u);
uint64_t generic_row0 = 0u;
uint64_t generic_rows = n_tok;
if (ds4_gpu_mpp_available() &&
nax_rows >= 32u &&
(in_dim % 64u) == 0 &&
(out_dim % 64u) == 0) {
uint64_t nax_tile_n = 32u;
if ((nax_rows % 128u) == 0) {
nax_tile_n = 128u;
} else if ((nax_rows % 64u) == 0) {
nax_tile_n = 64u;
}
const char *nax_fn = nax_tile_n == 128u
? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128"
: (nax_tile_n == 64u
? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64"
: "kernel_mul_mm_q8_0_f32_nax_direct_rhs");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false);
if (pipeline) {
ds4_gpu_mul_mm_args args =
ds4_gpu_make_mm_args(in_dim, out_dim, nax_rows, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(nax_rows / nax_tile_n),
(NSUInteger)out_dim / 64u,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (nax_rows == n_tok) {
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 NAX tensor matmul")) {
return 0;
}
return 1;
}
generic_row0 = nax_rows;
generic_rows = n_tok - nax_rows;
}
if (!pipeline) ds4_gpu_warn_mpp_fallback();
}
const bool bc_inp = (in_dim % 32u) != 0;
const bool bc_out =
(out_dim % 64u) != 0 || (generic_rows % 32u) != 0;
/* Amortize dequantization across token tiles with identical half
* operands. Default is scoped by the Qwen entry point to M3 Ultra. */
const int unpack_override = ds4_gpu_env_bool("DS4_QWEN4_Q8_PREFILL_UNPACK");
const bool prefill_unpack = generic_rows >= 32u &&
(unpack_override >= 0 ? unpack_override == 1 : prefill_default);
const char *fn = "kernel_mul_mm_q8_0_f32";
id<MTLBuffer> mat_weights = wbuf;
NSUInteger mat_offset = (NSUInteger)inner_offset;
uint64_t mat_row_bytes = row_bytes;
const bool unpack_fits = in_dim > 0u && out_dim > 0u && out_dim <= (64u << 20) / in_dim;
const uint64_t unpack_bytes = unpack_fits ? in_dim * out_dim * sizeof(uint16_t) : 0u;
if (prefill_unpack && unpack_fits) {
if (!ds4_gpu_ensure_scratch_buffer(&g_q8_prefill_scratch_buffer,
&g_q8_prefill_scratch_bytes, (NSUInteger)unpack_bytes, "q8 prefill weights")) return 0;
id<MTLComputePipelineState> unpack = ds4_gpu_get_pipeline("kernel_q8_prefill_unpack");
if (!unpack) return 0;
const uint32_t n_blocks = (uint32_t)(in_dim * out_dim / 32u);
id<MTLComputeCommandEncoder> unpack_enc = ds4_gpu_compute_encoder(cb);
[unpack_enc setComputePipelineState:unpack];
[unpack_enc setBytes:&n_blocks length:sizeof(n_blocks) atIndex:0];
[unpack_enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[unpack_enc setBuffer:g_q8_prefill_scratch_buffer offset:0 atIndex:2];
[unpack_enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_blocks * 8u + 255u)/256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, unpack_enc);
fn = "kernel_mul_mm_f16_f32";
mat_weights = g_q8_prefill_scratch_buffer;
mat_offset = 0;
mat_row_bytes = in_dim * sizeof(uint16_t);
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mm_pipeline(fn, bc_inp, bc_out);
if (!pipeline) return 0;
ds4_gpu_mul_mm_args args =
ds4_gpu_make_mm_args(in_dim, out_dim, generic_rows, mat_row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:mat_weights offset:mat_offset atIndex:1];
[enc setBuffer:xbuf
offset:ds4_gpu_tensor_offset(x) +
(NSUInteger)(generic_row0 * in_dim * sizeof(float))
atIndex:2];
[enc setBuffer:outbuf
offset:ds4_gpu_tensor_offset(out) +
(NSUInteger)(generic_row0 * out_dim * sizeof(float))
atIndex:3];
[enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)generic_rows + 31u) / 32u,
((NSUInteger)out_dim + 63u) / 64u,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor matmul")) {
return 0;
}
}
return 1;
}
static int ds4_gpu_matmul_q8_0_tensor_impl(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok, bool prefill_default,
int round_bf16, int *fused_round) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if ((in_dim & 31u) != 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) {
return 0;
}
const int profile_requested =
n_tok > 8u && ds4_gpu_env_bool("DS4_METAL_Q8_PREFILL_PROFILE") > 0;
int profile_prefill = 0;
int split_batch_for_profile = 0;
const char *profile_label = NULL;
char profile_label_buf[128];
char profile_fallback[128];
if (profile_requested) {
snprintf(profile_fallback, sizeof(profile_fallback),
"q8 weight_off=%llu in=%llu out=%llu tok=%llu",
(unsigned long long)weight_offset,
(unsigned long long)in_dim,
(unsigned long long)out_dim,
(unsigned long long)n_tok);
snprintf(profile_label_buf, sizeof(profile_label_buf), "%s", profile_fallback);
profile_label = profile_label_buf;
const char *profile_filter = getenv("DS4_METAL_Q8_PREFILL_PROFILE_FILTER");
profile_prefill =
profile_requested &&
(!profile_filter || !profile_filter[0] ||
strstr(profile_label, profile_filter) != NULL);
}
if (profile_prefill) {
if (g_batch_cb) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
split_batch_for_profile = 1;
}
}
const double profile_t0 = profile_prefill ? ds4_gpu_now_ms() : 0.0;
int ok = ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size,
weight_offset, in_dim, out_dim,
x, n_tok, false, prefill_default,
round_bf16, fused_round);
if (profile_prefill) {
if (split_batch_for_profile && ds4_gpu_end_commands() == 0) {
ok = 0;
}
const double elapsed_ms = ds4_gpu_now_ms() - profile_t0;
fprintf(stderr,
"ds4: Metal Q8_0 prefill profile %s in=%llu out=%llu tok=%llu %.3f ms\n",
profile_label ? profile_label : profile_fallback,
(unsigned long long)in_dim,
(unsigned long long)out_dim,
(unsigned long long)n_tok,
elapsed_ms);
if (split_batch_for_profile && ds4_gpu_begin_commands() == 0) {
ok = 0;
}
}
return ok;
}
int ds4_gpu_matmul_q8_0_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_q8_0_tensor_impl(out, model_map, model_size, weight_offset, in_dim, out_dim, x, n_tok, false, 0, NULL);
}
static int ds4_gpu_matmul_q8_0_rows_tensor_impl(
ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim,
const ds4_gpu_tensor *x, uint32_t n_rows, int round_bf16);
/* The matvec rounds its output to bf16 in the kernel where the single-row Q8
* path applies; any other path gets the separate rounding pass. */
int ds4_gpu_matmul_q8_0_bf16_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
int fused = 0;
if (n_tok >= 2 && n_tok <= 8 &&
ds4_gpu_matmul_q8_0_rows_tensor_impl(out, model_map, model_size, weight_offset,
in_dim, out_dim, x, (uint32_t)n_tok, 1)) return 1;
if (!ds4_gpu_matmul_q8_0_tensor_impl(out, model_map, model_size, weight_offset, in_dim, out_dim, x, n_tok, false, 1, &fused)) return 0;
return fused ? 1 : ds4_gpu_dsv41_quantize(out, (uint32_t)out_dim, (uint32_t)n_tok, DS4_V41_BF16);
}
/* As above, with the activation rounded to bf16 on load as well; rows other
* than one round the activation in place first. */
int ds4_gpu_matmul_q8_0_bf16io_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
ds4_gpu_tensor *x,
uint64_t n_tok) {
if (n_tok >= 2 && n_tok <= 8 &&
ds4_gpu_matmul_q8_0_rows_tensor_impl(out, model_map, model_size, weight_offset,
in_dim, out_dim, x, (uint32_t)n_tok, 2)) return 1;
if (n_tok != 1) {
if (!ds4_gpu_dsv41_quantize(x, (uint32_t)in_dim, (uint32_t)n_tok, DS4_V41_BF16)) return 0;
return ds4_gpu_matmul_q8_0_bf16_tensor(out, model_map, model_size, weight_offset, in_dim, out_dim, x, n_tok);
}
return ds4_gpu_matmul_q8_0_tensor_impl(out, model_map, model_size, weight_offset, in_dim, out_dim, x, 1, false, 2, NULL);
}
int ds4_gpu_qwen4_matmul_q8_0_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
if (!g_initialized && !ds4_gpu_init()) return 0;
return ds4_gpu_matmul_q8_0_tensor_impl(out, model_map, model_size, weight_offset, in_dim, out_dim, x, n_tok,
ds4_gpu_device_name_contains("M3 Ultra"), 0, NULL);
}
/* The matrix rows kernels round a few outputs differently from the row kernels. */
static int ds4_gpu_v41_rows_mma_off(void) {
static int off = -1;
if (off < 0) off = getenv("DS4_METAL_DISABLE_V41_ROWS_MMA") != NULL;
return off;
}
/* Two to eight rows share one pass over the weights (the rows kernels);
* each row's arithmetic is the single-row matvec's. round_bf16 as in
* ds4_gpu_matmul_q8_0_tensor_impl; rounding needs the rows kernel. */
static int ds4_gpu_matmul_q8_0_rows_tensor_impl(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows,
int round_bf16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !x || !model_map || n_rows == 0 ||
n_rows > INT32_MAX || in_dim == 0 || out_dim == 0 ||
(in_dim & 31u) != 0 || in_dim > UINT32_MAX ||
out_dim > UINT32_MAX ||
in_dim > UINT64_MAX / n_rows / sizeof(float) ||
out_dim > UINT64_MAX / n_rows / sizeof(float) ||
ds4_gpu_tensor_bytes(x) <
(uint64_t)n_rows * in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out) <
(uint64_t)n_rows * out_dim * sizeof(float)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t blocks = in_dim / 32u;
const uint64_t row_bytes = blocks * 34u;
if (!xbuf || !outbuf ||
out_dim > UINT64_MAX / row_bytes) {
return 0;
}
const uint64_t weight_bytes = out_dim * row_bytes;
if (weight_offset > model_size ||
weight_bytes > model_size - weight_offset) {
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(
model_map, model_size, weight_offset, weight_bytes,
&inner_offset);
if (!wbuf) return 0;
ds4_gpu_mv_dispatch dispatch = ds4_gpu_make_q8_0_mv_dispatch();
if (out_dim > 65536u) dispatch.nsg = 8;
ds4_gpu_q8_0_matvec_args args =
ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
args.ne11 = (int32_t)n_rows;
args.nb12 = (uint64_t)n_rows * in_dim * sizeof(float);
args.nb13 = args.nb12;
args.ne1 = (int32_t)n_rows;
args.nr0 = dispatch.nr0;
/* the rows one after another on the single-row layout, four outputs
* per simdgroup */
const bool rows_kernel = n_rows >= 2u && n_rows <= 8u;
if (!rows_kernel && round_bf16) return 0;
if (rows_kernel) { dispatch.nr0 = 4; args.nr0 = 4; }
const char *rows_name = round_bf16 == 2 ? "kernel_mul_mv_q8_0_f32_bf16io_rows_seq4" :
round_bf16 == 1 ? "kernel_mul_mv_q8_0_f32_bf16_rows_seq4" : "kernel_mul_mv_q8_0_f32_rows_seq4";
/* the matrix rows kernel pays off from three rows (four on the
* widest K); eight simdgroups split a K of 4096 or more */
const int mma_off = ds4_gpu_v41_rows_mma_off();
const int mma = in_dim >= 4096u ? 8 : 4;
const bool mma_kernel = !mma_off && rows_kernel && n_rows >= (in_dim >= 8192u ? 4u : 3u) &&
(in_dim % (mma == 8 ? 256u : 128u)) == 0 && (out_dim % 8u) == 0;
id<MTLComputePipelineState> pipeline = mma_kernel ?
ds4_gpu_get_pipeline(mma == 8 ?
(round_bf16 == 2 ? "kernel_mul_mv_q8_0_f32_bf16io_mma_rows8" :
round_bf16 == 1 ? "kernel_mul_mv_q8_0_f32_bf16_mma_rows8" : "kernel_mul_mv_q8_0_f32_mma_rows8") :
(round_bf16 == 2 ? "kernel_mul_mv_q8_0_f32_bf16io_mma_rows" :
round_bf16 == 1 ? "kernel_mul_mv_q8_0_f32_bf16_mma_rows" : "kernel_mul_mv_q8_0_f32_mma_rows")) :
ds4_gpu_get_mul_mv_pipeline(rows_kernel ? rows_name : dispatch.function_name, dispatch.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
if (mma_kernel) {
/* per simdgroup an activation and a weight block, then the partial tiles */
[enc setThreadgroupMemoryLength:(NSUInteger)mma * (2u * 8u * 32u + 64u) * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)out_dim / 8u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mma, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 decode-row matmul");
}
[enc setThreadgroupMemoryLength:(rows_kernel ? 32u * sizeof(float) * 4u : dispatch.smem) atIndex:0];
[enc dispatchThreadgroups:
MTLSizeMake(((NSUInteger)out_dim +
(NSUInteger)dispatch.nr0 - 1u) /
(NSUInteger)dispatch.nr0,
rows_kernel ? 1u : (NSUInteger)n_rows,
1)
threadsPerThreadgroup:
MTLSizeMake(32, (NSUInteger)dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(
cb, owned, "Q8_0 exact decode-row matvec");
}
}
int ds4_gpu_matmul_q8_0_decode_rows_exact_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows) {
return ds4_gpu_matmul_q8_0_rows_tensor_impl(out, model_map, model_size, weight_offset,
in_dim, out_dim, x, n_rows, 0);
}
int ds4_gpu_dsv41_project_pair_q8(ds4_gpu_tensor *out_a, ds4_gpu_tensor *out_b,
const void *model_map, uint64_t model_size,
uint64_t offset_a, uint64_t offset_b, uint32_t in_dim,
uint32_t out_a_dim, uint32_t out_b_dim, const ds4_gpu_tensor *x) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out_a || !out_b || !x || !model_map || !in_dim || (in_dim & 31u) ||
!out_a_dim || !out_b_dim || out_a_dim > 65536u || out_b_dim > 65536u ||
ds4_gpu_tensor_bytes(x) < (uint64_t)in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out_a) < (uint64_t)out_a_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out_b) < (uint64_t)out_b_dim * sizeof(float)) return 0;
@autoreleasepool {
const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u;
const uint64_t bytes_a = row_bytes * out_a_dim, bytes_b = row_bytes * out_b_dim;
if (offset_a > model_size || bytes_a > model_size - offset_a ||
offset_b > model_size || bytes_b > model_size - offset_b) return 0;
uint64_t inner_a = 0, inner_b = 0;
id<MTLBuffer> wa = ds4_gpu_wrap_model_range(model_map, model_size, offset_a, bytes_a, &inner_a);
id<MTLBuffer> wb = ds4_gpu_wrap_model_range(model_map, model_size, offset_b, bytes_b, &inner_b);
if (!wa || !wb) return 0;
const ds4_gpu_mv_dispatch dispatch = ds4_gpu_make_q8_0_mv_dispatch();
ds4_gpu_q8_0_matvec_args args_a = ds4_gpu_make_q8_0_mv_args(in_dim, out_a_dim);
ds4_gpu_q8_0_matvec_args args_b = ds4_gpu_make_q8_0_mv_args(in_dim, out_b_dim);
args_a.nr0 = args_b.nr0 = dispatch.nr0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_q8_0_f32_bf16_pair", dispatch.nsg);
if (!pipeline) return 0;
const NSUInteger nr0 = (NSUInteger)dispatch.nr0;
const NSUInteger groups = ((NSUInteger)out_a_dim + nr0 - 1u) / nr0 +
((NSUInteger)out_b_dim + nr0 - 1u) / nr0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args_a length:sizeof(args_a) atIndex:0];
[enc setBytes:&args_b length:sizeof(args_b) atIndex:1];
[enc setBuffer:wa offset:(NSUInteger)inner_a atIndex:2];
[enc setBuffer:wb offset:(NSUInteger)inner_b atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(out_a) offset:ds4_gpu_tensor_offset(out_a) atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(out_b) offset:ds4_gpu_tensor_offset(out_b) atIndex:6];
[enc setThreadgroupMemoryLength:dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 paired Q8_0 matvec");
}
}
int ds4_gpu_matmul_q8_0_decode_mpp_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size,
weight_offset, in_dim, out_dim,
x, n_tok, true, false, 0, NULL);
}
int ds4_gpu_matmul_q8_0_decode_mpp_model_view_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_q8_0_legacy_tensor(out, model_map, model_size,
weight_offset, in_dim, out_dim,
x, n_tok, true, false, 0, NULL);
}
static const char *ds4_gpu_q4_mv_ext_name(uint32_t weight_type, int16_t r1ptg) {
const char *prefix = NULL;
if (weight_type == DS4_METAL_TENSOR_Q4_K) {
prefix = "kernel_mul_mv_ext_q4_K_f32_r1_";
} else if (weight_type == DS4_METAL_TENSOR_Q4_0) {
prefix = "kernel_mul_mv_ext_q4_0_f32_r1_";
} else {
return NULL;
}
switch (r1ptg) {
case 1: return weight_type == DS4_METAL_TENSOR_Q4_K ?
"kernel_mul_mv_ext_q4_K_f32_r1_1" : "kernel_mul_mv_ext_q4_0_f32_r1_1";
case 2: return weight_type == DS4_METAL_TENSOR_Q4_K ?
"kernel_mul_mv_ext_q4_K_f32_r1_2" : "kernel_mul_mv_ext_q4_0_f32_r1_2";
case 3: return weight_type == DS4_METAL_TENSOR_Q4_K ?
"kernel_mul_mv_ext_q4_K_f32_r1_3" : "kernel_mul_mv_ext_q4_0_f32_r1_3";
case 4: return weight_type == DS4_METAL_TENSOR_Q4_K ?
"kernel_mul_mv_ext_q4_K_f32_r1_4" : "kernel_mul_mv_ext_q4_0_f32_r1_4";
case 5: return weight_type == DS4_METAL_TENSOR_Q4_K ?
"kernel_mul_mv_ext_q4_K_f32_r1_5" : "kernel_mul_mv_ext_q4_0_f32_r1_5";
default:
(void)prefix;
return NULL;
}
}
static const char *ds4_gpu_q4_mm_name(uint32_t weight_type) {
switch (weight_type) {
case DS4_METAL_TENSOR_Q4_0: return "kernel_mul_mm_q4_0_f32";
case DS4_METAL_TENSOR_Q4_K: return "kernel_mul_mm_q4_K_f32";
default: return NULL;
}
}
static const char *ds4_gpu_q4_nax_name(uint32_t weight_type, uint64_t tile_n) {
if (weight_type == DS4_METAL_TENSOR_Q4_0) {
return tile_n == 128u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n128" :
tile_n == 64u ? "kernel_mul_mm_q4_0_f32_nax_direct_rhs_n64" :
"kernel_mul_mm_q4_0_f32_nax_direct_rhs";
}
if (weight_type == DS4_METAL_TENSOR_Q4_K) {
return tile_n == 128u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n128" :
tile_n == 64u ? "kernel_mul_mm_q4_K_f32_nax_direct_rhs_n64" :
"kernel_mul_mm_q4_K_f32_nax_direct_rhs";
}
return NULL;
}
static int ds4_gpu_matmul_quant_impl_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
bool prefer_decode_mpp) {
if (weight_type == DS4_METAL_TENSOR_Q8_0) {
return ds4_gpu_matmul_q8_0_legacy_tensor(out,
model_map,
model_size,
weight_offset,
in_dim,
out_dim,
x,
n_tok,
prefer_decode_mpp, false, 0, NULL);
}
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !x || !model_map ||
in_dim == 0 || out_dim == 0 || n_tok == 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) {
return 0;
}
uint64_t row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(weight_type, (uint32_t)in_dim, &row_bytes)) {
fprintf(stderr, "ds4: Metal quant matmul received unsupported type/dim (%u, in=%llu)\n",
weight_type,
(unsigned long long)in_dim);
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal quant tensor matmul received undersized activation buffers\n");
return 0;
}
if (out_dim > UINT64_MAX / row_bytes) return 0;
const uint64_t weight_bytes = out_dim * row_bytes;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal quant tensor matmul range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
/*
* Small-batch Q4_K goes to the classic (llama.cpp-style) matvec:
* the mul_mv_ext family tops out around 220 GB/s on M5 for the GLM
* DenseQ4 decode shapes while this impl streams 530-650 GB/s
* (misc/q4mv_bench.m). Falls through to ext when unavailable.
*/
if (weight_type == DS4_METAL_TENSOR_Q4_K &&
n_tok <= 8u &&
(in_dim % 256u) == 0 &&
getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) {
const int16_t nsg = 2;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8);
if (pipeline) {
ds4_gpu_q8_0_matvec_args args = {
.ne00 = (int32_t)in_dim,
.ne01 = (int32_t)out_dim,
.ne02 = 1,
.nb00 = 1,
.nb01 = row_bytes,
.nb02 = row_bytes * out_dim,
.nb03 = row_bytes * out_dim,
.ne10 = (int32_t)in_dim,
.ne11 = (int32_t)n_tok,
.ne12 = 1,
.nb10 = sizeof(float),
.nb11 = in_dim * sizeof(float),
.nb12 = in_dim * n_tok * sizeof(float),
.nb13 = in_dim * n_tok * sizeof(float),
.ne0 = (int32_t)out_dim,
.ne1 = (int32_t)n_tok,
.nr0 = 2,
.r2 = 1,
.r3 = 1,
};
const uint64_t rows_ptg = (uint64_t)nsg * 2u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:32 atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg,
(NSUInteger)n_tok,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K classic mul_mv")) return 0;
return 1;
}
}
if (n_tok <= 8u && (in_dim % 128u) == 0) {
const int16_t nsg = 2;
const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
const int16_t r1ptg = (n_tok == 1u) ? 1 : ds4_gpu_mv_ext_r1ptg(n_tok);
const char *fn_name = ds4_gpu_q4_mv_ext_name(weight_type, r1ptg);
id<MTLComputePipelineState> pipeline =
fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
if (pipeline) {
const int16_t nypsg = 32 / nxpsg;
const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
ds4_gpu_mul_mv_ext_args args =
ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, row_bytes, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg,
((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor mul_mv_ext")) return 0;
return 1;
}
}
if (ds4_gpu_mpp_available() &&
n_tok >= 32u &&
(in_dim % 64u) == 0 &&
(out_dim % 64u) == 0 &&
(n_tok % 32u) == 0) {
uint64_t nax_tile_n = 32u;
if ((n_tok % 128u) == 0) {
nax_tile_n = 128u;
} else if ((n_tok % 64u) == 0) {
nax_tile_n = 64u;
}
const char *nax_fn = ds4_gpu_q4_nax_name(weight_type, nax_tile_n);
id<MTLComputePipelineState> pipeline =
nax_fn ? ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false) : nil;
if (pipeline) {
ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n),
(NSUInteger)out_dim / 64u,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 NAX tensor matmul")) return 0;
return 1;
}
ds4_gpu_warn_mpp_fallback();
}
const char *mm_fn = ds4_gpu_q4_mm_name(weight_type);
const bool bc_inp = (in_dim % 32u) != 0;
const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0;
id<MTLComputePipelineState> pipeline =
mm_fn ? ds4_gpu_get_mul_mm_pipeline(mm_fn, bc_inp, bc_out) : nil;
if (!pipeline) return 0;
ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u,
((NSUInteger)out_dim + 63u) / 64u,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4 tensor matmul")) return 0;
}
return 1;
}
int ds4_gpu_matmul_quant_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_quant_impl_tensor(out,
model_map,
model_size,
weight_offset,
weight_type,
in_dim,
out_dim,
x,
n_tok,
false);
}
int ds4_gpu_matmul_quant_decode_mpp_model_view_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
/* Classic Q8_0 matvec is substantially faster than MPP at n_tok=1. */
return ds4_gpu_matmul_quant_impl_tensor(out,
model_map,
model_size,
weight_offset,
weight_type,
in_dim,
out_dim,
x,
n_tok,
false);
}
int ds4_gpu_matmul_quant_rows_scalar_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_quant_tensor(out,
model_map,
model_size,
weight_offset,
weight_type,
in_dim,
out_dim,
x,
n_tok);
}
int ds4_gpu_matmul_q8_0_rows_scalar_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
if (n_tok == 1) {
return ds4_gpu_matmul_q8_0_tensor(out,
model_map,
model_size,
weight_offset,
in_dim,
out_dim,
x,
1);
}
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !x || !model_map ||
n_tok == 0 || n_tok > INT32_MAX ||
(in_dim & 31u) != 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX ||
in_dim > UINT64_MAX / sizeof(float) ||
out_dim > UINT64_MAX / sizeof(float)) {
return 0;
}
const uint64_t x_row_bytes = in_dim * sizeof(float);
const uint64_t out_row_bytes = out_dim * sizeof(float);
if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) ||
(out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t x_bytes = n_tok * x_row_bytes;
const uint64_t out_bytes = n_tok * out_row_bytes;
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul received undersized activation buffers\n");
return 0;
}
const uint64_t blocks = in_dim / 32;
if (blocks > UINT64_MAX / 34u) return 0;
const uint64_t row_bytes = blocks * 34u;
if (out_dim > UINT64_MAX / row_bytes) return 0;
const uint64_t weight_bytes = out_dim * row_bytes;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal Q8_0 scalar-row matmul range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf =
ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner_offset);
if (!wbuf) return 0;
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
if (out_dim > 65536u) mv_dispatch.nsg = 8;
mv_args.nr0 = mv_dispatch.nr0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
const NSUInteger x_base = ds4_gpu_tensor_offset(x);
const NSUInteger out_base = ds4_gpu_tensor_offset(out);
for (uint64_t t = 0; t < n_tok; t++) {
[enc setBuffer:xbuf offset:x_base + (NSUInteger)(t * x_row_bytes) atIndex:2];
[enc setBuffer:outbuf offset:out_base + (NSUInteger)(t * out_row_bytes) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
(NSUInteger)mv_dispatch.nr0,
1,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 scalar-row matvecs")) {
return 0;
}
}
return 1;
}
static int ds4_gpu_matmul_q8_0_pair_impl(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok, bool concat) {
if (concat && (!in_dim || ((out0_dim | out1_dim) & 1u))) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out0 || !out1 || !model_map || !x || n_tok != 1 ||
out0_dim == 0 || out1_dim == 0 || (in_dim & 31u) != 0 ||
in_dim > UINT32_MAX || out0_dim > UINT32_MAX || out1_dim > UINT32_MAX) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> out0buf = ds4_gpu_tensor_buffer(out0);
id<MTLBuffer> out1buf = ds4_gpu_tensor_buffer(out1);
const uint64_t x_bytes = in_dim * sizeof(float);
const uint64_t out0_bytes = out0_dim * sizeof(float);
const uint64_t out1_bytes = out1_dim * sizeof(float);
if (!xbuf || !out0buf || !out1buf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out0) < out0_bytes ||
ds4_gpu_tensor_bytes(out1) < out1_bytes) {
fprintf(stderr, "ds4: Metal paired Q8_0 matvec received undersized activation buffers\n");
return 0;
}
const uint64_t row_bytes = (in_dim / 32u) * 34u;
const uint64_t weight0_bytes = out0_dim * row_bytes;
const uint64_t weight1_bytes = out1_dim * row_bytes;
if (weight0_offset > model_size || weight0_bytes > model_size - weight0_offset ||
weight1_offset > model_size || weight1_bytes > model_size - weight1_offset) {
fprintf(stderr, "ds4: Metal paired Q8_0 matvec range is outside the mapped model\n");
return 0;
}
uint64_t inner0 = 0;
uint64_t inner1 = 0;
id<MTLBuffer> weight0buf =
ds4_gpu_wrap_model_range(model_map, model_size,
weight0_offset, weight0_bytes, &inner0);
id<MTLBuffer> weight1buf =
ds4_gpu_wrap_model_range(model_map, model_size,
weight1_offset, weight1_bytes, &inner1);
if (!weight0buf || !weight1buf) return 0;
ds4_gpu_mv_dispatch dispatch0 = ds4_gpu_make_q8_0_mv_dispatch();
ds4_gpu_mv_dispatch dispatch1 = ds4_gpu_make_q8_0_mv_dispatch();
if (out0_dim > 65536u) dispatch0.nsg = 8;
if (out1_dim > 65536u) dispatch1.nsg = 8;
/* A common threadgroup shape is required to retain each standalone
* reduction tree. Mixed 4/8-simdgroup extents use the existing fallback. */
if (dispatch0.nsg != dispatch1.nsg) return 0;
ds4_gpu_q8_0_matvec_args args0 = ds4_gpu_make_q8_0_mv_args(in_dim, out0_dim);
ds4_gpu_q8_0_matvec_args args1 = ds4_gpu_make_q8_0_mv_args(in_dim, out1_dim);
args0.nr0 = dispatch0.nr0;
args1.nr0 = dispatch1.nr0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(concat ? "kernel_qwen4_q8_concat" : "kernel_mul_mv_q8_0_f32_pair", dispatch0.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args0 length:sizeof(args0) atIndex:0];
[enc setBytes:&args1 length:sizeof(args1) atIndex:1];
[enc setBuffer:weight0buf offset:(NSUInteger)inner0 atIndex:2];
[enc setBuffer:weight1buf offset:(NSUInteger)inner1 atIndex:3];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
[enc setBuffer:out0buf offset:ds4_gpu_tensor_offset(out0) atIndex:5];
[enc setBuffer:out1buf offset:ds4_gpu_tensor_offset(out1) atIndex:6];
[enc setThreadgroupMemoryLength:2u * dispatch0.smem atIndex:0];
const uint64_t max_out_dim = concat ? out0_dim + out1_dim : (out0_dim > out1_dim ? out0_dim : out1_dim);
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)max_out_dim +
(NSUInteger)dispatch0.nr0 - 1u) /
(NSUInteger)dispatch0.nr0,
1,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dispatch0.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "paired Q8_0 matvec")) return 0;
}
return 1;
}
int ds4_gpu_matmul_q8_0_pair_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_q8_0_pair_impl(out0, out1, model_map, model_size, weight0_offset, weight1_offset, in_dim, out0_dim, out1_dim, x, n_tok, false);
}
int ds4_gpu_qwen4_q8_pair_tensor(
ds4_gpu_tensor *out0,
ds4_gpu_tensor *out1,
const void *model_map,
uint64_t model_size,
uint64_t weight0_offset,
uint64_t weight1_offset,
uint64_t in_dim,
uint64_t out0_dim,
uint64_t out1_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_q8_0_pair_impl(out0, out1, model_map, model_size, weight0_offset, weight1_offset, in_dim, out0_dim, out1_dim, x, n_tok, true);
}
int ds4_gpu_matmul_q8_0_f16_out_tensor(
ds4_gpu_tensor *out_h,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
(void)out_h; (void)model_map; (void)model_size; (void)weight_offset;
(void)in_dim; (void)out_dim; (void)x; (void)n_tok;
return 0;
}
static int ds4_gpu_shared_gate_up_swiglu_q8_0_impl(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp,
int store_gate_up, /* 1: gate and up too; 2: mid rounded to bf16 */
uint32_t rows) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!mid || !x || !model_map || !rows || rows > 65535u ||
(store_gate_up == 1 && (!gate || !up)) ||
(in_dim & 31u) != 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX ||
!isfinite(clamp) || clamp < 0.0f) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> gatebuf = store_gate_up == 1 ?
ds4_gpu_tensor_buffer(gate) : midbuf;
id<MTLBuffer> upbuf = store_gate_up == 1 ?
ds4_gpu_tensor_buffer(up) : midbuf;
const uint64_t x_bytes = rows * in_dim * sizeof(float);
const uint64_t out_bytes = rows * out_dim * sizeof(float);
if (!xbuf || !gatebuf || !upbuf || !midbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
(store_gate_up == 1 && ds4_gpu_tensor_bytes(gate) < out_bytes) ||
(store_gate_up == 1 && ds4_gpu_tensor_bytes(up) < out_bytes) ||
ds4_gpu_tensor_bytes(mid) < out_bytes) {
fprintf(stderr, "ds4: Metal shared expert fused gate/up received undersized activation buffers\n");
return 0;
}
const uint64_t blocks = in_dim / 32;
const uint64_t row_bytes = blocks * 34;
const uint64_t weight_bytes = out_dim * row_bytes;
if (gate_offset > model_size || weight_bytes > model_size - gate_offset ||
up_offset > model_size || weight_bytes > model_size - up_offset) {
fprintf(stderr, "ds4: Metal shared expert fused gate/up range is outside the mapped model\n");
return 0;
}
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
id<MTLBuffer> gate_wbuf = ds4_gpu_wrap_model_range(model_map,
model_size,
gate_offset,
weight_bytes,
&gate_inner);
id<MTLBuffer> up_wbuf = ds4_gpu_wrap_model_range(model_map,
model_size,
up_offset,
weight_bytes,
&up_inner);
if (!gate_wbuf || !up_wbuf) return 0;
ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
args.nr0 = mv_dispatch.nr0;
args.ne11 = (int32_t)rows;
args.nb12 = rows * in_dim * sizeof(float);
args.ne1 = (int32_t)rows;
const char *fn_name = store_gate_up == 1 ?
"kernel_dsv4_shared_gate_up_swiglu_q8_0" : store_gate_up == 2 ?
(rows > 1u ? "kernel_dsv4_shared_mid_swiglu_q8_0_bf16_rows" :
"kernel_dsv4_shared_mid_swiglu_q8_0_bf16") :
"kernel_dsv4_shared_mid_swiglu_q8_0";
const int mma_off = ds4_gpu_v41_rows_mma_off();
const bool mma_kernel = !mma_off && store_gate_up == 2 && rows >= 3u && rows <= 8u &&
(in_dim % 256u) == 0 && (out_dim % 8u) == 0;
id<MTLComputePipelineState> pipeline = mma_kernel ?
ds4_gpu_get_pipeline("kernel_dsv4_shared_mid_swiglu_q8_0_bf16_mma_rows8") :
ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1];
[enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:gatebuf offset:(store_gate_up == 1 ?
ds4_gpu_tensor_offset(gate) :
ds4_gpu_tensor_offset(mid)) atIndex:4];
[enc setBuffer:upbuf offset:(store_gate_up == 1 ?
ds4_gpu_tensor_offset(up) :
ds4_gpu_tensor_offset(mid)) atIndex:5];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
[enc setBytes:&clamp length:sizeof(clamp) atIndex:7];
if (mma_kernel) {
[enc setThreadgroupMemoryLength:(8u * 3u * 8u * 32u + 2u * 8u * 64u) * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)out_dim / 8u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "shared expert matrix rows");
}
[enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0];
[enc dispatchThreadgroups:(rows > 1u ?
MTLSizeMake((NSUInteger)rows,
((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, 1) :
MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, 1, 1))
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb,
owned,
store_gate_up == 1 ?
"shared expert fused gate/up" :
"shared expert fused mid")) {
return 0;
}
}
return 1;
}
/* Decode-only fusion of the router logits matvec with the shared-expert
* gate/up SwiGLU: one dispatch instead of two on the same normalized FFN
* input. Bit-exact by construction (see metal/dense.metal). Returns 1 on
* success, 0 when the shape is unsupported (caller falls back), -1 on
* error. */
int ds4_gpu_router_shared_gate_up_q8_0_tensor(
ds4_gpu_tensor *router_logits,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t router_weight_offset,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t router_out_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp,
bool router_only) {
if (!g_initialized && !ds4_gpu_init()) return -1;
if (!router_logits || !gate || !up || !mid || !model_map || !x ||
(in_dim & 31u) != 0 || in_dim > UINT32_MAX ||
router_out_dim > UINT32_MAX || out_dim > UINT32_MAX ||
(router_out_dim & 1u) != 0 || (out_dim & 3u) != 0 ||
!isfinite(clamp) || clamp < 0.0f) {
return -1;
}
if (in_dim != 4096u) return 0;
@autoreleasepool {
const uint64_t router_row_bytes = in_dim * sizeof(uint16_t);
const uint64_t router_weight_bytes = router_row_bytes * router_out_dim;
const uint64_t q8_row_bytes = (in_dim / 32u) * 34u;
const uint64_t q8_weight_bytes = q8_row_bytes * out_dim;
if (router_weight_offset > model_size ||
router_weight_bytes > model_size - router_weight_offset ||
gate_offset > model_size || q8_weight_bytes > model_size - gate_offset ||
up_offset > model_size || q8_weight_bytes > model_size - up_offset) {
return -1;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> rbuf = ds4_gpu_tensor_buffer(router_logits);
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
if (!xbuf || !rbuf || !gatebuf || !upbuf || !midbuf ||
ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(router_logits) < router_out_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(gate) < out_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(up) < out_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(mid) < out_dim * sizeof(float)) {
return -1;
}
uint64_t router_inner = 0, gate_inner = 0, up_inner = 0;
id<MTLBuffer> router_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, router_weight_offset, router_weight_bytes, &router_inner);
id<MTLBuffer> gate_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, q8_weight_bytes, &gate_inner);
id<MTLBuffer> up_wbuf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, q8_weight_bytes, &up_inner);
if (!router_wbuf || !gate_wbuf || !up_wbuf) return -1;
ds4_gpu_f16_matvec_args rargs = ds4_gpu_make_f16_mv_args(in_dim, router_out_dim);
rargs.nr0 = 2;
ds4_gpu_q8_0_matvec_args sargs = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
sargs.nr0 = 2;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_router_shared_gate_up_q8_0");
if (!pipeline) return -1;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return -1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&rargs length:sizeof(rargs) atIndex:0];
[enc setBytes:&sargs length:sizeof(sargs) atIndex:1];
[enc setBuffer:router_wbuf offset:(NSUInteger)router_inner atIndex:2];
[enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:3];
[enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:4];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:5];
[enc setBuffer:rbuf offset:ds4_gpu_tensor_offset(router_logits) atIndex:6];
[enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:7];
[enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:8];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:9];
[enc setBytes:&clamp length:sizeof(clamp) atIndex:10];
[enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0];
const NSUInteger router_tgs = ((NSUInteger)router_out_dim + 1u) / 2u;
const NSUInteger shared_tgs = ((NSUInteger)out_dim / 2u + 1u) / 2u;
[enc dispatchThreadgroups:MTLSizeMake(router_tgs +
(router_only ? 0u : shared_tgs),
1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "router + shared gate/up fused")) {
return -1;
}
}
return 1;
}
int ds4_gpu_router_project_select_fused_tensor(
ds4_gpu_tensor *router_logits,
ds4_gpu_tensor *probs,
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
const void *model_map,
uint64_t model_size,
uint64_t router_weight_offset,
uint64_t bias_offset,
bool has_bias,
const ds4_gpu_tensor *x) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!router_logits || !probs || !selected || !weights || !model_map || !x) {
return 0;
}
@autoreleasepool {
const uint32_t in_dim = 4096u;
const uint32_t n_expert = 256u;
const uint32_t n_used = 6u;
const uint64_t row_bytes = (uint64_t)in_dim * sizeof(uint16_t);
const uint64_t weight_bytes = row_bytes * n_expert;
const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
if (router_weight_offset > model_size ||
weight_bytes > model_size - router_weight_offset ||
(has_bias && (bias_offset > model_size ||
bias_bytes > model_size - bias_offset))) {
return 0;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(router_logits);
id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
if (!xbuf || !logitsbuf || !probsbuf || !selectedbuf || !weightsbuf ||
ds4_gpu_tensor_bytes(x) < (uint64_t)in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(router_logits) < bias_bytes ||
ds4_gpu_tensor_bytes(probs) < bias_bytes ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)n_used * sizeof(int32_t) ||
ds4_gpu_tensor_bytes(weights) < (uint64_t)n_used * sizeof(float)) {
return 0;
}
uint64_t router_inner = 0;
uint64_t bias_inner = 0;
id<MTLBuffer> routerbuf = ds4_gpu_wrap_model_range(
model_map, model_size, router_weight_offset,
weight_bytes, &router_inner);
id<MTLBuffer> biasbuf = has_bias ? ds4_gpu_wrap_model_range(
model_map, model_size, bias_offset, bias_bytes, &bias_inner) : nil;
if (!routerbuf || (has_bias && !biasbuf)) return 0;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
"kernel_dsv4_router_project_select_fused");
if (!pipeline || pipeline.maxTotalThreadsPerThreadgroup < 256u) return 0;
if (!g_dsv4_completion_cache) return 0;
NSString *completion_key = [NSString stringWithFormat:@"router:%p:%llu",
(void *)logitsbuf,
(unsigned long long)ds4_gpu_tensor_offset(router_logits)];
id<MTLBuffer> completion =
[g_dsv4_completion_cache objectForKey:completion_key];
if (!completion) {
completion = [g_device newBufferWithLength:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (!completion) return 0;
*((uint32_t *)[completion contents]) = 0u;
[g_dsv4_completion_cache setObject:completion
forKey:completion_key];
}
[g_transient_buffers addObject:completion];
ds4_gpu_f16_matvec_args rargs =
ds4_gpu_make_f16_mv_args(in_dim, n_expert);
rargs.nr0 = 2;
ds4_gpu_dsv4_router_select_one_args select_args = {
.has_bias = has_bias ? 1u : 0u,
.hash_mode = 0u,
.use_token_buffer = 0u,
.token = 0u,
.hash_rows = 0u,
};
const float zero_f32 = 0.0f;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&rargs length:sizeof(rargs) atIndex:0];
[enc setBytes:&select_args length:sizeof(select_args) atIndex:1];
[enc setBuffer:routerbuf offset:(NSUInteger)router_inner atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(router_logits) atIndex:4];
[enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5];
if (has_bias) {
[enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:6];
} else {
[enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:6];
}
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:7];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:8];
[enc setBuffer:completion offset:0 atIndex:9];
[enc setThreadgroupMemoryLength:1024u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(128, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "router project/select fused")) return -1;
}
return 1;
}
int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp) {
return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate,
up,
mid,
model_map,
model_size,
gate_offset,
up_offset,
in_dim,
out_dim,
x,
clamp,
1, 1u);
}
int ds4_gpu_shared_mid_swiglu_q8_0_tensor(
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp) {
return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(NULL,
NULL,
mid,
model_map,
model_size,
gate_offset,
up_offset,
in_dim,
out_dim,
x,
clamp,
0, 1u);
}
int ds4_gpu_dsv41_shared_swiglu(ds4_gpu_tensor *mid, ds4_gpu_tensor *gate, ds4_gpu_tensor *up,
const void *model_map, uint64_t model_size, uint64_t gate_offset,
uint64_t up_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x, float clamp, uint32_t rows) {
(void)gate; (void)up;
return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(NULL, NULL, mid, model_map, model_size,
gate_offset, up_offset, in_dim, out_dim,
x, clamp, 2, rows);
}
int ds4_gpu_shared_gate_up_swiglu_q8_0_model_view_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
float clamp) {
return ds4_gpu_shared_gate_up_swiglu_q8_0_impl(gate,
up,
mid,
model_map,
model_size,
gate_offset,
up_offset,
in_dim,
out_dim,
x,
clamp,
1, 1u);
}
int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
float clamp) {
if (n_tok == 1) {
return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate,
up,
mid,
model_map,
model_size,
gate_offset,
up_offset,
in_dim,
out_dim,
x,
clamp);
}
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!gate || !up || !mid || !x || !model_map ||
n_tok == 0 ||
(in_dim & 31u) != 0 || (in_dim % 128u) != 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX ||
!isfinite(clamp) || clamp < 0.0f) {
return 0;
}
const uint64_t mv_ext_max_tokens =
ds4_gpu_env_u64("DS4_METAL_Q8_MV_EXT_MAX_TOKENS", 16u, 2u, 128u);
if (n_tok > mv_ext_max_tokens) return 0;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
if (!xbuf || !gatebuf || !upbuf || !midbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(gate) < out_bytes ||
ds4_gpu_tensor_bytes(up) < out_bytes ||
ds4_gpu_tensor_bytes(mid) < out_bytes) {
fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows received undersized activation buffers\n");
return 0;
}
const uint64_t blocks = in_dim / 32;
const uint64_t row_bytes = blocks * 34;
const uint64_t weight_bytes = out_dim * row_bytes;
if (gate_offset > model_size || weight_bytes > model_size - gate_offset ||
up_offset > model_size || weight_bytes > model_size - up_offset) {
fprintf(stderr, "ds4: Metal fused Q8_0 gate/up rows range is outside the mapped model\n");
return 0;
}
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
id<MTLBuffer> gate_wbuf =
ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner);
id<MTLBuffer> up_wbuf =
ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner);
if (!gate_wbuf || !up_wbuf) return 0;
const int16_t nsg = 2;
const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok);
const char *fn_name = ds4_gpu_mv_ext_q8_pair_swiglu_name(r1ptg);
id<MTLComputePipelineState> pipeline =
fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
if (!pipeline) return 0;
const int16_t nypsg = 32 / nxpsg;
const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
ds4_gpu_mul_mv_ext_args args =
ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, 34, row_bytes);
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1];
[enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4];
[enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
[enc setBytes:&clamp length:sizeof(clamp) atIndex:7];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) /
(NSUInteger)r0ptg,
((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) /
(NSUInteger)r1ptg,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "fused Q8_0 gate/up rows")) {
return 0;
}
}
return 1;
}
int ds4_gpu_shared_gate_up_swiglu_q8_0_rows_scalar_tensor(
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
float clamp) {
if (n_tok == 1) {
return ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(gate,
up,
mid,
model_map,
model_size,
gate_offset,
up_offset,
in_dim,
out_dim,
x,
clamp);
}
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!gate || !up || !mid || !x || !model_map ||
n_tok == 0 ||
(in_dim & 31u) != 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX ||
in_dim > UINT64_MAX / sizeof(float) ||
out_dim > UINT64_MAX / sizeof(float) ||
!isfinite(clamp) || clamp < 0.0f) {
return 0;
}
const uint64_t x_row_bytes = in_dim * sizeof(float);
const uint64_t out_row_bytes = out_dim * sizeof(float);
if ((x_row_bytes != 0 && n_tok > UINT64_MAX / x_row_bytes) ||
(out_row_bytes != 0 && n_tok > UINT64_MAX / out_row_bytes)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
const uint64_t x_bytes = n_tok * x_row_bytes;
const uint64_t out_bytes = n_tok * out_row_bytes;
if (!xbuf || !gatebuf || !upbuf || !midbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(gate) < out_bytes ||
ds4_gpu_tensor_bytes(up) < out_bytes ||
ds4_gpu_tensor_bytes(mid) < out_bytes) {
fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up received undersized activation buffers\n");
return 0;
}
const uint64_t blocks = in_dim / 32;
if (blocks > UINT64_MAX / 34u) return 0;
const uint64_t row_bytes = blocks * 34u;
if (out_dim > UINT64_MAX / row_bytes) return 0;
const uint64_t weight_bytes = out_dim * row_bytes;
if (gate_offset > model_size || weight_bytes > model_size - gate_offset ||
up_offset > model_size || weight_bytes > model_size - up_offset) {
fprintf(stderr, "ds4: Metal shared expert scalar-row fused gate/up range is outside the mapped model\n");
return 0;
}
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
id<MTLBuffer> gate_wbuf =
ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, weight_bytes, &gate_inner);
id<MTLBuffer> up_wbuf =
ds4_gpu_wrap_model_range(model_map, model_size, up_offset, weight_bytes, &up_inner);
if (!gate_wbuf || !up_wbuf) return 0;
ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
args.nr0 = mv_dispatch.nr0;
args.ne11 = (int32_t)n_tok;
args.nb12 = n_tok * x_row_bytes;
args.nb13 = args.nb12;
args.ne1 = (int32_t)n_tok;
const char *fn_name = "kernel_dsv4_shared_gate_up_swiglu_q8_0";
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:gate_wbuf offset:(NSUInteger)gate_inner atIndex:1];
[enc setBuffer:up_wbuf offset:(NSUInteger)up_inner atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:4];
[enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:5];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
[enc setBytes:&clamp length:sizeof(clamp) atIndex:7];
[enc setThreadgroupMemoryLength:2u * mv_dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(
((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
(NSUInteger)mv_dispatch.nr0,
(NSUInteger)n_tok,
1)
threadsPerThreadgroup:
MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "shared expert scalar-row fused gate/up")) {
return 0;
}
}
return 1;
}
static int ds4_gpu_matmul_f16_tensor_impl(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok,
bool exact_rows,
int round_bf16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal F16 tensor matmul received undersized activation buffers\n");
return 0;
}
const uint64_t row_bytes = in_dim * sizeof(uint16_t);
const uint64_t weight_bytes = row_bytes * out_dim;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal F16 tensor matmul range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (n_tok == 1 || exact_rows) {
ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim);
mv_args.ne11 = mv_args.ne1 = (int32_t)n_tok;
mv_args.nb12 = mv_args.nb13 = n_tok * in_dim * sizeof(float);
ds4_gpu_mv_dispatch mv_dispatch =
ds4_gpu_make_plain_mv_dispatch(in_dim, 0);
if (!g_quality_mode && (out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) {
mv_dispatch.nr0 = 4;
mv_dispatch.smem = 32u * 4u * sizeof(float);
}
/* One row per SIMD group doubles the threadgroup count of narrow
* projections; every row keeps its K walk and reduction tree. */
if (ds4_gpu_plain_mv_single_row(out_dim)) {
mv_dispatch.nr0 = 1;
mv_dispatch.smem = 32u * sizeof(float);
}
mv_args.nr0 = mv_dispatch.nr0;
/* a few rows share the weight pass: one simdgroup group per row,
* at most 32 simdgroups per threadgroup, so more rows go in chunks */
const uint32_t chunk = mv_dispatch.smem && (in_dim % 4u) == 0 && mv_dispatch.nsg >= 1 && mv_dispatch.nsg <= 16 ?
32u / (uint32_t)mv_dispatch.nsg : 1u;
const bool rows_kernel = n_tok >= 2 && n_tok <= 8 && chunk >= 2u;
const char *fn_name = rows_kernel ?
(round_bf16 ? "kernel_mul_mv_f16_f32_4_rows_bf16" : "kernel_mul_mv_f16_f32_4_rows") : mv_dispatch.function_name;
if (round_bf16 && !rows_kernel) { /* only the vectorized kernels round in place */
if (strcmp(fn_name, "kernel_mul_mv_f16_f32_4") != 0) return 0;
fn_name = "kernel_mul_mv_f16_f32_4_bf16";
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(fn_name, mv_dispatch.nsg);
if (!pipeline) return 0;
for (uint64_t r = 0; r < n_tok; r += rows_kernel ? chunk : n_tok) {
const NSUInteger groups = rows_kernel ? (NSUInteger)(n_tok - r < chunk ? n_tok - r : chunk) : 1u;
if (rows_kernel) mv_args.ne11 = mv_args.ne1 = (int32_t)groups;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) + (NSUInteger)(r * in_dim * sizeof(float)) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) + (NSUInteger)(r * out_dim * sizeof(float)) atIndex:3];
if (mv_dispatch.smem) {
[enc setThreadgroupMemoryLength:mv_dispatch.smem * groups atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
rows_kernel ? 1u : (NSUInteger)n_tok,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg * groups, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matvec")) return 0;
return 1;
}
if (n_tok <= 8 && (in_dim % 128u) == 0) {
const int16_t nsg = ds4_gpu_mv_ext_nsg();
const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok);
const char *fn_name = ds4_gpu_mv_ext_name(0, r1ptg);
id<MTLComputePipelineState> pipeline =
fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
if (!pipeline) return 0;
const int16_t nypsg = 32 / nxpsg;
const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
ds4_gpu_mul_mv_ext_args args =
ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(uint16_t), row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg,
((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor mul_mv_ext")) return 0;
return 1;
}
/*
* Same direct-RHS TensorOps structure as Q8_0, but for F16 model
* matrices. The 128-token RHS tile is kept when the batch alignment
* allows it because the later tile_n=64 retest was neutral/slower.
*/
if (ds4_gpu_mpp_available() &&
n_tok >= 32u &&
(in_dim % 32u) == 0 &&
(out_dim % 64u) == 0 &&
(n_tok % 32u) == 0) {
uint64_t nax_tile_n = 32u;
if ((n_tok % 128u) == 0) {
nax_tile_n = 128u;
} else if ((n_tok % 64u) == 0) {
nax_tile_n = 64u;
}
const char *nax_fn = nax_tile_n == 128u
? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n128"
: (nax_tile_n == 64u
? "kernel_mul_mm_f16_f32_mpp_direct_rhs_n64"
: "kernel_mul_mm_f16_f32_mpp_direct_rhs");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mm_pipeline(nax_fn, false, false);
if (pipeline) {
ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_tok / nax_tile_n),
(NSUInteger)out_dim / 64u,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 NAX tensor matmul")) {
return 0;
}
return 1;
}
ds4_gpu_warn_mpp_fallback();
}
const bool bc_inp = (in_dim % 32u) != 0;
const bool bc_out = (out_dim % 64u) != 0 || (n_tok % 32u) != 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_f16_f32", bc_inp, bc_out);
if (!pipeline) return 0;
ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u) / 32u,
((NSUInteger)out_dim + 63u) / 64u,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 tensor matmul")) return 0;
}
return 1;
}
int ds4_gpu_matmul_f16_tensor(ds4_gpu_tensor *out,
const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint64_t in_dim,
uint64_t out_dim, const ds4_gpu_tensor *x,
uint64_t n_tok) {
return ds4_gpu_matmul_f16_tensor_impl(out, model_map, model_size,
weight_offset, in_dim, out_dim, x, n_tok, false, 0);
}
int ds4_gpu_dsv41_project_f16_bf16(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x) {
return ds4_gpu_matmul_f16_tensor_impl(out, model_map, model_size, weight_offset,
in_dim, out_dim, x, 1, false, 1);
}
int ds4_gpu_dsv41_projection_rows(ds4_gpu_tensor *out,
const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t width,
uint32_t outputs, uint32_t rows,
const ds4_gpu_tensor *in) {
if (!width || !outputs || !rows || rows > 8192u || !model_map) return 0;
return ds4_gpu_matmul_f16_tensor_impl(out, model_map, model_size,
weight_offset, width, outputs, in, rows, true, 0);
}
int ds4_gpu_matmul_f16_pair_tensor(
ds4_gpu_tensor *out_a,
ds4_gpu_tensor *out_b,
const void *model_map,
uint64_t model_size,
uint64_t weight_a_offset,
uint64_t weight_b_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok != 1 || (in_dim & 3u) != 0) return 0;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outabuf = ds4_gpu_tensor_buffer(out_a);
id<MTLBuffer> outbbuf = ds4_gpu_tensor_buffer(out_b);
const uint64_t x_bytes = in_dim * sizeof(float);
const uint64_t out_bytes = out_dim * sizeof(float);
if (!xbuf || !outabuf || !outbbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out_a) < out_bytes ||
ds4_gpu_tensor_bytes(out_b) < out_bytes) {
fprintf(stderr, "ds4: Metal F16 paired matvec received undersized activation buffers\n");
return 0;
}
const uint64_t row_bytes = in_dim * sizeof(uint16_t);
const uint64_t weight_bytes = row_bytes * out_dim;
if (weight_a_offset > model_size || weight_bytes > model_size - weight_a_offset ||
weight_b_offset > model_size || weight_bytes > model_size - weight_b_offset) {
fprintf(stderr, "ds4: Metal F16 paired matvec range is outside the mapped model\n");
return 0;
}
uint64_t inner_a = 0;
uint64_t inner_b = 0;
id<MTLBuffer> wabuf = ds4_gpu_wrap_model_range(model_map, model_size,
weight_a_offset, weight_bytes,
&inner_a);
id<MTLBuffer> wbbuf = ds4_gpu_wrap_model_range(model_map, model_size,
weight_b_offset, weight_bytes,
&inner_b);
if (!wabuf || !wbbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 0);
if (ds4_gpu_use_compressor_pair_nr4() &&
(out_dim == 512u || out_dim == 1024u) && in_dim >= 4096u) {
mv_dispatch.nr0 = 4;
mv_dispatch.smem = 32u * 4u * sizeof(float);
}
mv_args.nr0 = mv_dispatch.nr0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_f16_f32_pair_4", mv_dispatch.nsg);
if (!pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wabuf offset:(NSUInteger)inner_a atIndex:1];
[enc setBuffer:wbbuf offset:(NSUInteger)inner_b atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:outabuf offset:ds4_gpu_tensor_offset(out_a) atIndex:4];
[enc setBuffer:outbbuf offset:ds4_gpu_tensor_offset(out_b) atIndex:5];
if (mv_dispatch.smem) {
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
1,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "F16 paired matvec")) return 0;
}
return 1;
}
int ds4_gpu_matmul_f16_pair_compressor_store_tensor(
ds4_gpu_tensor *out_kv,
ds4_gpu_tensor *out_score,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t weight_kv_offset,
uint64_t weight_score_offset,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t in_dim,
uint32_t width,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos) {
if (!g_initialized && !ds4_gpu_init()) return -1;
if ((g_quality_mode ||
(!ds4_gpu_device_name_contains("M3") &&
!ds4_gpu_device_name_contains("M5"))) ||
getenv("DS4_METAL_DISABLE_COMPRESSOR_PAIR_PROJ") != NULL ||
getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") != NULL) {
return 0;
}
if (!out_kv || !out_score || !state_kv || !state_score || !model_map || !x ||
in_dim == 0 || width == 0 || ratio == 0 ||
(ape_type != 0u && ape_type != 1u)) {
return -1;
}
if (in_dim != 4096u ||
(width != 256u && width != 512u && width != 1024u) ||
(ratio != 4u && ratio != 128u)) {
return 0;
}
@autoreleasepool {
const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio;
const uint64_t row_bytes = in_dim * sizeof(uint16_t);
const uint64_t weight_bytes = row_bytes * width;
const uint64_t out_bytes = (uint64_t)width * sizeof(float);
const uint64_t state_bytes =
(uint64_t)state_rows * width * sizeof(float);
const uint64_t ape_elem = ape_type == 1u ? sizeof(uint16_t) : sizeof(float);
const uint64_t ape_bytes = (uint64_t)ratio * width * ape_elem;
if (weight_kv_offset > model_size ||
weight_bytes > model_size - weight_kv_offset ||
weight_score_offset > model_size ||
weight_bytes > model_size - weight_score_offset ||
ape_offset > model_size || ape_bytes > model_size - ape_offset) {
return -1;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outkvbuf = ds4_gpu_tensor_buffer(out_kv);
id<MTLBuffer> outscorebuf = ds4_gpu_tensor_buffer(out_score);
id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
if (!xbuf || !outkvbuf || !outscorebuf || !statekvbuf || !statescbuf ||
ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out_kv) < out_bytes ||
ds4_gpu_tensor_bytes(out_score) < out_bytes ||
ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
ds4_gpu_tensor_bytes(state_score) < state_bytes) {
return -1;
}
uint64_t weight_kv_inner = 0;
uint64_t weight_score_inner = 0;
uint64_t ape_inner = 0;
id<MTLBuffer> weightkvbuf = ds4_gpu_wrap_model_range(
model_map, model_size, weight_kv_offset, weight_bytes,
&weight_kv_inner);
id<MTLBuffer> weightscorebuf = ds4_gpu_wrap_model_range(
model_map, model_size, weight_score_offset, weight_bytes,
&weight_score_inner);
id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(
model_map, model_size, ape_offset, ape_bytes, &ape_inner);
if (!weightkvbuf || !weightscorebuf || !apebuf) return -1;
ds4_gpu_f16_matvec_args mv_args =
ds4_gpu_make_f16_mv_args(in_dim, width);
ds4_gpu_mv_dispatch mv_dispatch =
ds4_gpu_make_plain_mv_dispatch(in_dim, 0);
if (ds4_gpu_use_compressor_pair_nr4() &&
(width == 512u || width == 1024u) && in_dim >= 4096u) {
mv_dispatch.nr0 = 4;
mv_dispatch.smem = 32u * 4u * sizeof(float);
}
mv_args.nr0 = mv_dispatch.nr0;
ds4_gpu_dsv4_compressor_store_one_args store_args = {
.width = width,
.ratio = ratio,
.pos = pos,
.ape_type = ape_type,
};
id<MTLComputePipelineState> pipeline = ds4_gpu_get_mul_mv_pipeline(
"kernel_mul_mv_f16_f32_pair_compressor_store_4",
mv_dispatch.nsg);
if (!pipeline) return -1;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return -1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBytes:&store_args length:sizeof(store_args) atIndex:1];
[enc setBuffer:weightkvbuf offset:(NSUInteger)weight_kv_inner atIndex:2];
[enc setBuffer:weightscorebuf offset:(NSUInteger)weight_score_inner atIndex:3];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
[enc setBuffer:outkvbuf offset:ds4_gpu_tensor_offset(out_kv) atIndex:5];
[enc setBuffer:outscorebuf offset:ds4_gpu_tensor_offset(out_score) atIndex:6];
[enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:7];
[enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:8];
[enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:9];
if (mv_dispatch.smem) {
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(
((NSUInteger)width + (NSUInteger)mv_dispatch.nr0 - 1u) /
(NSUInteger)mv_dispatch.nr0,
1,
1)
threadsPerThreadgroup:MTLSizeMake(
32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "F16 paired matvec compressor state store")) {
return -1;
}
}
return 1;
}
/* Quad variant of the paired compressor projection: the attention compressor
* and indexer compressor pairs share the input activation and F16 matvec
* shape, so one dispatch covers all four matrices. Bit-exact by
* construction (see metal/dense.metal). Returns 1 when the fused dispatch
* ran, 0 when the caller should use the separate paths, -1 on error. */
int ds4_gpu_matmul_f16_quad_compressor_store_tensor(
ds4_gpu_tensor *out0_kv,
ds4_gpu_tensor *out0_score,
ds4_gpu_tensor *out1_kv,
ds4_gpu_tensor *out1_score,
ds4_gpu_tensor *state0_kv,
ds4_gpu_tensor *state0_score,
ds4_gpu_tensor *state1_kv,
ds4_gpu_tensor *state1_score,
const void *model_map,
uint64_t model_size,
uint64_t weight0_kv_offset,
uint64_t weight0_score_offset,
uint64_t weight1_kv_offset,
uint64_t weight1_score_offset,
uint64_t ape0_offset,
uint32_t ape0_type,
uint64_t ape1_offset,
uint32_t ape1_type,
uint64_t in_dim,
uint32_t width0,
uint32_t width1,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos) {
if (!g_initialized && !ds4_gpu_init()) return -1;
if (!out0_kv || !out0_score || !out1_kv || !out1_score ||
!state0_kv || !state0_score || !state1_kv || !state1_score ||
!model_map || !x ||
(ape0_type != 0u && ape0_type != 1u) ||
(ape1_type != 0u && ape1_type != 1u)) {
return -1;
}
if (in_dim != 4096u ||
(width0 != 512u && width0 != 1024u) ||
width1 != 256u || ratio != 4u) {
return 0;
}
@autoreleasepool {
const uint32_t state_rows = 2u * ratio;
const uint64_t row_bytes = in_dim * sizeof(uint16_t);
const uint64_t weight0_bytes = row_bytes * width0;
const uint64_t weight1_bytes = row_bytes * width1;
const uint64_t ape0_elem = ape0_type == 1u ? sizeof(uint16_t) : sizeof(float);
const uint64_t ape1_elem = ape1_type == 1u ? sizeof(uint16_t) : sizeof(float);
const uint64_t ape0_bytes = (uint64_t)ratio * width0 * ape0_elem;
const uint64_t ape1_bytes = (uint64_t)ratio * width1 * ape1_elem;
if (weight0_kv_offset > model_size || weight0_bytes > model_size - weight0_kv_offset ||
weight0_score_offset > model_size || weight0_bytes > model_size - weight0_score_offset ||
weight1_kv_offset > model_size || weight1_bytes > model_size - weight1_kv_offset ||
weight1_score_offset > model_size || weight1_bytes > model_size - weight1_score_offset ||
ape0_offset > model_size || ape0_bytes > model_size - ape0_offset ||
ape1_offset > model_size || ape1_bytes > model_size - ape1_offset) {
return -1;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> out0kvbuf = ds4_gpu_tensor_buffer(out0_kv);
id<MTLBuffer> out0scbuf = ds4_gpu_tensor_buffer(out0_score);
id<MTLBuffer> out1kvbuf = ds4_gpu_tensor_buffer(out1_kv);
id<MTLBuffer> out1scbuf = ds4_gpu_tensor_buffer(out1_score);
id<MTLBuffer> state0kvbuf = ds4_gpu_tensor_buffer(state0_kv);
id<MTLBuffer> state0scbuf = ds4_gpu_tensor_buffer(state0_score);
id<MTLBuffer> state1kvbuf = ds4_gpu_tensor_buffer(state1_kv);
id<MTLBuffer> state1scbuf = ds4_gpu_tensor_buffer(state1_score);
if (!xbuf || !out0kvbuf || !out0scbuf || !out1kvbuf || !out1scbuf ||
!state0kvbuf || !state0scbuf || !state1kvbuf || !state1scbuf ||
ds4_gpu_tensor_bytes(x) < in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out0_kv) < (uint64_t)width0 * sizeof(float) ||
ds4_gpu_tensor_bytes(out0_score) < (uint64_t)width0 * sizeof(float) ||
ds4_gpu_tensor_bytes(out1_kv) < (uint64_t)width1 * sizeof(float) ||
ds4_gpu_tensor_bytes(out1_score) < (uint64_t)width1 * sizeof(float) ||
ds4_gpu_tensor_bytes(state0_kv) < (uint64_t)state_rows * width0 * sizeof(float) ||
ds4_gpu_tensor_bytes(state0_score) < (uint64_t)state_rows * width0 * sizeof(float) ||
ds4_gpu_tensor_bytes(state1_kv) < (uint64_t)state_rows * width1 * sizeof(float) ||
ds4_gpu_tensor_bytes(state1_score) < (uint64_t)state_rows * width1 * sizeof(float)) {
return -1;
}
uint64_t w0kv_inner = 0, w0sc_inner = 0, w1kv_inner = 0, w1sc_inner = 0;
uint64_t ape0_inner = 0, ape1_inner = 0;
id<MTLBuffer> w0kvbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight0_kv_offset, weight0_bytes, &w0kv_inner);
id<MTLBuffer> w0scbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight0_score_offset, weight0_bytes, &w0sc_inner);
id<MTLBuffer> w1kvbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight1_kv_offset, weight1_bytes, &w1kv_inner);
id<MTLBuffer> w1scbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight1_score_offset, weight1_bytes, &w1sc_inner);
id<MTLBuffer> ape0buf = ds4_gpu_wrap_model_range(model_map, model_size, ape0_offset, ape0_bytes, &ape0_inner);
id<MTLBuffer> ape1buf = ds4_gpu_wrap_model_range(model_map, model_size, ape1_offset, ape1_bytes, &ape1_inner);
if (!w0kvbuf || !w0scbuf || !w1kvbuf || !w1scbuf || !ape0buf || !ape1buf) return -1;
ds4_gpu_f16_matvec_args mv_args = ds4_gpu_make_f16_mv_args(in_dim, width0);
mv_args.nr0 = 2;
ds4_gpu_dsv4_compressor_store_one_args store0_args = {
.width = width0,
.ratio = ratio,
.pos = pos,
.ape_type = ape0_type,
};
ds4_gpu_dsv4_compressor_store_one_args store1_args = {
.width = width1,
.ratio = ratio,
.pos = pos,
.ape_type = ape1_type,
};
id<MTLComputePipelineState> pipeline = ds4_gpu_get_mul_mv_pipeline(
"kernel_mul_mv_f16_f32_quad_compressor_store_4", 8);
if (!pipeline) return -1;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return -1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBytes:&store0_args length:sizeof(store0_args) atIndex:1];
[enc setBytes:&store1_args length:sizeof(store1_args) atIndex:2];
[enc setBuffer:w0kvbuf offset:(NSUInteger)w0kv_inner atIndex:3];
[enc setBuffer:w0scbuf offset:(NSUInteger)w0sc_inner atIndex:4];
[enc setBuffer:w1kvbuf offset:(NSUInteger)w1kv_inner atIndex:5];
[enc setBuffer:w1scbuf offset:(NSUInteger)w1sc_inner atIndex:6];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:7];
[enc setBuffer:out0kvbuf offset:ds4_gpu_tensor_offset(out0_kv) atIndex:8];
[enc setBuffer:out0scbuf offset:ds4_gpu_tensor_offset(out0_score) atIndex:9];
[enc setBuffer:out1kvbuf offset:(NSUInteger)ds4_gpu_tensor_offset(out1_kv) atIndex:10];
[enc setBuffer:out1scbuf offset:(NSUInteger)ds4_gpu_tensor_offset(out1_score) atIndex:11];
[enc setBuffer:ape0buf offset:(NSUInteger)ape0_inner atIndex:12];
[enc setBuffer:ape1buf offset:(NSUInteger)ape1_inner atIndex:13];
[enc setBuffer:state0kvbuf offset:ds4_gpu_tensor_offset(state0_kv) atIndex:14];
[enc setBuffer:state0scbuf offset:ds4_gpu_tensor_offset(state0_score) atIndex:15];
[enc setBuffer:state1kvbuf offset:ds4_gpu_tensor_offset(state1_kv) atIndex:16];
[enc setBuffer:state1scbuf offset:ds4_gpu_tensor_offset(state1_score) atIndex:17];
[enc setThreadgroupMemoryLength:32u * 2u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(
((NSUInteger)width0 + (NSUInteger)width1 + 1u) / 2u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "F16 quad matvec compressor state store")) {
return -1;
}
}
return 1;
}
/* Decode-only emit-path fusion: finalize the freshly pooled attention and
* indexer compressor rows (norm + rope + fp8/commit + qat) in one dispatch
* instead of seven. Bit-exact vs the separate dispatches; see the kernel
* comment. Returns 1 when the fused dispatch ran, 0 to fall back. */
int ds4_gpu_dsv4_comp_row_finalize_tensor(
ds4_gpu_tensor *attn_stage,
ds4_gpu_tensor *attn_cache,
uint32_t attn_comp_row,
uint64_t attn_norm_offset,
ds4_gpu_tensor *index_cache,
uint32_t index_comp_row,
uint64_t index_norm_offset,
ds4_gpu_tensor *attn_state_kv,
ds4_gpu_tensor *attn_state_score,
ds4_gpu_tensor *index_state_kv,
ds4_gpu_tensor *index_state_score,
const void *model_map,
uint64_t model_size,
uint32_t pos,
uint32_t n_rot,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!attn_stage || !attn_cache || !index_cache || !model_map) return 0;
@autoreleasepool {
const uint64_t attn_norm_bytes = 512u * sizeof(float);
const uint64_t index_norm_bytes = 128u * sizeof(float);
if (attn_norm_offset > model_size ||
attn_norm_bytes > model_size - attn_norm_offset ||
index_norm_offset > model_size ||
index_norm_bytes > model_size - index_norm_offset) {
return 0;
}
if (!attn_state_kv || !attn_state_score || !index_state_kv ||
!index_state_score) return 0;
if (ds4_gpu_tensor_bytes(attn_state_kv) < 8u * 1024u * sizeof(float) ||
ds4_gpu_tensor_bytes(attn_state_score) < 8u * 1024u * sizeof(float) ||
ds4_gpu_tensor_bytes(index_state_kv) < 8u * 256u * sizeof(float) ||
ds4_gpu_tensor_bytes(index_state_score) < 8u * 256u * sizeof(float)) {
return 0;
}
if (ds4_gpu_tensor_bytes(attn_stage) < 512u * sizeof(float) ||
ds4_gpu_tensor_bytes(index_cache) <
((uint64_t)index_comp_row + 1u) * 128u * sizeof(float) ||
ds4_gpu_tensor_bytes(attn_cache) <
((uint64_t)attn_comp_row + 1u) * 512u * sizeof(uint16_t)) {
return 0;
}
uint64_t attn_norm_inner = 0, index_norm_inner = 0;
id<MTLBuffer> anbuf = ds4_gpu_wrap_model_range(
model_map, model_size, attn_norm_offset, attn_norm_bytes, &attn_norm_inner);
id<MTLBuffer> inbuf = ds4_gpu_wrap_model_range(
model_map, model_size, index_norm_offset, index_norm_bytes, &index_norm_inner);
if (!anbuf || !inbuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_comp_row_finalize_f32");
if (!pipeline) return 0;
ds4_gpu_comp_finalize_args fin_args = {
.rope = {
.row_bytes = 0,
.token_bytes = 0,
.head_dim = 0,
.n_dims = (int32_t)n_rot,
.n_ctx_orig = (int32_t)n_ctx_orig,
.inverse = 0,
.pos0 = pos,
.pos_step = 0,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
},
.rms_eps = rms_eps,
.pad0 = 0,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&fin_args length:sizeof(fin_args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(attn_stage)
offset:ds4_gpu_tensor_offset(attn_stage) atIndex:1];
[enc setBuffer:anbuf offset:(NSUInteger)attn_norm_inner atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(attn_cache)
offset:ds4_gpu_tensor_offset(attn_cache) +
(NSUInteger)attn_comp_row * 512u * sizeof(uint16_t) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(index_cache)
offset:ds4_gpu_tensor_offset(index_cache) +
(NSUInteger)index_comp_row * 128u * sizeof(float) atIndex:4];
[enc setBuffer:inbuf offset:(NSUInteger)index_norm_inner atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(attn_state_kv)
offset:ds4_gpu_tensor_offset(attn_state_kv) atIndex:6];
[enc setBuffer:ds4_gpu_tensor_buffer(attn_state_score)
offset:ds4_gpu_tensor_offset(attn_state_score) atIndex:7];
[enc setBuffer:ds4_gpu_tensor_buffer(index_state_kv)
offset:ds4_gpu_tensor_offset(index_state_kv) atIndex:8];
[enc setBuffer:ds4_gpu_tensor_buffer(index_state_score)
offset:ds4_gpu_tensor_offset(index_state_score) atIndex:9];
[enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(2u + (4u * 1024u + 4u * 256u + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "compressor row finalize")) {
return 0;
}
}
return 1;
}
/* Decode-only fusion: q_a/kv Q8 pair projection + F16 quad compressor
* projection/store in one dispatch (both read the same normalized input).
* Bit-exact vs the separate dispatches; see the kernel comment. Returns
* 1 when the fused dispatch ran, 0 when the caller must use the separate
* paths, -1 on error. */
int ds4_gpu_qkv_pair_quad_compressor_store_tensor(
ds4_gpu_tensor *qr,
ds4_gpu_tensor *kv_raw,
ds4_gpu_tensor *out0_kv,
ds4_gpu_tensor *out0_score,
ds4_gpu_tensor *out1_kv,
ds4_gpu_tensor *out1_score,
ds4_gpu_tensor *state0_kv,
ds4_gpu_tensor *state0_score,
ds4_gpu_tensor *state1_kv,
ds4_gpu_tensor *state1_score,
const void *model_map,
uint64_t model_size,
uint64_t q_a_offset,
uint64_t kv_offset,
uint64_t weight0_kv_offset,
uint64_t weight0_score_offset,
uint64_t weight1_kv_offset,
uint64_t weight1_score_offset,
uint64_t ape0_offset,
uint32_t ape0_type,
uint64_t ape1_offset,
uint32_t ape1_type,
uint32_t in_dim,
uint32_t q_rank,
uint32_t kv_dim,
uint32_t width0,
uint32_t width1,
const ds4_gpu_tensor *x,
uint32_t ratio,
uint32_t pos) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!qr || !kv_raw || !out0_kv || !out0_score || !out1_kv || !out1_score ||
!state0_kv || !state0_score || !state1_kv || !state1_score ||
!model_map || !x || ratio == 0u || (in_dim & 31u) != 0 ||
(width0 & 1u) != 0 || (width1 & 1u) != 0) {
return 0;
}
@autoreleasepool {
const uint64_t q8_row_bytes = (in_dim / 32u) * 34u;
const uint64_t f16_row_bytes = (uint64_t)in_dim * sizeof(uint16_t);
const uint64_t q_a_bytes = (uint64_t)q_rank * q8_row_bytes;
const uint64_t kv_bytes = (uint64_t)kv_dim * q8_row_bytes;
const uint64_t weight0_bytes = (uint64_t)width0 * f16_row_bytes;
const uint64_t weight1_bytes = (uint64_t)width1 * f16_row_bytes;
const uint64_t state_rows = ratio == 4u ? 2u * ratio : ratio;
const uint64_t ape0_bytes = (uint64_t)width0 * ratio *
(ape0_type == 1u ? 2u : 4u);
const uint64_t ape1_bytes = (uint64_t)width1 * ratio *
(ape1_type == 1u ? 2u : 4u);
if (q_a_offset > model_size || q_a_bytes > model_size - q_a_offset ||
kv_offset > model_size || kv_bytes > model_size - kv_offset ||
weight0_kv_offset > model_size ||
weight0_bytes > model_size - weight0_kv_offset ||
weight0_score_offset > model_size ||
weight0_bytes > model_size - weight0_score_offset ||
weight1_kv_offset > model_size ||
weight1_bytes > model_size - weight1_kv_offset ||
weight1_score_offset > model_size ||
weight1_bytes > model_size - weight1_score_offset ||
ape0_offset > model_size || ape0_bytes > model_size - ape0_offset ||
ape1_offset > model_size || ape1_bytes > model_size - ape1_offset) {
return -1;
}
if (ds4_gpu_tensor_bytes(qr) < (uint64_t)q_rank * sizeof(float) ||
ds4_gpu_tensor_bytes(kv_raw) < (uint64_t)kv_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out0_kv) < (uint64_t)width0 * sizeof(float) ||
ds4_gpu_tensor_bytes(out0_score) < (uint64_t)width0 * sizeof(float) ||
ds4_gpu_tensor_bytes(out1_kv) < (uint64_t)width1 * sizeof(float) ||
ds4_gpu_tensor_bytes(out1_score) < (uint64_t)width1 * sizeof(float) ||
ds4_gpu_tensor_bytes(x) < (uint64_t)in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(state0_kv) < state_rows * width0 * sizeof(float) ||
ds4_gpu_tensor_bytes(state0_score) < state_rows * width0 * sizeof(float) ||
ds4_gpu_tensor_bytes(state1_kv) < state_rows * width1 * sizeof(float) ||
ds4_gpu_tensor_bytes(state1_score) < state_rows * width1 * sizeof(float)) {
return -1;
}
uint64_t q_a_inner = 0, kv_inner = 0;
uint64_t w0kv_inner = 0, w0sc_inner = 0, w1kv_inner = 0, w1sc_inner = 0;
uint64_t ape0_inner = 0, ape1_inner = 0;
id<MTLBuffer> qw0buf = ds4_gpu_wrap_model_range(model_map, model_size, q_a_offset, q_a_bytes, &q_a_inner);
id<MTLBuffer> qw1buf = ds4_gpu_wrap_model_range(model_map, model_size, kv_offset, kv_bytes, &kv_inner);
id<MTLBuffer> w0kvbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight0_kv_offset, weight0_bytes, &w0kv_inner);
id<MTLBuffer> w0scbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight0_score_offset, weight0_bytes, &w0sc_inner);
id<MTLBuffer> w1kvbuf = width1 != 0u
? ds4_gpu_wrap_model_range(model_map, model_size, weight1_kv_offset,
weight1_bytes, &w1kv_inner)
: w0kvbuf;
id<MTLBuffer> w1scbuf = width1 != 0u
? ds4_gpu_wrap_model_range(model_map, model_size, weight1_score_offset,
weight1_bytes, &w1sc_inner)
: w0scbuf;
id<MTLBuffer> ape0buf = ds4_gpu_wrap_model_range(model_map, model_size, ape0_offset, ape0_bytes, &ape0_inner);
id<MTLBuffer> ape1buf = width1 != 0u
? ds4_gpu_wrap_model_range(model_map, model_size, ape1_offset,
ape1_bytes, &ape1_inner)
: ape0buf;
if (width1 == 0u) {
w1kv_inner = w0kv_inner;
w1sc_inner = w0sc_inner;
ape1_inner = ape0_inner;
}
if (!qw0buf || !qw1buf || !w0kvbuf || !w0scbuf ||
!w1kvbuf || !w1scbuf || !ape0buf || !ape1buf) return -1;
ds4_gpu_q8_0_matvec_args args0 = ds4_gpu_make_q8_0_mv_args(in_dim, q_rank);
ds4_gpu_q8_0_matvec_args args1 = ds4_gpu_make_q8_0_mv_args(in_dim, kv_dim);
args0.nr0 = 2;
args1.nr0 = 2;
ds4_gpu_f16_matvec_args cargs = ds4_gpu_make_f16_mv_args(in_dim, width0);
cargs.nr0 = 2;
ds4_gpu_dsv4_compressor_store_one_args store0_args = {
.width = width0,
.ratio = ratio,
.pos = pos,
.ape_type = ape0_type,
};
ds4_gpu_dsv4_compressor_store_one_args store1_args = {
.width = width1,
.ratio = ratio,
.pos = pos,
.ape_type = ape1_type,
};
const uint32_t max_out = q_rank > kv_dim ? q_rank : kv_dim;
const uint32_t pair_vtgs = (max_out + 1u) / 2u;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_mul_mv_pipeline(
"kernel_dsv4_qkv_pair_quad_compressor_store_q8_0", 8);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args0 length:sizeof(args0) atIndex:0];
[enc setBytes:&args1 length:sizeof(args1) atIndex:1];
[enc setBytes:&cargs length:sizeof(cargs) atIndex:2];
[enc setBytes:&store0_args length:sizeof(store0_args) atIndex:3];
[enc setBytes:&store1_args length:sizeof(store1_args) atIndex:4];
[enc setBytes:&pair_vtgs length:sizeof(pair_vtgs) atIndex:5];
[enc setBuffer:qw0buf offset:(NSUInteger)q_a_inner atIndex:6];
[enc setBuffer:qw1buf offset:(NSUInteger)kv_inner atIndex:7];
[enc setBuffer:w0kvbuf offset:(NSUInteger)w0kv_inner atIndex:8];
[enc setBuffer:w0scbuf offset:(NSUInteger)w0sc_inner atIndex:9];
[enc setBuffer:w1kvbuf offset:(NSUInteger)w1kv_inner atIndex:10];
[enc setBuffer:w1scbuf offset:(NSUInteger)w1sc_inner atIndex:11];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:12];
[enc setBuffer:ds4_gpu_tensor_buffer(qr) offset:ds4_gpu_tensor_offset(qr) atIndex:13];
[enc setBuffer:ds4_gpu_tensor_buffer(kv_raw) offset:ds4_gpu_tensor_offset(kv_raw) atIndex:14];
[enc setBuffer:ds4_gpu_tensor_buffer(out0_kv) offset:ds4_gpu_tensor_offset(out0_kv) atIndex:15];
[enc setBuffer:ds4_gpu_tensor_buffer(out0_score) offset:ds4_gpu_tensor_offset(out0_score) atIndex:16];
[enc setBuffer:ds4_gpu_tensor_buffer(out1_kv) offset:ds4_gpu_tensor_offset(out1_kv) atIndex:17];
[enc setBuffer:ds4_gpu_tensor_buffer(out1_score) offset:ds4_gpu_tensor_offset(out1_score) atIndex:18];
[enc setBuffer:ape0buf offset:(NSUInteger)ape0_inner atIndex:19];
[enc setBuffer:ape1buf offset:(NSUInteger)ape1_inner atIndex:20];
[enc setBuffer:ds4_gpu_tensor_buffer(state0_kv) offset:ds4_gpu_tensor_offset(state0_kv) atIndex:21];
[enc setBuffer:ds4_gpu_tensor_buffer(state0_score) offset:ds4_gpu_tensor_offset(state0_score) atIndex:22];
[enc setBuffer:ds4_gpu_tensor_buffer(state1_kv) offset:ds4_gpu_tensor_offset(state1_kv) atIndex:23];
[enc setBuffer:ds4_gpu_tensor_buffer(state1_score) offset:ds4_gpu_tensor_offset(state1_score) atIndex:24];
[enc setThreadgroupMemoryLength:2u * 2u * 2u * 32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(
(NSUInteger)(pair_vtgs + 1u) / 2u +
((NSUInteger)width0 + (NSUInteger)width1 + 1u) / 2u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "qkv pair + quad compressor store")) {
return -1;
}
}
return 1;
}
int ds4_gpu_matmul_f32_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t n_tok) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok == 0 || n_tok > UINT32_MAX) return 0;
if (in_dim > UINT64_MAX / n_tok / sizeof(float) ||
out_dim > UINT64_MAX / n_tok / sizeof(float)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t x_bytes = n_tok * in_dim * sizeof(float);
const uint64_t out_bytes = n_tok * out_dim * sizeof(float);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal F32 tensor matmul received undersized activation buffers\n");
return 0;
}
const uint64_t row_bytes = in_dim * sizeof(float);
const uint64_t weight_bytes = row_bytes * out_dim;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal F32 tensor matmul range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (n_tok == 1) {
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, 1);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1);
if (ds4_gpu_plain_mv_single_row(out_dim)) {
mv_dispatch.nr0 = 1;
mv_dispatch.smem = 32u * sizeof(float);
}
mv_args.nr0 = mv_dispatch.nr0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
if (!pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
if (mv_dispatch.smem) {
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
1,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matvec")) return 0;
return 1;
}
const bool router_batch = in_dim == 5120u && out_dim == 384u &&
n_tok >= 32u && !g_quality_mode &&
!getenv("DS4_METAL_DISABLE_V41_ROUTER_BATCH");
if (router_batch && n_tok >= 256u && !getenv("DS4_METAL_DISABLE_V41_ROUTER_MM")) {
const bool bc_out = (n_tok % 32u) != 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mm_pipeline("kernel_mul_mm_f32_f32_full", false, bc_out);
if (!pipeline) return 0;
ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(in_dim, out_dim, n_tok, row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:12288u atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tok + 31u)/32u,
(NSUInteger)out_dim/64u, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "F32 router prefill matmul");
}
if ((n_tok <= 8 || router_batch) && (in_dim % 128u) == 0) {
const int16_t nsg = router_batch ? 2 : ds4_gpu_mv_ext_nsg();
const int16_t nxpsg = ds4_gpu_mv_ext_nxpsg(in_dim, n_tok);
const int16_t r1ptg = ds4_gpu_mv_ext_r1ptg(n_tok);
const char *fn_name = ds4_gpu_mv_ext_f32_name(r1ptg);
id<MTLComputePipelineState> pipeline =
fn_name ? ds4_gpu_get_mul_mv_ext_pipeline(fn_name, nsg, nxpsg) : nil;
if (!pipeline) return 0;
const int16_t nypsg = 32 / nxpsg;
const uint64_t r0ptg = (uint64_t)nypsg * (uint64_t)nsg;
ds4_gpu_mul_mv_ext_args args =
ds4_gpu_make_mv_ext_args(in_dim, out_dim, n_tok, sizeof(float), row_bytes);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)r0ptg - 1u) / (NSUInteger)r0ptg,
((NSUInteger)n_tok + (NSUInteger)r1ptg - 1u) / (NSUInteger)r1ptg,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor mul_mv_ext")) return 0;
return 1;
}
/* Generic multi-row path (GLM prefill shapes: one grid row per
* token through the plain matvec pipeline). */
{
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, n_tok);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1);
mv_args.nr0 = mv_dispatch.nr0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
if (!pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
if (mv_dispatch.smem) {
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
(NSUInteger)n_tok,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "F32 tensor matmul")) return 0;
}
}
return 1;
}
int ds4_gpu_repeat_hc_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *row,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !row || n_embd == 0 || n_hc == 0) return 0;
@autoreleasepool {
id<MTLBuffer> rowbuf = ds4_gpu_tensor_buffer(row);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t out_bytes = row_bytes * n_hc;
if (!rowbuf || !outbuf ||
ds4_gpu_tensor_bytes(row) < row_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal HC repeat received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_repeat_hc_embedding(cb,
rowbuf,
ds4_gpu_tensor_offset(row),
outbuf,
ds4_gpu_tensor_offset(out),
1,
n_embd,
n_hc)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC repeat")) return 0;
}
return 1;
}
int ds4_gpu_repeat_hc_rows_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *rows,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !rows || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0;
@autoreleasepool {
id<MTLBuffer> rowsbuf = ds4_gpu_tensor_buffer(rows);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t rows_bytes = (uint64_t)n_tokens * row_bytes;
const uint64_t out_bytes = rows_bytes * n_hc;
if (!rowsbuf || !outbuf ||
ds4_gpu_tensor_bytes(rows) < rows_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal HC row repeat received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_repeat_hc_embedding(cb,
rowsbuf,
ds4_gpu_tensor_offset(rows),
outbuf,
ds4_gpu_tensor_offset(out),
n_tokens,
n_embd,
n_hc)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC row repeat")) return 0;
}
return 1;
}
int ds4_gpu_rms_norm_plain_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
uint32_t n,
float eps) {
return ds4_gpu_rms_norm_plain_rows_tensor(out, x, n, 1, eps);
}
int ds4_gpu_rms_norm_plain_rows_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
uint32_t n,
uint32_t rows,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (n == 0 || rows == 0 || (n & 3u) != 0) return 0;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t bytes = (uint64_t)n * rows * sizeof(float);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < bytes ||
ds4_gpu_tensor_bytes(out) < bytes) {
fprintf(stderr, "ds4: Metal plain RMS norm received undersized activation buffers\n");
return 0;
}
ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps);
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_rms_norm_plain_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "plain RMS norm")) return 0;
}
return 1;
}
static int ds4_gpu_hc_rms_scale_project_mode(
uint32_t in_dim,
uint32_t out_dim,
uint32_t n_rows) {
const bool hard_shape = n_rows > 8u &&
(in_dim == 16384u || in_dim == 20480u || in_dim == 28672u) && out_dim == 24u;
if (!hard_shape) return 0;
if (in_dim == 20480u && n_rows < 512u) return 0;
if ((in_dim != 20480u && !ds4_gpu_device_name_contains("M3") &&
(g_test_flags & DS4_GPU_TEST_HC_RMS_SCALE_PROJ) == 0u) ||
getenv("DS4_METAL_DISABLE_HC_RMS_SCALE_PROJ") != NULL) {
return 0;
}
if (g_rms_norm_scale_pipeline == nil) {
return 0;
}
return 1;
}
int ds4_gpu_hc_rms_scale_project_f16_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *scale_scratch,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !scale_scratch || !model_map || !x ||
in_dim == 0u || out_dim == 0u || n_rows == 0u ||
n_rows > INT32_MAX) {
return 0;
}
const int mode =
ds4_gpu_hc_rms_scale_project_mode(in_dim, out_dim, n_rows);
if (mode < 0) return 0;
if (mode == 0) {
return ds4_gpu_rms_norm_plain_rows_tensor(
scale_scratch, x, in_dim, n_rows, eps) != 0 &&
ds4_gpu_matmul_f16_tensor(
out, model_map, model_size, weight_offset,
in_dim, out_dim, scale_scratch, n_rows) != 0;
}
@autoreleasepool {
const uint64_t x_elems = (uint64_t)in_dim * n_rows;
const uint64_t out_elems = (uint64_t)out_dim * n_rows;
if (x_elems > UINT64_MAX / sizeof(float) ||
out_elems > UINT64_MAX / sizeof(float)) {
return 0;
}
const uint64_t x_bytes = x_elems * sizeof(float);
const uint64_t out_bytes = out_elems * sizeof(float);
const uint64_t scale_bytes = (uint64_t)n_rows * sizeof(float);
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> scalebuf = ds4_gpu_tensor_buffer(scale_scratch);
if (!xbuf || !outbuf || !scalebuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(scale_scratch) < scale_bytes) {
fprintf(stderr, "ds4: Metal HC RMS scale projection received undersized activation buffers\n");
return 0;
}
const uint64_t weight_row_bytes = (uint64_t)in_dim * sizeof(uint16_t);
if ((uint64_t)out_dim > UINT64_MAX / weight_row_bytes) return 0;
const uint64_t weight_bytes = (uint64_t)out_dim * weight_row_bytes;
if (weight_offset > model_size ||
weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal HC RMS scale projection weight range is outside the mapped model\n");
return 0;
}
const bool bc_inp = (in_dim % 32u) != 0u;
const bool bc_out = (out_dim % 64u) != 0u ||
(n_rows % 32u) != 0u;
id<MTLComputePipelineState> mm_pipeline =
ds4_gpu_get_mul_mm_pipeline(
"kernel_mul_mm_f16_f32_scaled", bc_inp, bc_out);
if (!mm_pipeline) {
return ds4_gpu_rms_norm_plain_rows_tensor(
scale_scratch, x, in_dim, n_rows, eps) != 0 &&
ds4_gpu_matmul_f16_tensor(
out, model_map, model_size, weight_offset,
in_dim, out_dim, scale_scratch, n_rows) != 0;
}
id<MTLComputePipelineState> scale_pipeline = ds4_gpu_hot_pipeline(
g_rms_norm_scale_pipeline, "kernel_rms_norm_scale_f32_4");
if (!scale_pipeline) return 0;
uint64_t weight_inner = 0;
id<MTLBuffer> weightbuf = ds4_gpu_wrap_model_range(
model_map, model_size, weight_offset, weight_bytes, &weight_inner);
if (!weightbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_rms_norm_args norm_args =
ds4_gpu_make_rms_norm_args(in_dim, n_rows, eps);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:scale_pipeline];
[enc setBytes:&norm_args length:sizeof(norm_args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:scalebuf
offset:ds4_gpu_tensor_offset(scale_scratch)
atIndex:2];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(
ds4_gpu_rms_norm_threads(in_dim), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
ds4_gpu_mul_mm_args mm_args = ds4_gpu_make_mm_args(
in_dim, out_dim, n_rows, weight_row_bytes);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:mm_pipeline];
[enc setBytes:&mm_args length:sizeof(mm_args) atIndex:0];
[enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setBuffer:scalebuf
offset:ds4_gpu_tensor_offset(scale_scratch)
atIndex:4];
[enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(
((NSUInteger)n_rows + 31u) / 32u,
((NSUInteger)out_dim + 63u) / 64u,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "HC RMS scale F16 projection")) {
return 0;
}
}
return 1;
}
int ds4_gpu_rms_norm_weight_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
float eps) {
return ds4_gpu_rms_norm_weight_rows_tensor(out, x, model_map, model_size, weight_offset, n, 1, eps);
}
static int ds4_gpu_rms_norm_weight_rows_impl(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
uint32_t rows,
float eps,
int round_bf16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (n == 0 || rows == 0 || (n & 3u) != 0) return 0;
id<MTLComputePipelineState> pipeline = g_rms_norm_pipeline;
if (round_bf16) {
static id<MTLComputePipelineState> rounded;
if (!rounded) rounded = ds4_gpu_get_pipeline("kernel_rms_norm_mul_f32_4_bf16");
pipeline = rounded;
}
if (!pipeline) return 0;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t row_bytes = (uint64_t)n * sizeof(float);
const uint64_t bytes = row_bytes * rows;
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < bytes ||
ds4_gpu_tensor_bytes(out) < bytes) {
fprintf(stderr, "ds4: Metal weighted RMS norm received undersized activation buffers\n");
return 0;
}
if (weight_offset > model_size || row_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal weighted RMS norm range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
row_bytes,
&inner_offset);
if (!wbuf) return 0;
ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, rows, eps);
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "weighted RMS norm")) return 0;
}
return 1;
}
int ds4_gpu_rms_norm_weight_rows_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
uint32_t rows,
float eps) {
return ds4_gpu_rms_norm_weight_rows_impl(out, x, model_map, model_size, weight_offset, n, rows, eps, 0);
}
int ds4_gpu_rms_norm_weight_bf16_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
float eps) {
return ds4_gpu_rms_norm_weight_rows_impl(out, x, model_map, model_size, weight_offset, n, 1, eps, 1);
}
int ds4_gpu_add_rms_norm_weight_tensor(
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *sum_out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!norm_out || !sum_out || !a || !b || n == 0 || (n & 3u) != 0) return 0;
@autoreleasepool {
id<MTLBuffer> abuf = ds4_gpu_tensor_buffer(a);
id<MTLBuffer> bbuf = ds4_gpu_tensor_buffer(b);
id<MTLBuffer> sumbuf = ds4_gpu_tensor_buffer(sum_out);
id<MTLBuffer> normbuf = ds4_gpu_tensor_buffer(norm_out);
const uint64_t row_bytes = (uint64_t)n * sizeof(float);
if (!abuf || !bbuf || !sumbuf || !normbuf ||
ds4_gpu_tensor_bytes(a) < row_bytes ||
ds4_gpu_tensor_bytes(b) < row_bytes ||
ds4_gpu_tensor_bytes(sum_out) < row_bytes ||
ds4_gpu_tensor_bytes(norm_out) < row_bytes) {
fprintf(stderr, "ds4: Metal add+RMS norm received undersized activation buffers\n");
return 0;
}
if (weight_offset > model_size || row_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal add+RMS norm range is outside the mapped model\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
row_bytes,
&inner_offset);
if (!wbuf) return 0;
ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_args(n, 1, eps);
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_add_rms_norm_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1];
[enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:3];
[enc setBuffer:sumbuf offset:ds4_gpu_tensor_offset(sum_out) atIndex:4];
[enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:5];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(n), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "add+RMS norm")) return 0;
}
return 1;
}
int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_out,
const ds4_gpu_tensor *kv,
uint64_t kv_weight_offset,
uint32_t kv_n,
uint32_t rows,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!q_out || !q || !kv_out || !kv || q_n == 0 || kv_n == 0 || rows == 0 ||
(q_n & 3u) != 0 || (kv_n & 3u) != 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> qoutbuf = ds4_gpu_tensor_buffer(q_out);
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
id<MTLBuffer> kvoutbuf = ds4_gpu_tensor_buffer(kv_out);
const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float);
const uint64_t kv_row_bytes = (uint64_t)kv_n * sizeof(float);
if (!qbuf || !qoutbuf || !kvbuf || !kvoutbuf ||
ds4_gpu_tensor_bytes(q) < q_row_bytes * rows ||
ds4_gpu_tensor_bytes(q_out) < q_row_bytes * rows ||
ds4_gpu_tensor_bytes(kv) < kv_row_bytes * rows ||
ds4_gpu_tensor_bytes(kv_out) < kv_row_bytes * rows) {
fprintf(stderr, "ds4: Metal fused q/kv RMS norm received undersized activation buffers\n");
return 0;
}
if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset ||
kv_weight_offset > model_size || kv_row_bytes > model_size - kv_weight_offset) {
fprintf(stderr, "ds4: Metal fused q/kv RMS norm weight range is outside the mapped model\n");
return 0;
}
uint64_t q_inner_offset = 0;
uint64_t kv_inner_offset = 0;
id<MTLBuffer> q_wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
q_weight_offset, q_row_bytes,
&q_inner_offset);
if (!q_wbuf) return 0;
id<MTLBuffer> kv_wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
kv_weight_offset, kv_row_bytes,
&kv_inner_offset);
if (!kv_wbuf) return 0;
ds4_gpu_qkv_rms_norm_args args = {
.q_n = (int32_t)q_n,
.q_n4 = (int32_t)(q_n / 4u),
.kv_n = (int32_t)kv_n,
.kv_n4 = (int32_t)(kv_n / 4u),
.q_row_stride = q_row_bytes,
.kv_row_stride = kv_row_bytes,
.eps = eps,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_qkv_rms_norm_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:q_wbuf offset:(NSUInteger)q_inner_offset atIndex:2];
[enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3];
[enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:4];
[enc setBuffer:kv_wbuf offset:(NSUInteger)kv_inner_offset atIndex:5];
[enc setBuffer:kvoutbuf offset:ds4_gpu_tensor_offset(kv_out) atIndex:6];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(rows, 2, 1)
threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "fused q/kv RMS norm")) return 0;
}
return 1;
}
/* Decode-only triple fusion: q/kv RMS norm + KV RoPE tail + FP8/raw store.
* Same arithmetic, order and rounding as the three separate dispatches (see
* metal/norm.metal); gated and verified against full-vocabulary logits.
* Returns 1 on success, 0 on unsupported shape (caller falls back). */
/* Deferred kv norm/RoPE/FP8 task: the q/kv norm dispatch runs its q task
* only and the kv task is folded into the KV staging kernel of the same
* layer (ds4_gpu_kv_norm_task_flush runs it standalone if no staging
* consumed it). */
static struct {
int pending;
id<MTLBuffer> kv, kv_out, kv_w, q, q_out, q_w, raw;
NSUInteger kv_off, kv_out_off, kv_w_off, q_off, q_out_off, q_w_off, raw_off;
ds4_gpu_qkv_rms_norm_args nargs;
ds4_gpu_rope_affine_pair_args rope;
ds4_gpu_dsv4_kv_fp8_store_args store;
NSUInteger norm_threads;
} g_kv_task;
static int g_qkv_norm_defer_kv;
void ds4_gpu_dsv4_qkv_norm_defer_kv_next(void) {
g_qkv_norm_defer_kv = 1;
}
int ds4_gpu_kv_norm_task_pending(void) {
return g_kv_task.pending;
}
/* Run a deferred kv task standalone (the full q/kv norm kernel; the q task
* recomputes identical values). Used when no KV staging consumed it. */
int ds4_gpu_kv_norm_task_flush(void) {
if (!g_kv_task.pending) return 1;
g_kv_task.pending = 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_qkv_rms_norm_kv_rope_fp8_store_f32");
if (!pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&g_kv_task.nargs length:sizeof(g_kv_task.nargs) atIndex:0];
[enc setBytes:&g_kv_task.rope length:sizeof(g_kv_task.rope) atIndex:1];
[enc setBytes:&g_kv_task.store length:sizeof(g_kv_task.store) atIndex:2];
[enc setBuffer:g_kv_task.q offset:g_kv_task.q_off atIndex:3];
[enc setBuffer:g_kv_task.q_w offset:g_kv_task.q_w_off atIndex:4];
[enc setBuffer:g_kv_task.q_out offset:g_kv_task.q_out_off atIndex:5];
[enc setBuffer:g_kv_task.kv offset:g_kv_task.kv_off atIndex:6];
[enc setBuffer:g_kv_task.kv_w offset:g_kv_task.kv_w_off atIndex:7];
[enc setBuffer:g_kv_task.kv_out offset:g_kv_task.kv_out_off atIndex:8];
[enc setBuffer:g_kv_task.raw offset:g_kv_task.raw_off atIndex:9];
[enc setThreadgroupMemoryLength:(32u + 64u) * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 2, 1)
threadsPerThreadgroup:MTLSizeMake(g_kv_task.norm_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "deferred kv norm task");
}
/* Run the deferred kv task in a concurrent encoder left open for the
* caller's next dispatch (the q_b projection), so its ~9 us of serial work
* hides under that stream. End closes the section. */
int ds4_gpu_kv_norm_task_begin_concurrent(void) {
if (!g_kv_task.pending || !g_batch_cb || g_batch_encoder_concurrent) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_qkv_rms_norm_kv_rope_fp8_store_f32");
if (!pipeline) return 0;
g_kv_task.pending = 0;
ds4_gpu_close_batch_encoder();
g_batch_encoder_concurrent = YES;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(g_batch_cb);
if (!enc || enc.dispatchType != MTLDispatchTypeConcurrent) {
ds4_gpu_close_batch_encoder();
g_batch_encoder_concurrent = NO;
g_kv_task.pending = 1;
return 0;
}
[enc setComputePipelineState:pipeline];
[enc setBytes:&g_kv_task.nargs length:sizeof(g_kv_task.nargs) atIndex:0];
[enc setBytes:&g_kv_task.rope length:sizeof(g_kv_task.rope) atIndex:1];
[enc setBytes:&g_kv_task.store length:sizeof(g_kv_task.store) atIndex:2];
[enc setBuffer:g_kv_task.q offset:g_kv_task.q_off atIndex:3];
[enc setBuffer:g_kv_task.q_w offset:g_kv_task.q_w_off atIndex:4];
[enc setBuffer:g_kv_task.q_out offset:g_kv_task.q_out_off atIndex:5];
[enc setBuffer:g_kv_task.kv offset:g_kv_task.kv_off atIndex:6];
[enc setBuffer:g_kv_task.kv_w offset:g_kv_task.kv_w_off atIndex:7];
[enc setBuffer:g_kv_task.kv_out offset:g_kv_task.kv_out_off atIndex:8];
[enc setBuffer:g_kv_task.raw offset:g_kv_task.raw_off atIndex:9];
[enc setThreadgroupMemoryLength:(32u + 64u) * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 2, 1)
threadsPerThreadgroup:MTLSizeMake(g_kv_task.norm_threads, 1, 1)];
ds4_gpu_end_compute_encoder(g_batch_cb, enc);
return 1;
}
void ds4_gpu_kv_norm_task_end_concurrent(void) {
if (!g_batch_encoder_concurrent) return;
ds4_gpu_close_batch_encoder();
g_batch_encoder_concurrent = NO;
}
int ds4_gpu_dsv4_qkv_rms_norm_kv_rope_fp8_store_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_out,
const ds4_gpu_tensor *kv,
uint64_t kv_weight_offset,
uint32_t kv_n,
ds4_gpu_tensor *raw_cache,
uint64_t raw_cap,
uint32_t raw_row,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!q_out || !q || !kv_out || !kv || !raw_cache ||
q_n == 0 || kv_n == 0 || (q_n & 3u) != 0 || (kv_n & 3u) != 0 ||
n_rot > kv_n || (n_rot & 1u) != 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> qoutbuf = ds4_gpu_tensor_buffer(q_out);
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
id<MTLBuffer> kvoutbuf = ds4_gpu_tensor_buffer(kv_out);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_cache);
const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float);
const uint64_t kv_row_bytes = (uint64_t)kv_n * sizeof(float);
const uint64_t raw_bytes = raw_cap * kv_row_bytes;
if (!qbuf || !qoutbuf || !kvbuf || !kvoutbuf || !rawbuf ||
ds4_gpu_tensor_bytes(q) < q_row_bytes ||
ds4_gpu_tensor_bytes(q_out) < q_row_bytes ||
ds4_gpu_tensor_bytes(kv) < kv_row_bytes ||
ds4_gpu_tensor_bytes(kv_out) < kv_row_bytes ||
ds4_gpu_tensor_bytes(raw_cache) < raw_bytes ||
(uint64_t)raw_row >= raw_cap) {
fprintf(stderr, "ds4: Metal fused q/kv norm RoPE store received undersized buffers\n");
return 0;
}
if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset ||
kv_weight_offset > model_size || kv_row_bytes > model_size - kv_weight_offset) {
fprintf(stderr, "ds4: Metal fused q/kv norm RoPE store weight range is outside the mapped model\n");
return 0;
}
uint64_t q_inner_offset = 0;
uint64_t kv_inner_offset = 0;
id<MTLBuffer> q_wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
q_weight_offset, q_row_bytes,
&q_inner_offset);
if (!q_wbuf) return 0;
id<MTLBuffer> kv_wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
kv_weight_offset, kv_row_bytes,
&kv_inner_offset);
if (!kv_wbuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_qkv_rms_norm_kv_rope_fp8_store_f32");
if (!pipeline) return 0;
ds4_gpu_qkv_rms_norm_args args = {
.q_n = (int32_t)q_n,
.q_n4 = (int32_t)(q_n / 4u),
.kv_n = (int32_t)kv_n,
.kv_n4 = (int32_t)(kv_n / 4u),
.q_row_stride = q_row_bytes,
.kv_row_stride = kv_row_bytes,
.eps = eps,
};
ds4_gpu_rope_affine_pair_args rope = {
.row_bytes = kv_row_bytes,
.token_bytes = kv_row_bytes,
.head_dim = (int32_t)kv_n,
.n_dims = (int32_t)n_rot,
.n_ctx_orig = (int32_t)n_ctx_orig,
.inverse = 0,
.pos0 = pos0,
.pos_step = 1,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
};
ds4_gpu_dsv4_kv_fp8_store_args store = {
.head_dim = (int32_t)kv_n,
.n_rot = (int32_t)n_rot,
.raw_row = (int32_t)raw_row,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBytes:&rope length:sizeof(rope) atIndex:1];
[enc setBytes:&store length:sizeof(store) atIndex:2];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:3];
[enc setBuffer:q_wbuf offset:(NSUInteger)q_inner_offset atIndex:4];
[enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:5];
[enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:6];
[enc setBuffer:kv_wbuf offset:(NSUInteger)kv_inner_offset atIndex:7];
[enc setBuffer:kvoutbuf offset:ds4_gpu_tensor_offset(kv_out) atIndex:8];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:9];
[enc setThreadgroupMemoryLength:(32u + 64u) * sizeof(float) atIndex:0];
const NSUInteger norm_threads = ds4_gpu_rms_norm_threads(q_n);
const bool defer_kv = g_qkv_norm_defer_kv && !g_kv_task.pending;
g_qkv_norm_defer_kv = 0;
if (defer_kv) {
args.tasks = 1u; /* q task now; kv task later, concurrently with q_b */
[enc setBytes:&args length:sizeof(args) atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(1, defer_kv ? 1 : 2, 1)
threadsPerThreadgroup:MTLSizeMake(norm_threads, 1, 1)];
if (defer_kv) {
g_kv_task.pending = 1;
g_kv_task.kv = kvbuf; g_kv_task.kv_off = ds4_gpu_tensor_offset(kv);
g_kv_task.kv_out = kvoutbuf; g_kv_task.kv_out_off = ds4_gpu_tensor_offset(kv_out);
g_kv_task.kv_w = kv_wbuf; g_kv_task.kv_w_off = (NSUInteger)kv_inner_offset;
g_kv_task.q = qbuf; g_kv_task.q_off = ds4_gpu_tensor_offset(q);
g_kv_task.q_out = qoutbuf; g_kv_task.q_out_off = ds4_gpu_tensor_offset(q_out);
g_kv_task.q_w = q_wbuf; g_kv_task.q_w_off = (NSUInteger)q_inner_offset;
g_kv_task.raw = rawbuf; g_kv_task.raw_off = ds4_gpu_tensor_offset(raw_cache);
g_kv_task.nargs = args;
g_kv_task.nargs.tasks = 2u; /* kv task only */
g_kv_task.rope = rope;
g_kv_task.store = store;
g_kv_task.norm_threads = norm_threads;
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "fused q/kv norm RoPE store")) return 0;
}
return 1;
}
int ds4_gpu_head_rms_norm_tensor(
ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || (head_dim & 3u) != 0) return 0;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float);
if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
fprintf(stderr, "ds4: Metal head RMS norm received undersized activation buffer\n");
return 0;
}
ds4_gpu_rms_norm_args args = ds4_gpu_make_rms_norm_3d_args(head_dim, n_head, n_tok, eps);
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_rms_norm_plain_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1)
threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_pipeline_threads(head_dim, g_rms_norm_plain_pipeline), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "head RMS norm")) return 0;
}
return 1;
}
int ds4_gpu_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!x || n_tok == 0 || n_head == 0 || head_dim == 0 || n_rot > head_dim || (n_rot & 1u) != 0) {
return 0;
}
if (n_rot == 0) return 1;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
const uint64_t bytes = (uint64_t)n_tok * n_head * head_dim * sizeof(float);
if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
fprintf(stderr, "ds4: Metal RoPE received undersized activation buffer\n");
return 0;
}
ds4_gpu_rope_tail_batch_args args = ds4_gpu_make_rope_tail_args(
n_tok, n_head, head_dim, n_rot, n_ctx_orig, inverse,
freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_rope_tail_inplace(cb,
xbuf,
ds4_gpu_tensor_offset(x),
&args,
n_tok,
n_head,
head_dim,
pos0,
1)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "RoPE tail")) return 0;
}
return 1;
}
int ds4_gpu_head_rms_norm_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!x || n_tok == 0u || n_head == 0u || head_dim == 0u ||
(head_dim & 3u) != 0u || n_rot > head_dim ||
pos0 > (uint32_t)INT32_MAX - n_tok) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
const uint64_t bytes =
(uint64_t)n_tok * n_head * head_dim * sizeof(float);
if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
fprintf(stderr,
"ds4: Metal fused head norm/RoPE received undersized activation buffer\n");
return 0;
}
id<MTLComputePipelineState> pipeline =
g_use_dsv4_head_rms_norm_rope_tail_pipeline
? g_dsv4_head_rms_norm_rope_tail_pipeline
: ds4_gpu_get_pipeline(
"kernel_dsv4_head_rms_norm_rope_tail_f32");
if (!pipeline) return 0;
struct {
int32_t n_head;
int32_t head_dim;
int32_t head_dim4;
int32_t n_dims;
int32_t n_ctx_orig;
int32_t pos0;
int32_t inverse;
float eps;
float freq_base;
float freq_scale;
float ext_factor;
float attn_factor;
float beta_fast;
float beta_slow;
} args = {
.n_head = (int32_t)n_head,
.head_dim = (int32_t)head_dim,
.head_dim4 = (int32_t)(head_dim / 4u),
.n_dims = (int32_t)n_rot,
.n_ctx_orig = (int32_t)n_ctx_orig,
.pos0 = (int32_t)pos0,
.inverse = inverse ? 1 : 0,
.eps = eps,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_head, n_tok, 1)
threadsPerThreadgroup:MTLSizeMake(
ds4_gpu_rms_norm_pipeline_threads(head_dim, pipeline),
1,
1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "fused head norm/RoPE")) {
return 0;
}
}
return 1;
}
int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *q_half,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t n_head,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float eps) {
(void)out; (void)q_half; (void)model_map; (void)model_size;
(void)weight_offset; (void)in_dim; (void)out_dim; (void)x;
(void)n_tok; (void)n_head; (void)head_dim; (void)n_rot; (void)pos0;
(void)n_ctx_orig; (void)inverse; (void)freq_base; (void)freq_scale;
(void)ext_factor; (void)attn_factor; (void)beta_fast; (void)beta_slow;
(void)eps;
return 0;
}
int ds4_gpu_dsv4_fp8_kv_quantize_tensor(
ds4_gpu_tensor *x,
uint32_t n_tok,
uint32_t head_dim,
uint32_t n_rot) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!x || n_tok == 0 || head_dim == 0 || n_rot > head_dim) return 0;
if (n_rot == head_dim) return 1;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
const uint64_t bytes = (uint64_t)n_tok * head_dim * sizeof(float);
if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
fprintf(stderr, "ds4: Metal DSV4 FP8 KV quantize received undersized activation buffer\n");
return 0;
}
ds4_gpu_dsv4_fp8_kv_quantize_args args = {
.ne00 = head_dim,
.ne01 = n_tok,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = (uint64_t)head_dim * sizeof(float),
.nb02 = (uint64_t)n_tok * head_dim * sizeof(float),
.nb03 = (uint64_t)n_tok * head_dim * sizeof(float),
.nb0 = sizeof(float),
.nb1 = (uint64_t)head_dim * sizeof(float),
.nb2 = (uint64_t)n_tok * head_dim * sizeof(float),
.nb3 = (uint64_t)n_tok * head_dim * sizeof(float),
.n_rot = (int32_t)n_rot,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_fp8_kv_quantize_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_tok, 1, 1)
threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 FP8 KV quantize")) return 0;
}
return 1;
}
int ds4_gpu_dsv4_indexer_qat_tensor(
ds4_gpu_tensor *x,
uint32_t n_rows,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!x || n_rows == 0 || head_dim != 128u) return 0;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
const uint64_t bytes = (uint64_t)n_rows * head_dim * sizeof(float);
if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
fprintf(stderr, "ds4: Metal DSV4 indexer QAT received undersized activation buffer\n");
return 0;
}
ds4_gpu_dsv4_indexer_qat_args args = {
.n_rows = n_rows,
.head_dim = head_dim,
.row_stride = (uint64_t)head_dim * sizeof(float),
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_indexer_qat_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setThreadgroupMemoryLength:256u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "DSV4 indexer Hadamard+FP4")) return 0;
}
return 1;
}
static void ds4_gpu_set_rows_thread_shape(
uint32_t width,
NSUInteger *nth_out,
NSUInteger *nrptg_out) {
const NSUInteger nk0 = width ? (NSUInteger)width : 1u;
const NSUInteger max_threads = g_set_rows_f32_i32_pipeline
? (NSUInteger)g_set_rows_f32_i32_pipeline.maxTotalThreadsPerThreadgroup
: 1024u;
NSUInteger nth = 32u;
while (nth < nk0 && nth < max_threads) {
nth *= 2u;
}
NSUInteger nrptg = 1u;
if (nth > nk0) {
nrptg = (nth + nk0 - 1u) / nk0;
nth = nk0;
if (nrptg * nth > max_threads) {
nrptg--;
}
}
if (nth > nk0) nth = nk0;
if (nth == 0u) nth = 1u;
if (nrptg == 0u) nrptg = 1u;
*nth_out = nth;
*nrptg_out = nrptg;
}
static int ds4_gpu_encode_f16_round_copy_for_raw_store(
id<MTLCommandBuffer> cb,
const ds4_gpu_tensor *src,
uint32_t n) {
id<MTLBuffer> srcbuf = ds4_gpu_tensor_buffer(src);
const uint64_t src_bytes = (uint64_t)n * sizeof(float);
if (!srcbuf || ds4_gpu_tensor_bytes(src) < src_bytes) {
fprintf(stderr, "ds4: Metal raw KV store received undersized source buffer\n");
return 0;
}
if (!ds4_gpu_ensure_scratch_buffer(&g_f16_round_scratch_buffer,
&g_f16_round_scratch_bytes,
(NSUInteger)n * sizeof(uint16_t),
"ds4_f16_round_scratch") ||
!ds4_gpu_ensure_scratch_buffer(&g_raw_store_round_buffer,
&g_raw_store_round_bytes,
(NSUInteger)n * sizeof(float),
"ds4_raw_store_round")) {
return 0;
}
if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
srcbuf,
ds4_gpu_tensor_offset(src),
g_f16_round_scratch_buffer,
0,
n)) {
return 0;
}
return ds4_gpu_encode_cpy_f16_f32_1d(cb,
g_f16_round_scratch_buffer,
0,
g_raw_store_round_buffer,
0,
n);
}
static int ds4_gpu_encode_set_rows_f32_i32(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *dst,
id<MTLBuffer> srcbuf,
NSUInteger src_off,
const int32_t *rows,
uint32_t n_rows,
uint32_t dst_rows,
uint32_t width) {
id<MTLBuffer> dstbuf = ds4_gpu_tensor_buffer(dst);
const uint64_t dst_bytes = (uint64_t)dst_rows * width * sizeof(float);
const uint64_t src_bytes = (uint64_t)n_rows * width * sizeof(float);
if (!dstbuf || !srcbuf || !rows || n_rows == 0 || width == 0 ||
ds4_gpu_tensor_bytes(dst) < dst_bytes ||
src_bytes > NSUIntegerMax - src_off) {
fprintf(stderr, "ds4: Metal DS4 set_rows received invalid buffers\n");
return 0;
}
const uint64_t row_bytes = (uint64_t)width * sizeof(float);
const uint64_t rows_bytes = (uint64_t)n_rows * sizeof(int32_t);
ds4_gpu_set_rows_args args = {
.nk0 = (int32_t)width,
.ne01 = (int32_t)n_rows,
.nb01 = row_bytes,
.nb02 = (uint64_t)n_rows * row_bytes,
.nb03 = (uint64_t)n_rows * row_bytes,
.ne11 = 1,
.ne12 = 1,
.nb10 = sizeof(int32_t),
.nb11 = rows_bytes,
.nb12 = rows_bytes,
.nb1 = row_bytes,
.nb2 = (uint64_t)dst_rows * row_bytes,
.nb3 = (uint64_t)dst_rows * row_bytes,
};
NSUInteger nth;
NSUInteger nrptg;
ds4_gpu_set_rows_thread_shape(width, &nth, &nrptg);
id<MTLBuffer> rowsbuf = nil;
if (rows_bytes > 4096u) {
rowsbuf = ds4_gpu_new_transient_buffer((NSUInteger)rows_bytes, "ds4_set_rows_indices");
if (!rowsbuf) return 0;
memcpy([rowsbuf contents], rows, (NSUInteger)rows_bytes);
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_set_rows_f32_i32_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:srcbuf offset:src_off atIndex:1];
if (rowsbuf) {
[enc setBuffer:rowsbuf offset:0 atIndex:2];
} else {
[enc setBytes:rows length:(NSUInteger)rows_bytes atIndex:2];
}
[enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + nrptg - 1u) / nrptg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, nrptg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_add_f32_1d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> a,
NSUInteger a_off,
id<MTLBuffer> b,
NSUInteger b_off,
id<MTLBuffer> out,
NSUInteger out_off,
uint32_t n) {
if (!cb || !a || !b || !out || n == 0) return 0;
ds4_gpu_add_flat_args args = { .n = n };
NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > (NSUInteger)n) nth = (NSUInteger)n;
if (nth == 0u) nth = 1u;
const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_add2_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:a offset:a_off atIndex:1];
[enc setBuffer:b offset:b_off atIndex:2];
[enc setBuffer:out offset:out_off atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
int ds4_gpu_store_raw_kv_tensor(
ds4_gpu_tensor *raw_cache,
const ds4_gpu_tensor *kv,
uint32_t raw_cap,
uint32_t row,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!raw_cache || !kv || raw_cap == 0 || row >= raw_cap || head_dim == 0 || raw_cap > INT32_MAX) return 0;
@autoreleasepool {
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) {
fprintf(stderr, "ds4: Metal raw KV store received undersized destination buffer\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const int32_t row_i32 = (int32_t)row;
if (!ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, head_dim) ||
!ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache,
g_raw_store_round_buffer,
0,
&row_i32,
1,
raw_cap,
head_dim)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV DS4 set_rows store")) return 0;
}
return 1;
}
/* Release decode fused KV finalizer. Reference paths are selected by the C
* graph driver; this Objective-C entry point always means "use the fused
* Metal kernel." */
/* The fused KV RoPE/FP8 kernel replicates the affine decode RoPE specialisation,
* so it is only valid where that specialisation is the path in use. */
int ds4_gpu_kv_rope_fp8_fuse_available(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (g_dsv4_kv_rope_fp8_store_pipeline == nil) return 0;
if (g_rope_tail_inplace_pair_affine_pipeline == nil) return 0;
if (getenv("DS4_METAL_DISABLE_INPLACE_ROPE_PAIR") != NULL) return 0;
if (getenv("DS4_METAL_DISABLE_AFFINE_ROPE_PAIR") != NULL) return 0;
if (!ds4_gpu_device_name_contains("M3") && !ds4_gpu_device_name_contains("M5")) return 0;
return 1;
}
/* Decode-only fusion: one dispatch does the KV RoPE tail and the FP8/raw
* finalizer that previously cost two. Same arithmetic, same order; gated and
* verified against full-vocabulary logits. */
int ds4_gpu_kv_rope_fp8_store_raw_tensor(
ds4_gpu_tensor *kv,
ds4_gpu_tensor *raw_cache,
uint32_t raw_cap,
uint32_t row,
uint32_t head_dim,
uint32_t n_rot,
uint32_t pos0,
uint32_t n_ctx_orig,
bool inverse,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!kv || !raw_cache || raw_cap == 0 || row >= raw_cap || head_dim == 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 || raw_cap > INT32_MAX ||
g_dsv4_kv_rope_fp8_store_pipeline == nil) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_cache);
const uint64_t kv_bytes = (uint64_t)head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
if (!kvbuf || !rawbuf ||
ds4_gpu_tensor_bytes(kv) < kv_bytes ||
ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) {
fprintf(stderr, "ds4: Metal fused KV RoPE/FP8 store received undersized buffers\n");
return 0;
}
ds4_gpu_dsv4_kv_fp8_store_args args = {
.head_dim = (int32_t)head_dim,
.n_rot = (int32_t)n_rot,
.raw_row = (int32_t)row,
};
const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float);
ds4_gpu_rope_affine_pair_args rope = {
.row_bytes = row_bytes,
.token_bytes = row_bytes,
.head_dim = (int32_t)head_dim,
.n_dims = (int32_t)n_rot,
.n_ctx_orig = (int32_t)n_ctx_orig,
.inverse = inverse ? 1 : 0,
.pos0 = pos0,
.pos_step = 1,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_kv_rope_fp8_store_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBytes:&rope length:sizeof(rope) atIndex:1];
[enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:2];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:3];
[enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "KV RoPE/FP8 store fused")) return 0;
}
return 1;
}
int ds4_gpu_kv_fp8_store_raw_tensor(
ds4_gpu_tensor *kv,
ds4_gpu_tensor *raw_cache,
uint32_t raw_cap,
uint32_t row,
uint32_t head_dim,
uint32_t n_rot) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!kv || !raw_cache || raw_cap == 0 || row >= raw_cap || head_dim == 0 ||
n_rot > head_dim || raw_cap > INT32_MAX) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_cache);
const uint64_t kv_bytes = (uint64_t)head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
if (!kvbuf || !rawbuf ||
ds4_gpu_tensor_bytes(kv) < kv_bytes ||
ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) {
fprintf(stderr, "ds4: Metal fused KV FP8/raw-store received undersized buffers\n");
return 0;
}
ds4_gpu_dsv4_kv_fp8_store_args args = {
.head_dim = (int32_t)head_dim,
.n_rot = (int32_t)n_rot,
.raw_row = (int32_t)row,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_kv_fp8_store_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_cache) atIndex:2];
[enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(64, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "KV FP8/raw-store fused")) return 0;
}
return 1;
}
int ds4_gpu_store_raw_kv_batch_tensor(
ds4_gpu_tensor *raw_cache,
const ds4_gpu_tensor *kv,
uint32_t raw_cap,
uint32_t pos0,
uint32_t n_tokens,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!raw_cache || !kv || raw_cap == 0 || n_tokens == 0 || head_dim == 0 || raw_cap > INT32_MAX) return 0;
@autoreleasepool {
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
if (ds4_gpu_tensor_bytes(raw_cache) < raw_bytes) {
fprintf(stderr, "ds4: Metal raw KV batch store received undersized destination buffer\n");
return 0;
}
int32_t rows_stack[512];
int32_t *rows = rows_stack;
if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) {
rows = malloc((size_t)n_tokens * sizeof(*rows));
if (!rows) {
fprintf(stderr, "ds4: failed to allocate raw KV set_rows index list\n");
return 0;
}
}
for (uint32_t t = 0; t < n_tokens; t++) {
rows[t] = (int32_t)((pos0 + t) % raw_cap);
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) {
if (rows != rows_stack) free(rows);
return 0;
}
const uint64_t n = (uint64_t)n_tokens * head_dim;
const int ok = n <= UINT32_MAX &&
ds4_gpu_encode_f16_round_copy_for_raw_store(cb, kv, (uint32_t)n) &&
ds4_gpu_encode_set_rows_f32_i32(cb, raw_cache,
g_raw_store_round_buffer,
0,
rows,
n_tokens,
raw_cap,
head_dim);
if (rows != rows_stack) free(rows);
if (!ok) return 0;
if (!ds4_gpu_finish_command_buffer(cb, owned, "raw KV batch DS4 set_rows store")) return 0;
}
return 1;
}
typedef struct {
uint32_t n_total;
uint32_t width;
uint32_t ratio;
uint32_t pos0;
} ds4_gpu_cpy_tile_rows_args;
/* dst[t*width + e] = ape[((pos0 + t) % ratio)*width + e] for t < n_tokens,
* converting f16 -> f32 when ape_type == 1. One dispatch for the whole
* batch (the previous per-segment copies were n_tokens/ratio encoders). */
static int ds4_gpu_encode_cpy_ape_tile_rows(
id<MTLCommandBuffer> cb,
id<MTLBuffer> ape,
NSUInteger ape_off,
uint32_t ape_type,
id<MTLBuffer> dst,
uint32_t width,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens) {
if (!cb || !ape || !dst || width == 0 || ratio == 0 || n_tokens == 0) return 0;
const uint64_t n_total64 = (uint64_t)n_tokens * width;
if (n_total64 > UINT32_MAX) return 0;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
ape_type == 1u ? "kernel_cpy_tile_rows_f16_f32" : "kernel_cpy_tile_rows_f32_f32");
if (!pipeline) return 0;
ds4_gpu_cpy_tile_rows_args args = { (uint32_t)n_total64, width, ratio, pos0 };
const NSUInteger nth = 256;
const NSUInteger groups = ((NSUInteger)args.n_total + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ape offset:ape_off atIndex:1];
[enc setBuffer:dst offset:0 atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_compressor_score_with_ape(
id<MTLCommandBuffer> cb,
id<MTLBuffer> score_src,
NSUInteger score_src_offset,
id<MTLBuffer> score_dst,
NSUInteger score_dst_offset,
id<MTLBuffer> apebuf,
NSUInteger ape_offset,
uint32_t ape_type,
uint32_t width,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens) {
if (!cb || !score_src || !score_dst || !apebuf ||
width == 0 || ratio == 0 || n_tokens == 0 ||
(ape_type != 0u && ape_type != 1u)) {
return 0;
}
const uint64_t total_elems64 = (uint64_t)n_tokens * width;
if (total_elems64 > UINT32_MAX) {
fprintf(stderr, "ds4: Metal compressor APE add received too many elements\n");
return 0;
}
const uint32_t total_elems = (uint32_t)total_elems64;
const bool use_fused =
ds4_gpu_device_name_contains("M3") &&
getenv("DS4_METAL_DISABLE_COMPRESSOR_APE_ADD") == NULL;
if (use_fused) {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
ape_type == 1u ? "kernel_dsv4_compressor_score_ape_f16"
: "kernel_dsv4_compressor_score_ape_f32");
if (pipeline) {
ds4_gpu_dsv4_compressor_score_ape_args args = {
.width = width,
.ratio = ratio,
.pos0 = pos0,
.n_tokens = n_tokens,
};
NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth == 0) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:score_src offset:score_src_offset atIndex:1];
[enc setBuffer:apebuf offset:ape_offset atIndex:2];
[enc setBuffer:score_dst offset:score_dst_offset atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(
((NSUInteger)total_elems + nth - 1u) / nth, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
}
const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer,
&g_compressor_store_ape_bytes,
scratch_bytes,
"ds4_compressor_store_ape")) {
return 0;
}
if (!ds4_gpu_encode_cpy_ape_tile_rows(cb, apebuf, ape_offset, ape_type,
g_compressor_store_ape_buffer,
width, ratio, pos0, n_tokens)) {
return 0;
}
return ds4_gpu_encode_add_f32_1d(cb,
score_src,
score_src_offset,
g_compressor_store_ape_buffer,
0,
score_dst,
score_dst_offset,
total_elems);
}
static int ds4_gpu_encode_compressor_set_rows_projected(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
id<MTLBuffer> kvbuf,
NSUInteger kv_offset,
id<MTLBuffer> scorebuf,
NSUInteger score_offset,
id<MTLBuffer> apebuf,
NSUInteger ape_offset,
uint32_t ape_type,
uint32_t width,
uint32_t ratio,
uint32_t pos0,
const int32_t *rows,
uint32_t n_rows,
uint32_t state_rows) {
if (!cb || !state_kv || !state_score || !kvbuf || !scorebuf ||
!apebuf || !rows || width == 0 || n_rows == 0 || state_rows == 0) {
return 0;
}
const NSUInteger score_scratch_bytes = (NSUInteger)n_rows * width * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
&g_compressor_store_score_bytes,
score_scratch_bytes,
"ds4_compressor_store_score")) {
return 0;
}
return ds4_gpu_encode_compressor_score_with_ape(cb,
scorebuf,
score_offset,
g_compressor_store_score_buffer,
0,
apebuf,
ape_offset,
ape_type,
width,
ratio,
pos0,
n_rows) &&
ds4_gpu_encode_set_rows_f32_i32(cb,
state_kv,
kvbuf,
kv_offset,
rows,
n_rows,
state_rows,
width) &&
ds4_gpu_encode_set_rows_f32_i32(cb,
state_score,
g_compressor_store_score_buffer,
0,
rows,
n_rows,
state_rows,
width);
}
static int ds4_gpu_compressor_store_one_tensor(
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t width,
uint32_t ratio,
uint32_t pos) {
if (!kv || !sc || !state_kv || !state_score || !model_map ||
width == 0 || ratio == 0 || (ape_type != 0u && ape_type != 1u)) {
return 0;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_dsv4_compressor_store_one_pipeline,
"kernel_dsv4_compressor_store_one");
if (!pipeline) return 0;
const uint32_t state_rows = ratio == 4u ? 2u * ratio : ratio;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t row_bytes = (uint64_t)width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * row_bytes;
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
ds4_gpu_tensor_bytes(kv) < row_bytes ||
ds4_gpu_tensor_bytes(sc) < row_bytes ||
ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
ds4_gpu_tensor_bytes(state_score) < state_bytes) {
return 0;
}
uint64_t ape_inner = 0;
id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size,
ape_offset, ape_bytes,
&ape_inner);
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc);
id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
if (!apebuf || !kvbuf || !scbuf || !statekvbuf || !statescbuf) return 0;
ds4_gpu_dsv4_compressor_store_one_args args = {
.width = width,
.ratio = ratio,
.pos = pos,
.ape_type = ape_type,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const NSUInteger nth = 256u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv) atIndex:1];
[enc setBuffer:scbuf offset:ds4_gpu_tensor_offset(sc) atIndex:2];
[enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:3];
[enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:4];
[enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:5];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)width + nth - 1u) / nth, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "compressor one-row store");
}
int ds4_gpu_compressor_store_batch_tensor(
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!kv || !sc || !state_kv || !state_score || !model_map ||
head_dim == 0 || ratio == 0 || n_tokens == 0 ||
(ape_type != 0u && ape_type != 1u)) {
return 0;
}
@autoreleasepool {
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t width = coff * head_dim;
const uint32_t state_rows = coff * ratio;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
if (ape_offset > model_size || ape_bytes > model_size - ape_offset) {
fprintf(stderr, "ds4: Metal compressor batch APE range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc);
if (!kvbuf || !scbuf ||
ds4_gpu_tensor_bytes(kv) < kv_bytes ||
ds4_gpu_tensor_bytes(sc) < kv_bytes ||
ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
ds4_gpu_tensor_bytes(state_score) < state_bytes) {
fprintf(stderr, "ds4: Metal compressor batch store received undersized buffers\n");
return 0;
}
uint64_t ape_inner = 0;
id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner);
if (!apebuf) return 0;
const uint64_t total_elems64 = (uint64_t)n_tokens * width;
if (total_elems64 > UINT32_MAX || state_rows > INT32_MAX) {
fprintf(stderr, "ds4: Metal compressor batch store received too many elements\n");
return 0;
}
const uint32_t total_elems = (uint32_t)total_elems64;
const NSUInteger scratch_bytes = (NSUInteger)total_elems * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_ape_buffer,
&g_compressor_store_ape_bytes,
scratch_bytes,
"ds4_compressor_store_ape") ||
!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
&g_compressor_store_score_bytes,
scratch_bytes,
"ds4_compressor_store_score")) {
return 0;
}
int32_t rows_stack[16];
int32_t *rows = rows_stack;
if (n_tokens > (uint32_t)(sizeof(rows_stack) / sizeof(rows_stack[0]))) {
rows = malloc((size_t)n_tokens * sizeof(*rows));
if (!rows) {
fprintf(stderr, "ds4: failed to allocate compressor set_rows index list\n");
return 0;
}
}
for (uint32_t t = 0; t < n_tokens; t++) {
const uint32_t pos_mod = (pos0 + t) % ratio;
rows[t] = (int32_t)(ratio == 4u ? ratio + pos_mod : pos_mod);
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) {
if (rows != rows_stack) free(rows);
return 0;
}
int ok = ds4_gpu_encode_cpy_ape_tile_rows(cb, apebuf, (NSUInteger)ape_inner, ape_type,
g_compressor_store_ape_buffer,
width, ratio, pos0, n_tokens);
if (ok) {
ok = ds4_gpu_encode_add_f32_1d(cb,
scbuf,
ds4_gpu_tensor_offset(sc),
g_compressor_store_ape_buffer,
0,
g_compressor_store_score_buffer,
0,
total_elems);
}
if (ok) {
ok = ds4_gpu_encode_set_rows_f32_i32(cb,
state_kv,
kvbuf,
ds4_gpu_tensor_offset(kv),
rows,
n_tokens,
state_rows,
width);
}
if (ok) {
ok = ds4_gpu_encode_set_rows_f32_i32(cb,
state_score,
g_compressor_store_score_buffer,
0,
rows,
n_tokens,
state_rows,
width);
}
if (rows != rows_stack) free(rows);
if (!ok) return 0;
if (!ds4_gpu_finish_command_buffer(cb, owned, "compressor batch DS4 store")) return 0;
}
return 1;
}
static ds4_gpu_bin_args ds4_gpu_make_bin_contiguous_3d_args(
uint32_t cols,
uint32_t rows,
uint32_t planes) {
const uint64_t row_bytes = (uint64_t)cols * sizeof(float);
const uint64_t plane_bytes = (uint64_t)rows * row_bytes;
return (ds4_gpu_bin_args) {
.ne00 = (int32_t)cols,
.ne01 = (int32_t)rows,
.ne02 = (int32_t)planes,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = row_bytes,
.nb02 = plane_bytes,
.nb03 = (uint64_t)planes * plane_bytes,
.ne10 = (int32_t)cols,
.ne11 = (int32_t)rows,
.ne12 = (int32_t)planes,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = row_bytes,
.nb12 = plane_bytes,
.nb13 = (uint64_t)planes * plane_bytes,
.ne0 = (int32_t)cols,
.ne1 = (int32_t)rows,
.ne2 = (int32_t)planes,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = row_bytes,
.nb2 = plane_bytes,
.nb3 = (uint64_t)planes * plane_bytes,
.offs = 0,
.o1 = { 0 },
};
}
static int ds4_gpu_encode_softmax_f32_contiguous(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t width,
uint32_t rows,
uint32_t planes) {
if (!cb || !src || !dst || width == 0 || rows == 0 || planes == 0) return 0;
const uint64_t row_bytes = (uint64_t)width * sizeof(float);
const uint64_t plane_bytes = (uint64_t)rows * row_bytes;
ds4_gpu_softmax_args args = {
.ne00 = (int32_t)width,
.ne01 = (int32_t)rows,
.ne02 = (int32_t)planes,
.nb01 = row_bytes,
.nb02 = plane_bytes,
.nb03 = (uint64_t)planes * plane_bytes,
.ne11 = (int32_t)width,
.ne12 = (int32_t)rows,
.ne13 = (int32_t)planes,
.nb11 = row_bytes,
.nb12 = plane_bytes,
.nb13 = (uint64_t)planes * plane_bytes,
.nb1 = row_bytes,
.nb2 = plane_bytes,
.nb3 = (uint64_t)planes * plane_bytes,
.scale = 1.0f,
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 1,
};
id<MTLComputePipelineState> pipeline =
(width % 4u) == 0 ? g_soft_max_f32_4_pipeline : g_soft_max_f32_pipeline;
if (!pipeline) return 0;
NSUInteger nth = 32u;
if ((width % 4u) == 0) {
while (nth < (NSUInteger)(width / 4u) &&
nth * (NSUInteger)rows * (NSUInteger)planes < 256u) {
nth *= 2u;
}
} else {
while (nth < (NSUInteger)width &&
nth * (NSUInteger)rows * (NSUInteger)planes < 256u) {
nth *= 2u;
}
}
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth == 0) nth = 1u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:src offset:src_off atIndex:2];
[enc setBuffer:src offset:src_off atIndex:3];
[enc setBuffer:dst offset:dst_off atIndex:4];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(rows, planes, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_ensure_compressor_pool_ggml_scratch(
uint32_t n_rows,
uint32_t head_dim) {
const uint64_t pack_elems = (uint64_t)n_rows * head_dim;
if (n_rows == 0u || head_dim == 0u ||
pack_elems > NSUIntegerMax / sizeof(float)) {
return 0;
}
const NSUInteger pack_bytes =
(NSUInteger)pack_elems * sizeof(float);
return
ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_product_buffer,
&g_compressor_pool_product_bytes,
pack_bytes,
"ds4_compressor_pool_product") &&
ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_cont_buffer,
&g_compressor_pool_score_cont_bytes,
pack_bytes,
"ds4_compressor_pool_score_cont") &&
ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_softmax_buffer,
&g_compressor_pool_softmax_bytes,
pack_bytes,
"ds4_compressor_pool_softmax");
}
static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml_reduce(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *out,
uint32_t n_rows,
uint32_t head_dim,
bool decode_ratio4_one_comp) {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
if (!cb || !outbuf || n_rows == 0u || head_dim == 0u ||
ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * sizeof(float) ||
!ds4_gpu_ensure_compressor_pool_ggml_scratch(n_rows, head_dim)) {
return 0;
}
const bool use_pre_m5_exact_reduction_fusion_default =
!g_ssd_streaming_mode &&
(ds4_gpu_device_is_pre_m5_apple_silicon() ||
ds4_gpu_device_is_m5_apple_silicon()) &&
getenv("DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_EXACT_REDUCTION_FUSION") == NULL;
const bool request_exact_reduction_fusion =
decode_ratio4_one_comp &&
n_rows == 8u &&
(head_dim == 128u || head_dim == 512u) &&
use_pre_m5_exact_reduction_fusion_default;
id<MTLComputePipelineState> exact_reduction_fusion_pipeline = nil;
if (request_exact_reduction_fusion) {
exact_reduction_fusion_pipeline = ds4_gpu_hot_pipeline(
g_dsv4_compressor_exact_softmax_product_pipeline,
"kernel_dsv4_compressor_exact_softmax_product_ratio4");
}
const bool use_exact_reduction_fusion =
exact_reduction_fusion_pipeline != nil &&
exact_reduction_fusion_pipeline.threadExecutionWidth == 32u &&
exact_reduction_fusion_pipeline.maxTotalThreadsPerThreadgroup >= 32u;
if (use_exact_reduction_fusion) {
if (getenv(
"DS4_METAL_TEST_POISON_COMPRESSOR_EXACT_REDUCTION_SCRATCH") != NULL) {
const uint32_t poison_bits = 0x7fc01234u;
float poison;
memcpy(&poison, &poison_bits, sizeof(poison));
if (!ds4_gpu_encode_fill_f32_rows(cb,
g_compressor_pool_softmax_buffer,
0,
n_rows,
head_dim,
poison)) {
return 0;
}
}
ds4_gpu_dsv4_compressor_pack_ratio4_args args = {
.head_dim = head_dim,
.n_comp = 1u,
.replay = 1u,
.n_threads = 32u,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:exact_reduction_fusion_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:g_compressor_pool_product_buffer offset:0 atIndex:1];
[enc setBuffer:g_compressor_pool_score_cont_buffer offset:0 atIndex:2];
[enc setBuffer:g_compressor_pool_softmax_buffer offset:0 atIndex:3];
[enc setBuffer:g_compressor_pool_product_buffer offset:0 atIndex:4];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(head_dim, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_encode_sum_rows_f32(cb,
g_compressor_pool_product_buffer,
0,
outbuf,
ds4_gpu_tensor_offset(out),
n_rows,
head_dim);
}
/*
* Preserve the exact graph reduction sequence and its contiguous
* [head_dim, n_rows] scratch layout. Only the preceding copy/transpose
* may be fused by callers.
*/
ds4_gpu_bin_args mul_args =
ds4_gpu_make_bin_contiguous_3d_args(n_rows, head_dim, 1);
return
ds4_gpu_encode_softmax_f32_contiguous(cb,
g_compressor_pool_score_cont_buffer,
0,
g_compressor_pool_softmax_buffer,
0,
n_rows,
head_dim,
1) &&
ds4_gpu_encode_bin_f32_rows(cb,
g_mul_pipeline,
&mul_args,
g_compressor_pool_product_buffer,
0,
g_compressor_pool_softmax_buffer,
0,
g_compressor_pool_product_buffer,
0) &&
ds4_gpu_encode_sum_rows_f32(cb,
g_compressor_pool_product_buffer,
0,
outbuf,
ds4_gpu_tensor_offset(out),
n_rows,
head_dim);
}
static int ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *out,
id<MTLBuffer> kvbuf,
NSUInteger kv_offset,
uint64_t kv_nb0,
uint64_t kv_nb1,
uint64_t kv_nb2,
id<MTLBuffer> scorebuf,
NSUInteger score_offset,
uint64_t score_nb0,
uint64_t score_nb1,
uint64_t score_nb2,
uint32_t n_rows,
uint32_t head_dim) {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
if (!cb || !outbuf || !kvbuf || !scorebuf || n_rows == 0 || head_dim == 0 ||
ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * sizeof(float)) {
return 0;
}
if (!ds4_gpu_ensure_compressor_pool_ggml_scratch(n_rows, head_dim)) {
return 0;
}
const uint64_t cont_row_stride = (uint64_t)n_rows * sizeof(float);
const uint64_t cont_plane_stride = (uint64_t)head_dim * cont_row_stride;
return
ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb,
kvbuf,
kv_offset,
g_compressor_pool_product_buffer,
0,
n_rows,
head_dim,
1,
kv_nb0,
kv_nb1,
kv_nb2,
cont_row_stride,
cont_plane_stride) &&
ds4_gpu_encode_cpy_f32_f32_3d_src_strided(cb,
scorebuf,
score_offset,
g_compressor_pool_score_cont_buffer,
0,
n_rows,
head_dim,
1,
score_nb0,
score_nb1,
score_nb2,
cont_row_stride,
cont_plane_stride) &&
ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml_reduce(cb,
out,
n_rows,
head_dim,
false);
}
static int ds4_gpu_encode_dsv4_softmax_pool(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *out,
id<MTLBuffer> kvbuf,
NSUInteger kv_offset,
uint64_t kv_nb0,
uint64_t kv_nb1,
uint64_t kv_nb2,
id<MTLBuffer> scorebuf,
NSUInteger score_offset,
uint64_t score_nb0,
uint64_t score_nb1,
uint64_t score_nb2,
uint32_t n_rows,
uint32_t head_dim,
uint32_t n_comp) {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
if (!cb || !outbuf || !kvbuf || !scorebuf ||
n_rows == 0 || head_dim == 0 || n_comp == 0 ||
ds4_gpu_tensor_bytes(out) < (uint64_t)head_dim * n_comp * sizeof(float)) {
return 0;
}
if (n_comp == 1) {
return ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml(cb,
out,
kvbuf,
kv_offset,
kv_nb0,
kv_nb1,
kv_nb2,
scorebuf,
score_offset,
score_nb0,
score_nb1,
score_nb2,
n_rows,
head_dim);
}
ds4_gpu_dsv4_softmax_pool_args args = {
.ne00 = (int64_t)n_rows,
.ne01 = (int64_t)head_dim,
.ne02 = (int64_t)n_comp,
.nb00 = kv_nb0,
.nb01 = kv_nb1,
.nb02 = kv_nb2,
.nb10 = score_nb0,
.nb11 = score_nb1,
.nb12 = score_nb2,
.ne0 = (int64_t)head_dim,
.ne1 = (int64_t)n_comp,
.nb0 = sizeof(float),
.nb1 = (uint64_t)head_dim * sizeof(float),
};
const uint64_t n = (uint64_t)head_dim * n_comp;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_softmax_pool_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:kvbuf offset:kv_offset atIndex:1];
[enc setBuffer:scorebuf offset:score_offset atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_concat_f32_dim1(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src0,
NSUInteger src0_offset,
uint32_t src0_rows,
uint64_t src0_row_stride,
id<MTLBuffer> src1,
NSUInteger src1_offset,
uint32_t src1_rows,
uint64_t src1_row_stride,
id<MTLBuffer> dst,
NSUInteger dst_offset,
uint32_t cols,
uint64_t dst_row_stride) {
if (!cb || !src0 || !src1 || !dst || cols == 0 || src0_rows == 0 || src1_rows == 0) {
return 0;
}
const uint32_t rows = src0_rows + src1_rows;
const uint64_t src0_plane = (uint64_t)src0_rows * src0_row_stride;
const uint64_t src1_plane = (uint64_t)src1_rows * src1_row_stride;
const uint64_t dst_plane = (uint64_t)rows * dst_row_stride;
ds4_gpu_concat_args args = {
.ne00 = (int32_t)cols,
.ne01 = (int32_t)src0_rows,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = src0_row_stride,
.nb02 = src0_plane,
.nb03 = src0_plane,
.ne10 = (int32_t)cols,
.ne11 = (int32_t)src1_rows,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = src1_row_stride,
.nb12 = src1_plane,
.nb13 = src1_plane,
.ne0 = (int32_t)cols,
.ne1 = (int32_t)rows,
.ne2 = 1,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = dst_row_stride,
.nb2 = dst_plane,
.nb3 = dst_plane,
.dim = 1,
};
NSUInteger nth = cols < 1024u ? (NSUInteger)cols : 1024u;
const NSUInteger max_threads = g_concat_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth == 0) nth = 1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_concat_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src0 offset:src0_offset atIndex:1];
[enc setBuffer:src1 offset:src1_offset atIndex:2];
[enc setBuffer:dst offset:dst_offset atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_compressor_pack_ratio4_fusion_mode(uint32_t head_dim) {
const bool default_shape = head_dim == 128u || head_dim == 512u;
if (getenv("DS4_METAL_DISABLE_COMPRESSOR_RATIO4_PACK_FUSION") != NULL ||
!(ds4_gpu_device_name_contains("M3") && default_shape)) {
return 0;
}
if (g_dsv4_compressor_pack_ratio4_pipeline == nil) {
return 0;
}
return 1;
}
static int ds4_gpu_compressor_ratio4_decode_pack_mode(uint32_t head_dim) {
const bool default_shape = head_dim == 128u || head_dim == 512u;
const bool default_pre_m5 =
(ds4_gpu_device_is_pre_m5_apple_silicon() ||
ds4_gpu_device_is_m5_apple_silicon()) && default_shape;
if (getenv("DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_RATIO4_DECODE_PACK_FUSION") != NULL ||
!default_pre_m5) {
return 0;
}
if (g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline == nil) {
return 0;
}
return 1;
}
static bool ds4_gpu_buffer_ranges_overlap(
id<MTLBuffer> a,
NSUInteger a_offset,
uint64_t a_bytes,
id<MTLBuffer> b,
NSUInteger b_offset,
uint64_t b_bytes) {
if (a != b || a_bytes == 0u || b_bytes == 0u) return false;
if (a_offset <= b_offset) {
return a_bytes > (uint64_t)(b_offset - a_offset);
}
return b_bytes > (uint64_t)(a_offset - b_offset);
}
static int ds4_gpu_compressor_ratio4_direct_pool_mode(
uint32_t head_dim,
uint32_t n_comp) {
// One compressed row intentionally uses the legacy GGML reduction graph.
if (n_comp <= 1u) return 0;
const bool default_shape = head_dim == 128u || head_dim == 512u;
if (getenv("DS4_METAL_DISABLE_COMPRESSOR_RATIO4_DIRECT_POOL") != NULL ||
!(ds4_gpu_device_name_contains("M3") && default_shape)) {
return 0;
}
if (g_dsv4_softmax_pool_ratio4_direct_pipeline == nil) {
return 0;
}
return 1;
}
static int ds4_gpu_encode_compressor_ratio4_direct_pool(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *out,
id<MTLBuffer> kvbuf,
NSUInteger kv_offset,
id<MTLBuffer> scorebuf,
NSUInteger score_offset,
id<MTLBuffer> statekvbuf,
NSUInteger state_kv_offset,
id<MTLBuffer> statescbuf,
NSUInteger state_score_offset,
uint32_t head_dim,
uint32_t n_comp,
bool replay) {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t n = (uint64_t)head_dim * n_comp;
if (!cb || !outbuf || !kvbuf || !scorebuf || !statekvbuf || !statescbuf ||
head_dim == 0u || n_comp <= 1u || n > UINT32_MAX ||
ds4_gpu_tensor_bytes(out) < n * sizeof(float)) {
return 0;
}
id<MTLComputePipelineState> pipeline = ds4_gpu_hot_pipeline(
g_dsv4_softmax_pool_ratio4_direct_pipeline,
"kernel_dsv4_softmax_pool_ratio4_direct");
if (!pipeline) return 0;
ds4_gpu_dsv4_softmax_pool_ratio4_direct_args args = {
.n_rows = 8,
.head_dim = head_dim,
.n_comp = n_comp,
.replay = replay ? 1u : 0u,
.pad = 0u,
};
NSUInteger nth = 256u;
if (nth > pipeline.maxTotalThreadsPerThreadgroup) {
nth = pipeline.maxTotalThreadsPerThreadgroup;
}
if (nth == 0u) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:kvbuf offset:kv_offset atIndex:1];
[enc setBuffer:scorebuf offset:score_offset atIndex:2];
[enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3];
[enc setBuffer:statescbuf offset:state_score_offset atIndex:4];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:5];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + nth - 1u) / nth, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_compressor_pack_ratio4(
id<MTLCommandBuffer> cb,
id<MTLBuffer> kvbuf,
NSUInteger kv_offset,
id<MTLBuffer> scorebuf,
NSUInteger score_offset,
id<MTLBuffer> statekvbuf,
NSUInteger state_kv_offset,
id<MTLBuffer> statescbuf,
NSUInteger state_score_offset,
uint32_t head_dim,
uint32_t n_comp,
bool replay) {
if (!cb || !kvbuf || !scorebuf || !statekvbuf || !statescbuf ||
!g_compressor_pool_kv_buffer || !g_compressor_pool_score_buffer ||
head_dim == 0 || n_comp == 0 || head_dim > UINT32_MAX / 2u) {
return 0;
}
const uint64_t total_elems64 = (uint64_t)n_comp * 8u * head_dim;
if (total_elems64 > UINT32_MAX || n_comp > UINT32_MAX / 8u) return 0;
id<MTLComputePipelineState> pipeline = ds4_gpu_hot_pipeline(
g_dsv4_compressor_pack_ratio4_pipeline,
"kernel_dsv4_compressor_pack_ratio4");
if (!pipeline) return 0;
NSUInteger nth = head_dim;
if (nth > 256u) nth = 256u;
if (nth > pipeline.maxTotalThreadsPerThreadgroup) {
nth = pipeline.maxTotalThreadsPerThreadgroup;
}
if (nth == 0) return 0;
ds4_gpu_dsv4_compressor_pack_ratio4_args args = {
.head_dim = head_dim,
.n_comp = n_comp,
.replay = replay ? 1u : 0u,
.n_threads = (uint32_t)nth,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:kvbuf offset:kv_offset atIndex:1];
[enc setBuffer:scorebuf offset:score_offset atIndex:2];
[enc setBuffer:statekvbuf offset:state_kv_offset atIndex:3];
[enc setBuffer:statescbuf offset:state_score_offset atIndex:4];
[enc setBuffer:g_compressor_pool_kv_buffer offset:0 atIndex:5];
[enc setBuffer:g_compressor_pool_score_buffer offset:0 atIndex:6];
[enc dispatchThreadgroups:MTLSizeMake(n_comp, 8u, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_compressor_ratio4_decode_pack_ggml(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *out,
id<MTLBuffer> statekvbuf,
NSUInteger state_kv_offset,
id<MTLBuffer> statescbuf,
NSUInteger state_score_offset,
uint32_t head_dim) {
if (!cb || !out || !statekvbuf || !statescbuf || head_dim == 0u ||
!ds4_gpu_ensure_compressor_pool_ggml_scratch(8u, head_dim)) {
return 0;
}
const bool exact_pool_enabled =
ds4_gpu_ported_m5_decode_feature_enabled(
"DS4_METAL_DISABLE_PRE_M5_COMPRESSOR_EXACT_POOL_RATIO4",
"DS4_METAL_DISABLE_M5_COMPRESSOR_EXACT_POOL_RATIO4");
if ((head_dim == 128u || head_dim == 512u) && exact_pool_enabled) {
id<MTLComputePipelineState> exact_pool_pipeline = ds4_gpu_hot_pipeline(
g_dsv4_compressor_exact_pool_ratio4_pipeline,
"kernel_dsv4_compressor_exact_pool_ratio4_decode_ggml");
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const NSUInteger out_offset = ds4_gpu_tensor_offset(out);
const uint64_t state_bytes = 16ull * head_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)head_dim * sizeof(float);
const bool overlap =
ds4_gpu_buffer_ranges_overlap(outbuf, out_offset, out_bytes,
statekvbuf, state_kv_offset,
state_bytes) ||
ds4_gpu_buffer_ranges_overlap(outbuf, out_offset, out_bytes,
statescbuf, state_score_offset,
state_bytes);
if (exact_pool_pipeline && outbuf && !overlap &&
exact_pool_pipeline.threadExecutionWidth == 32u &&
exact_pool_pipeline.maxTotalThreadsPerThreadgroup >= 32u) {
if (getenv(
"DS4_METAL_TEST_POISON_COMPRESSOR_EXACT_REDUCTION_SCRATCH") != NULL) {
const uint32_t poison_bits = 0x7fc01234u;
float poison;
memcpy(&poison, &poison_bits, sizeof(poison));
if (!ds4_gpu_encode_fill_f32_rows(cb,
g_compressor_pool_softmax_buffer,
0,
8u,
head_dim,
poison)) {
return 0;
}
}
ds4_gpu_dsv4_compressor_pack_ratio4_args args = {
.head_dim = head_dim,
.n_comp = 1u,
.replay = 1u,
.n_threads = 32u,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:exact_pool_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:statekvbuf offset:state_kv_offset atIndex:1];
[enc setBuffer:statescbuf offset:state_score_offset atIndex:2];
[enc setBuffer:g_compressor_pool_softmax_buffer offset:0 atIndex:3];
[enc setBuffer:g_compressor_pool_product_buffer offset:0 atIndex:4];
[enc setBuffer:outbuf offset:out_offset atIndex:5];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(head_dim, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
}
id<MTLComputePipelineState> pipeline = ds4_gpu_hot_pipeline(
g_dsv4_compressor_pack_ratio4_decode_ggml_pipeline,
"kernel_dsv4_compressor_pack_ratio4_decode_ggml");
if (!pipeline) return 0;
NSUInteger nth = head_dim;
if (nth > 256u) nth = 256u;
if (nth > pipeline.maxTotalThreadsPerThreadgroup) {
nth = pipeline.maxTotalThreadsPerThreadgroup;
}
if (nth == 0u) return 0;
ds4_gpu_dsv4_compressor_pack_ratio4_args args = {
.head_dim = head_dim,
.n_comp = 1u,
.replay = 1u,
.n_threads = (uint32_t)nth,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:statekvbuf offset:state_kv_offset atIndex:1];
[enc setBuffer:statescbuf offset:state_score_offset atIndex:2];
[enc setBuffer:g_compressor_pool_product_buffer offset:0 atIndex:3];
[enc setBuffer:g_compressor_pool_score_cont_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(8u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_encode_dsv4_softmax_pool_one_comp_ggml_reduce(cb,
out,
8u,
head_dim,
true);
}
static int ds4_gpu_encode_compressor_pool(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *out,
const ds4_gpu_tensor *state_kv,
const ds4_gpu_tensor *state_score,
uint32_t head_dim,
uint32_t ratio,
bool decode_one_token) {
id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
if (!cb || !out || !statekvbuf || !statescbuf || head_dim == 0 || ratio == 0) return 0;
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t width = coff * head_dim;
const uint32_t rows = coff * ratio;
const uint64_t state_bytes = (uint64_t)width * rows * sizeof(float);
if (ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
ds4_gpu_tensor_bytes(state_score) < state_bytes) {
return 0;
}
if (ratio != 4u) {
const uint64_t row_stride = (uint64_t)width * sizeof(float);
return ds4_gpu_encode_dsv4_softmax_pool(cb,
out,
statekvbuf,
ds4_gpu_tensor_offset(state_kv),
row_stride,
sizeof(float),
(uint64_t)rows * row_stride,
statescbuf,
ds4_gpu_tensor_offset(state_score),
row_stride,
sizeof(float),
(uint64_t)rows * row_stride,
ratio,
head_dim,
1);
}
const int decode_pack_mode = decode_one_token
? ds4_gpu_compressor_ratio4_decode_pack_mode(head_dim)
: 0;
if (decode_pack_mode < 0) return 0;
if (decode_pack_mode > 0) {
return ds4_gpu_encode_compressor_ratio4_decode_pack_ggml(
cb,
out,
statekvbuf,
ds4_gpu_tensor_offset(state_kv),
statescbuf,
ds4_gpu_tensor_offset(state_score),
head_dim);
}
const NSUInteger packed_bytes = (NSUInteger)8u * head_dim * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer,
&g_compressor_pool_kv_bytes,
packed_bytes,
"ds4_compressor_pool_kv") ||
!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer,
&g_compressor_pool_score_bytes,
packed_bytes,
"ds4_compressor_pool_score")) {
return 0;
}
const uint64_t state_row_stride = (uint64_t)width * sizeof(float);
const uint64_t pool_row_stride = (uint64_t)head_dim * sizeof(float);
const NSUInteger curr_offset = (NSUInteger)4u * state_row_stride +
(NSUInteger)head_dim * sizeof(float);
if (!ds4_gpu_encode_concat_f32_dim1(cb,
statekvbuf,
ds4_gpu_tensor_offset(state_kv),
4,
state_row_stride,
statekvbuf,
ds4_gpu_tensor_offset(state_kv) + curr_offset,
4,
state_row_stride,
g_compressor_pool_kv_buffer,
0,
head_dim,
pool_row_stride) ||
!ds4_gpu_encode_concat_f32_dim1(cb,
statescbuf,
ds4_gpu_tensor_offset(state_score),
4,
state_row_stride,
statescbuf,
ds4_gpu_tensor_offset(state_score) + curr_offset,
4,
state_row_stride,
g_compressor_pool_score_buffer,
0,
head_dim,
pool_row_stride)) {
return 0;
}
return ds4_gpu_encode_dsv4_softmax_pool(cb,
out,
g_compressor_pool_kv_buffer,
0,
pool_row_stride,
sizeof(float),
packed_bytes,
g_compressor_pool_score_buffer,
0,
pool_row_stride,
sizeof(float),
packed_bytes,
8,
head_dim,
1);
}
static int ds4_gpu_encode_compressor_shift_ratio4(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
uint32_t width) {
id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
if (!cb || !statekvbuf || !statescbuf || !g_dsv4_ratio4_shift_pipeline || width == 0) return 0;
ds4_gpu_dsv4_ratio4_shift_args args = { .width = width };
const uint32_t n = 4u * width;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_ratio4_shift_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:statekvbuf offset:ds4_gpu_tensor_offset(state_kv) atIndex:1];
[enc setBuffer:statescbuf offset:ds4_gpu_tensor_offset(state_score) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
int ds4_gpu_compressor_prefill_tensor(
ds4_gpu_tensor *comp_cache,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos0,
uint32_t n_tokens,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool quantize_fp8,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
head_dim == 0 || ratio == 0 || n_tokens == 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
(ape_type != 0u && ape_type != 1u) ||
norm_type != 0u) {
return 0;
}
@autoreleasepool {
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t width = coff * head_dim;
const uint32_t state_rows = coff * ratio;
const uint32_t n_comp = n_tokens / ratio;
const uint32_t cutoff = n_comp * ratio;
const uint32_t rem = n_tokens - cutoff;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
norm_offset > model_size || norm_bytes > model_size - norm_offset) {
fprintf(stderr, "ds4: Metal compressor prefill tensor range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc);
id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_cache);
id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf ||
ds4_gpu_tensor_bytes(kv) < kv_bytes ||
ds4_gpu_tensor_bytes(sc) < kv_bytes ||
ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
ds4_gpu_tensor_bytes(state_score) < state_bytes ||
(n_comp && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) {
fprintf(stderr, "ds4: Metal compressor prefill received undersized buffers\n");
return 0;
}
uint64_t ape_inner = 0;
id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner);
if (!apebuf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
int ok = 1;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) ok = 0;
if (ok) {
ok = ds4_gpu_encode_fill_f32_rows(cb,
statekvbuf,
ds4_gpu_tensor_offset(state_kv),
width,
state_rows,
0.0f) &&
ds4_gpu_encode_fill_f32_rows(cb,
statescbuf,
ds4_gpu_tensor_offset(state_score),
width,
state_rows,
ds4_gpu_negative_infinity());
}
if (ok && ratio == 4u) {
int32_t rows_prev[4] = { 0, 1, 2, 3 };
const int have_prev = cutoff >= ratio ? 1 : 0;
const uint32_t prev_start = rem == 0 ? cutoff - ratio : cutoff - ratio;
if (have_prev) {
ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
state_kv,
state_score,
kvbuf,
ds4_gpu_tensor_offset(kv) +
(NSUInteger)prev_start * width * sizeof(float),
scbuf,
ds4_gpu_tensor_offset(sc) +
(NSUInteger)prev_start * width * sizeof(float),
apebuf,
(NSUInteger)ape_inner,
ape_type,
width,
ratio,
pos0 + prev_start,
rows_prev,
4,
state_rows);
}
if (ok && rem != 0) {
int32_t rows_cur[4];
for (uint32_t i = 0; i < rem; i++) rows_cur[i] = (int32_t)(ratio + i);
ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
state_kv,
state_score,
kvbuf,
ds4_gpu_tensor_offset(kv) +
(NSUInteger)cutoff * width * sizeof(float),
scbuf,
ds4_gpu_tensor_offset(sc) +
(NSUInteger)cutoff * width * sizeof(float),
apebuf,
(NSUInteger)ape_inner,
ape_type,
width,
ratio,
pos0 + cutoff,
rows_cur,
rem,
state_rows);
}
} else if (ok && rem != 0) {
int32_t rows[128];
if (rem > (uint32_t)(sizeof(rows) / sizeof(rows[0]))) {
fprintf(stderr, "ds4: Metal compressor prefill remainder exceeds local row list\n");
ok = 0;
} else {
for (uint32_t i = 0; i < rem; i++) rows[i] = (int32_t)i;
ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
state_kv,
state_score,
kvbuf,
ds4_gpu_tensor_offset(kv) +
(NSUInteger)cutoff * width * sizeof(float),
scbuf,
ds4_gpu_tensor_offset(sc) +
(NSUInteger)cutoff * width * sizeof(float),
apebuf,
(NSUInteger)ape_inner,
ape_type,
width,
ratio,
pos0 + cutoff,
rows,
rem,
state_rows);
}
}
if (ok && n_comp != 0) {
const NSUInteger score_bytes = (NSUInteger)cutoff * width * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
&g_compressor_store_score_bytes,
score_bytes,
"ds4_compressor_store_score")) {
ok = 0;
}
if (ok) {
ok = ds4_gpu_encode_compressor_score_with_ape(cb,
scbuf,
ds4_gpu_tensor_offset(sc),
g_compressor_store_score_buffer,
0,
apebuf,
(NSUInteger)ape_inner,
ape_type,
width,
ratio,
pos0,
cutoff);
}
if (ok && ratio == 4u) {
const int direct_pool_mode =
ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp);
if (ok && direct_pool_mode < 0) ok = 0;
const uint64_t direct_output_bytes =
(uint64_t)n_comp * head_dim * sizeof(float);
const uint64_t direct_input_bytes =
(uint64_t)n_comp * 4u * width * sizeof(float);
const bool direct_pool_overlap = direct_pool_mode > 0 &&
ds4_gpu_buffer_ranges_overlap(
compbuf,
ds4_gpu_tensor_offset(comp_cache),
direct_output_bytes,
kvbuf,
ds4_gpu_tensor_offset(kv),
direct_input_bytes);
const bool use_direct_pool =
direct_pool_mode > 0 && !direct_pool_overlap;
const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float);
if (ok && !use_direct_pool &&
(!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer,
&g_compressor_pool_kv_bytes,
pack_bytes,
"ds4_compressor_pool_kv") ||
!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer,
&g_compressor_pool_score_bytes,
pack_bytes,
"ds4_compressor_pool_score"))) {
ok = 0;
}
const int pack_fusion_mode = use_direct_pool ? 0 :
ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim);
if (ok && pack_fusion_mode < 0) ok = 0;
const bool use_pack_fusion = pack_fusion_mode > 0;
if (ok && use_direct_pool) {
ok = ds4_gpu_encode_compressor_ratio4_direct_pool(
cb,
comp_cache,
kvbuf,
ds4_gpu_tensor_offset(kv),
g_compressor_store_score_buffer,
0,
statekvbuf,
ds4_gpu_tensor_offset(state_kv),
statescbuf,
ds4_gpu_tensor_offset(state_score),
head_dim,
n_comp,
false);
}
if (ok && use_pack_fusion) {
ok = ds4_gpu_encode_compressor_pack_ratio4(
cb,
kvbuf,
ds4_gpu_tensor_offset(kv),
g_compressor_store_score_buffer,
0,
statekvbuf,
ds4_gpu_tensor_offset(state_kv),
statescbuf,
ds4_gpu_tensor_offset(state_score),
head_dim,
n_comp,
false);
}
if (ok && !use_direct_pool && !use_pack_fusion) {
ok = ds4_gpu_encode_fill_f32_rows(cb,
g_compressor_pool_kv_buffer,
0,
head_dim,
8u * n_comp,
0.0f) &&
ds4_gpu_encode_fill_f32_rows(cb,
g_compressor_pool_score_buffer,
0,
head_dim,
8u * n_comp,
ds4_gpu_negative_infinity());
}
if (ok && !use_direct_pool && !use_pack_fusion) {
const uint64_t src_row_stride = (uint64_t)width * sizeof(float);
const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride;
const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float);
const uint64_t dst_plane_stride = 8ull * dst_row_stride;
ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
kvbuf,
ds4_gpu_tensor_offset(kv) +
(NSUInteger)head_dim * sizeof(float),
g_compressor_pool_kv_buffer,
(NSUInteger)4u * head_dim * sizeof(float),
head_dim,
ratio,
n_comp,
src_row_stride,
src_plane_stride,
dst_row_stride,
dst_plane_stride) &&
ds4_gpu_encode_cpy_f32_f32_3d(cb,
g_compressor_store_score_buffer,
(NSUInteger)head_dim * sizeof(float),
g_compressor_pool_score_buffer,
(NSUInteger)4u * head_dim * sizeof(float),
head_dim,
ratio,
n_comp,
src_row_stride,
src_plane_stride,
dst_row_stride,
dst_plane_stride);
}
if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) {
const uint64_t src_row_stride = (uint64_t)width * sizeof(float);
const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride;
const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float);
const uint64_t dst_plane_stride = 8ull * dst_row_stride;
ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
kvbuf,
ds4_gpu_tensor_offset(kv),
g_compressor_pool_kv_buffer,
dst_plane_stride,
head_dim,
ratio,
n_comp - 1u,
src_row_stride,
src_plane_stride,
dst_row_stride,
dst_plane_stride) &&
ds4_gpu_encode_cpy_f32_f32_3d(cb,
g_compressor_store_score_buffer,
0,
g_compressor_pool_score_buffer,
dst_plane_stride,
head_dim,
ratio,
n_comp - 1u,
src_row_stride,
src_plane_stride,
dst_row_stride,
dst_plane_stride);
}
if (ok && !use_direct_pool) {
ok = ds4_gpu_encode_dsv4_softmax_pool(cb,
comp_cache,
g_compressor_pool_kv_buffer,
0,
(uint64_t)head_dim * sizeof(float),
sizeof(float),
8ull * head_dim * sizeof(float),
g_compressor_pool_score_buffer,
0,
(uint64_t)head_dim * sizeof(float),
sizeof(float),
8ull * head_dim * sizeof(float),
8,
head_dim,
n_comp);
}
} else if (ok) {
const uint64_t row_stride = (uint64_t)width * sizeof(float);
ok = ds4_gpu_encode_dsv4_softmax_pool(cb,
comp_cache,
kvbuf,
ds4_gpu_tensor_offset(kv),
row_stride,
sizeof(float),
(uint64_t)ratio * row_stride,
g_compressor_store_score_buffer,
0,
row_stride,
sizeof(float),
(uint64_t)ratio * row_stride,
ratio,
head_dim,
n_comp);
}
}
if (ok && n_comp != 0) {
ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache,
comp_cache,
model_map,
model_size,
norm_offset,
head_dim,
n_comp,
rms_eps) != 0;
}
if (ok && n_comp != 0 && n_rot != 0) {
ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args(
n_comp, 1, head_dim, n_rot, n_ctx_orig, false,
freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
cb = ds4_gpu_command_buffer(&owned);
ok = cb && !owned &&
ds4_gpu_encode_rope_tail_inplace(cb,
compbuf,
ds4_gpu_tensor_offset(comp_cache),
&rope_args,
n_comp,
1,
head_dim,
pos0,
ratio);
}
if (ok && n_comp != 0 && quantize_fp8) {
ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0;
}
if (!had_batch) {
const int end_ok = ds4_gpu_end_commands();
ok = end_ok && ok;
}
return ok ? 1 : 0;
}
}
int ds4_gpu_compressor_prefill_ratio4_replay_tensor(
ds4_gpu_tensor *comp_cache,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *sc,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t pos0,
uint32_t n_tokens,
uint32_t n_rot,
uint32_t n_ctx_orig,
bool quantize_fp8,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
(ape_type != 0u && ape_type != 1u) ||
norm_type != 0u) {
return 0;
}
@autoreleasepool {
const uint32_t ratio = 4u;
const uint32_t width = 2u * head_dim;
const uint32_t state_rows = 8u;
const uint32_t n_comp = n_tokens / ratio;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t kv_bytes = (uint64_t)n_tokens * width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim * sizeof(float);
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
norm_offset > model_size || norm_bytes > model_size - norm_offset) {
fprintf(stderr, "ds4: Metal compressor replay tensor range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv);
id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc);
id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_cache);
id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
if (!kvbuf || !scbuf || !compbuf || !statekvbuf || !statescbuf ||
ds4_gpu_tensor_bytes(kv) < kv_bytes ||
ds4_gpu_tensor_bytes(sc) < kv_bytes ||
ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
ds4_gpu_tensor_bytes(state_score) < state_bytes ||
ds4_gpu_tensor_bytes(comp_cache) < comp_bytes) {
fprintf(stderr, "ds4: Metal compressor replay received undersized buffers\n");
return 0;
}
uint64_t ape_inner = 0;
id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner);
if (!apebuf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
int ok = 1;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) ok = 0;
const NSUInteger score_bytes = (NSUInteger)n_tokens * width * sizeof(float);
const NSUInteger pack_bytes = (NSUInteger)n_comp * 8u * head_dim * sizeof(float);
const int direct_pool_mode =
ds4_gpu_compressor_ratio4_direct_pool_mode(head_dim, n_comp);
if (ok && direct_pool_mode < 0) ok = 0;
const uint64_t direct_output_bytes =
(uint64_t)n_comp * head_dim * sizeof(float);
const uint64_t direct_input_bytes =
(uint64_t)n_tokens * width * sizeof(float);
const uint64_t direct_state_bytes =
(uint64_t)4u * width * sizeof(float);
const bool direct_pool_overlap = direct_pool_mode > 0 &&
(ds4_gpu_buffer_ranges_overlap(
compbuf,
ds4_gpu_tensor_offset(comp_cache),
direct_output_bytes,
kvbuf,
ds4_gpu_tensor_offset(kv),
direct_input_bytes) ||
ds4_gpu_buffer_ranges_overlap(
compbuf,
ds4_gpu_tensor_offset(comp_cache),
direct_output_bytes,
statekvbuf,
ds4_gpu_tensor_offset(state_kv),
direct_state_bytes) ||
ds4_gpu_buffer_ranges_overlap(
compbuf,
ds4_gpu_tensor_offset(comp_cache),
direct_output_bytes,
statescbuf,
ds4_gpu_tensor_offset(state_score),
direct_state_bytes));
const bool use_direct_pool =
direct_pool_mode > 0 && !direct_pool_overlap;
if (ok && (!ds4_gpu_ensure_scratch_buffer(&g_compressor_store_score_buffer,
&g_compressor_store_score_bytes,
score_bytes,
"ds4_compressor_store_score") ||
(!use_direct_pool &&
(!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_kv_buffer,
&g_compressor_pool_kv_bytes,
pack_bytes,
"ds4_compressor_pool_kv") ||
!ds4_gpu_ensure_scratch_buffer(&g_compressor_pool_score_buffer,
&g_compressor_pool_score_bytes,
pack_bytes,
"ds4_compressor_pool_score"))))) {
ok = 0;
}
if (ok) {
ok = ds4_gpu_encode_compressor_score_with_ape(cb,
scbuf,
ds4_gpu_tensor_offset(sc),
g_compressor_store_score_buffer,
0,
apebuf,
(NSUInteger)ape_inner,
ape_type,
width,
ratio,
pos0,
n_tokens);
}
const uint64_t src_row_stride = (uint64_t)width * sizeof(float);
const uint64_t src_plane_stride = (uint64_t)ratio * src_row_stride;
const uint64_t dst_row_stride = (uint64_t)head_dim * sizeof(float);
const uint64_t dst_plane_stride = 8ull * dst_row_stride;
const NSUInteger state_off = ds4_gpu_tensor_offset(state_kv);
const NSUInteger state_score_off = ds4_gpu_tensor_offset(state_score);
const int pack_fusion_mode = use_direct_pool ? 0 :
ds4_gpu_compressor_pack_ratio4_fusion_mode(head_dim);
if (ok && pack_fusion_mode < 0) ok = 0;
const bool use_pack_fusion = pack_fusion_mode > 0;
if (ok && use_direct_pool) {
ok = ds4_gpu_encode_compressor_ratio4_direct_pool(
cb,
comp_cache,
kvbuf,
ds4_gpu_tensor_offset(kv),
g_compressor_store_score_buffer,
0,
statekvbuf,
state_off,
statescbuf,
state_score_off,
head_dim,
n_comp,
true);
}
if (ok && use_pack_fusion) {
ok = ds4_gpu_encode_compressor_pack_ratio4(
cb,
kvbuf,
ds4_gpu_tensor_offset(kv),
g_compressor_store_score_buffer,
0,
statekvbuf,
state_off,
statescbuf,
state_score_off,
head_dim,
n_comp,
true);
}
if (ok && !use_direct_pool && !use_pack_fusion) {
ok = ds4_gpu_encode_fill_f32_rows(cb,
g_compressor_pool_kv_buffer,
0,
head_dim,
8u * n_comp,
0.0f) &&
ds4_gpu_encode_fill_f32_rows(cb,
g_compressor_pool_score_buffer,
0,
head_dim,
8u * n_comp,
ds4_gpu_negative_infinity());
}
if (ok && !use_direct_pool && !use_pack_fusion) {
/*
* The aligned nonzero ratio-4 path replays the current ubatch
* compressor, but seeds the first compressed row with the previous
* compressor state. Rows 0..3 are the previous half, rows 4..7 are
* the current half.
*/
ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
statekvbuf,
state_off,
g_compressor_pool_kv_buffer,
0,
head_dim,
ratio,
1,
src_row_stride,
(uint64_t)ratio * src_row_stride,
dst_row_stride,
dst_plane_stride) &&
ds4_gpu_encode_cpy_f32_f32_3d(cb,
statescbuf,
state_score_off,
g_compressor_pool_score_buffer,
0,
head_dim,
ratio,
1,
src_row_stride,
(uint64_t)ratio * src_row_stride,
dst_row_stride,
dst_plane_stride);
}
if (ok && !use_direct_pool && !use_pack_fusion) {
ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
kvbuf,
ds4_gpu_tensor_offset(kv) +
(NSUInteger)head_dim * sizeof(float),
g_compressor_pool_kv_buffer,
(NSUInteger)4u * head_dim * sizeof(float),
head_dim,
ratio,
n_comp,
src_row_stride,
src_plane_stride,
dst_row_stride,
dst_plane_stride) &&
ds4_gpu_encode_cpy_f32_f32_3d(cb,
g_compressor_store_score_buffer,
(NSUInteger)head_dim * sizeof(float),
g_compressor_pool_score_buffer,
(NSUInteger)4u * head_dim * sizeof(float),
head_dim,
ratio,
n_comp,
src_row_stride,
src_plane_stride,
dst_row_stride,
dst_plane_stride);
}
if (ok && !use_direct_pool && !use_pack_fusion && n_comp > 1u) {
ok = ds4_gpu_encode_cpy_f32_f32_3d(cb,
kvbuf,
ds4_gpu_tensor_offset(kv),
g_compressor_pool_kv_buffer,
dst_plane_stride,
head_dim,
ratio,
n_comp - 1u,
src_row_stride,
src_plane_stride,
dst_row_stride,
dst_plane_stride) &&
ds4_gpu_encode_cpy_f32_f32_3d(cb,
g_compressor_store_score_buffer,
0,
g_compressor_pool_score_buffer,
dst_plane_stride,
head_dim,
ratio,
n_comp - 1u,
src_row_stride,
src_plane_stride,
dst_row_stride,
dst_plane_stride);
}
if (ok && !use_direct_pool) {
ok = ds4_gpu_encode_dsv4_softmax_pool(cb,
comp_cache,
g_compressor_pool_kv_buffer,
0,
dst_row_stride,
sizeof(float),
dst_plane_stride,
g_compressor_pool_score_buffer,
0,
dst_row_stride,
sizeof(float),
dst_plane_stride,
8,
head_dim,
n_comp);
}
if (ok) {
ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_cache,
comp_cache,
model_map,
model_size,
norm_offset,
head_dim,
n_comp,
rms_eps) != 0;
}
if (ok && n_rot != 0) {
ds4_gpu_rope_tail_batch_args rope_args = ds4_gpu_make_rope_tail_args(
n_comp, 1, head_dim, n_rot, n_ctx_orig, false,
freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
cb = ds4_gpu_command_buffer(&owned);
ok = cb && !owned &&
ds4_gpu_encode_rope_tail_inplace(cb,
compbuf,
ds4_gpu_tensor_offset(comp_cache),
&rope_args,
n_comp,
1,
head_dim,
pos0,
ratio);
}
if (ok && quantize_fp8) {
ok = ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot) != 0;
}
if (ok) {
ok = ds4_gpu_encode_fill_f32_rows(cb,
statekvbuf,
state_off,
width,
state_rows,
0.0f) &&
ds4_gpu_encode_fill_f32_rows(cb,
statescbuf,
state_score_off,
width,
state_rows,
ds4_gpu_negative_infinity());
}
if (ok) {
int32_t rows_prev[4] = { 0, 1, 2, 3 };
const uint32_t prev_start = n_tokens - ratio;
ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
state_kv,
state_score,
kvbuf,
ds4_gpu_tensor_offset(kv) +
(NSUInteger)prev_start * width * sizeof(float),
scbuf,
ds4_gpu_tensor_offset(sc) +
(NSUInteger)prev_start * width * sizeof(float),
apebuf,
(NSUInteger)ape_inner,
ape_type,
width,
ratio,
pos0 + prev_start,
rows_prev,
ratio,
state_rows);
}
if (!had_batch) {
const int end_ok = ds4_gpu_end_commands();
ok = end_ok && ok;
}
return ok ? 1 : 0;
}
}
int ds4_gpu_compressor_prefill_state_ratio4_tensor(
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
const ds4_gpu_tensor *kv_tail,
const ds4_gpu_tensor *sc_tail,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint32_t head_dim,
uint32_t pos0) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map ||
head_dim == 0 || (ape_type != 0u && ape_type != 1u)) {
return 0;
}
@autoreleasepool {
const uint32_t ratio = 4u;
const uint32_t width = 2u * head_dim;
const uint32_t state_rows = 8u;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t tail_bytes = (uint64_t)ratio * width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t ape_bytes = (uint64_t)ratio * width * elem_ape;
if (ape_offset > model_size || ape_bytes > model_size - ape_offset) {
fprintf(stderr, "ds4: Metal compressor prefill-state APE range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv_tail);
id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc_tail);
id<MTLBuffer> statekvbuf = ds4_gpu_tensor_buffer(state_kv);
id<MTLBuffer> statescbuf = ds4_gpu_tensor_buffer(state_score);
if (!kvbuf || !scbuf || !statekvbuf || !statescbuf ||
ds4_gpu_tensor_bytes(kv_tail) < tail_bytes ||
ds4_gpu_tensor_bytes(sc_tail) < tail_bytes ||
ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
ds4_gpu_tensor_bytes(state_score) < state_bytes) {
fprintf(stderr, "ds4: Metal compressor prefill-state received undersized buffers\n");
return 0;
}
uint64_t ape_inner = 0;
id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(model_map, model_size, ape_offset, ape_bytes, &ape_inner);
if (!apebuf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
int ok = 1;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) ok = 0;
if (ok) {
ok = ds4_gpu_encode_fill_f32_rows(cb,
statekvbuf,
ds4_gpu_tensor_offset(state_kv),
width,
state_rows,
0.0f) &&
ds4_gpu_encode_fill_f32_rows(cb,
statescbuf,
ds4_gpu_tensor_offset(state_score),
width,
state_rows,
ds4_gpu_negative_infinity());
}
if (ok) {
int32_t rows[4] = { 0, 1, 2, 3 };
ok = ds4_gpu_encode_compressor_set_rows_projected(cb,
state_kv,
state_score,
kvbuf,
ds4_gpu_tensor_offset(kv_tail),
scbuf,
ds4_gpu_tensor_offset(sc_tail),
apebuf,
(NSUInteger)ape_inner,
ape_type,
width,
ratio,
pos0,
rows,
ratio,
state_rows);
}
if (!had_batch) {
const int end_ok = ds4_gpu_end_commands();
ok = end_ok && ok;
}
return ok ? 1 : 0;
}
}
int ds4_gpu_compressor_update_tensor(
const ds4_gpu_tensor *kv_cur,
const ds4_gpu_tensor *sc_cur,
ds4_gpu_tensor *state_kv,
ds4_gpu_tensor *state_score,
ds4_gpu_tensor *comp_cache,
const void *model_map,
uint64_t model_size,
uint64_t ape_offset,
uint32_t ape_type,
uint64_t norm_offset,
uint32_t norm_type,
uint32_t head_dim,
uint32_t ratio,
uint32_t pos,
uint32_t comp_row,
uint32_t n_rot,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
float rms_eps,
bool state_already_stored,
bool decode_one_token,
bool defer_finalize) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache ||
!model_map || head_dim == 0 || ratio == 0 ||
n_rot > head_dim || (n_rot & 1u) != 0 ||
(ape_type != 0u && ape_type != 1u) ||
norm_type != 0u) {
return 0;
}
@autoreleasepool {
const uint32_t coff = ratio == 4u ? 2u : 1u;
const uint32_t width = coff * head_dim;
const uint32_t state_rows = coff * ratio;
const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u;
const uint64_t elem_ape = ape_type == 1u ? 2u : 4u;
const uint64_t kv_bytes = (uint64_t)width * sizeof(float);
const uint64_t state_bytes = (uint64_t)state_rows * width * sizeof(float);
const uint64_t comp_bytes = (uint64_t)(comp_row + (emit ? 1u : 0u)) * head_dim * sizeof(float);
const uint64_t ape_bytes = (uint64_t)width * ratio * elem_ape;
const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
if (ape_offset > model_size || ape_bytes > model_size - ape_offset ||
norm_offset > model_size || norm_bytes > model_size - norm_offset) {
fprintf(stderr, "ds4: Metal compressor tensor range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv_cur);
id<MTLBuffer> scbuf = ds4_gpu_tensor_buffer(sc_cur);
id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_cache);
if (!kvbuf || !scbuf || !compbuf ||
ds4_gpu_tensor_bytes(kv_cur) < kv_bytes ||
ds4_gpu_tensor_bytes(sc_cur) < kv_bytes ||
ds4_gpu_tensor_bytes(state_kv) < state_bytes ||
ds4_gpu_tensor_bytes(state_score) < state_bytes ||
(emit && ds4_gpu_tensor_bytes(comp_cache) < comp_bytes)) {
fprintf(stderr, "ds4: Metal compressor update received undersized buffers\n");
return 0;
}
if (!state_already_stored) {
const bool use_store_one =
getenv("DS4_METAL_DISABLE_COMPRESSOR_STORE_ONE") == NULL;
const int store_ok = use_store_one
? ds4_gpu_compressor_store_one_tensor(kv_cur,
sc_cur,
state_kv,
state_score,
model_map,
model_size,
ape_offset,
ape_type,
width,
ratio,
pos)
: ds4_gpu_compressor_store_batch_tensor(kv_cur,
sc_cur,
state_kv,
state_score,
model_map,
model_size,
ape_offset,
ape_type,
head_dim,
ratio,
pos,
1);
if (!store_ok) {
return 0;
}
}
if (!emit) return 1;
ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view(
comp_cache,
(uint64_t)comp_row * head_dim * sizeof(float),
(uint64_t)head_dim * sizeof(float));
if (!comp_row_view) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
int ok = cb &&
ds4_gpu_encode_compressor_pool(cb,
comp_row_view,
state_kv,
state_score,
head_dim,
ratio,
decode_one_token);
if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor DS4 softmax pool");
/* When the caller runs the fused finalize dispatch afterwards, the
* pooled row is consumed un-normalized there; skip the standalone
* norm+rope so the work happens exactly once. The state shift still
* runs: it only moves the rolling state rows. */
if (ok && !defer_finalize) {
ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view,
comp_row_view,
model_map,
model_size,
norm_offset,
head_dim,
1,
rms_eps) != 0;
}
if (ok && !defer_finalize) {
const uint32_t comp_pos = pos + 1u - ratio;
ok = ds4_gpu_rope_tail_tensor(comp_row_view,
1,
1,
head_dim,
n_rot,
comp_pos,
n_ctx_orig,
false,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow) != 0;
}
if (ok && ratio == 4u && !defer_finalize) {
cb = ds4_gpu_command_buffer(&owned);
ok = cb &&
ds4_gpu_encode_compressor_shift_ratio4(cb,
state_kv,
state_score,
width);
if (ok) ok = ds4_gpu_finish_command_buffer(cb, owned, "compressor ratio4 state shift");
}
ds4_gpu_tensor_free(comp_row_view);
if (!ok) return 0;
}
return 1;
}
static int ds4_gpu_encode_fill_f32_rows(
id<MTLCommandBuffer> cb,
id<MTLBuffer> buf,
NSUInteger offset,
uint32_t width,
uint32_t rows,
float value) {
if (!cb || !buf || width == 0 || rows == 0 || (width & 3u) != 0) return 0;
ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, 1, 0.0f, 0.0f);
args.val = value;
NSUInteger nth_max = g_unary_fill_pipeline.maxTotalThreadsPerThreadgroup;
if (nth_max > 256u) nth_max = 256u;
NSUInteger nth = (NSUInteger)args.ne00;
if (nth > nth_max) nth = nth_max;
if (nth == 0) nth = 1u;
const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_unary_fill_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:buf offset:offset atIndex:1];
[enc setBuffer:buf offset:offset atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01,
(NSUInteger)args.ne02,
(NSUInteger)args.ne03)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_matmul_q8_0_kslice_exact_rows(
ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint64_t full_in_dim, uint64_t k_off,
uint64_t k_cnt, uint64_t out_dim, const ds4_gpu_tensor *x,
uint64_t x_elem_off, uint32_t n_rows);
static int ds4_gpu_matmul_q8_0_kslice_rows_impl(
ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint64_t full_in_dim, uint64_t out_dim,
uint64_t k_off, uint64_t k_cnt, const ds4_gpu_tensor *x,
uint64_t n_rows, uint64_t x_stride, uint64_t x_off);
static int ds4_gpu_attention_output_q8_batch_impl(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
ds4_gpu_tensor *group_tmp,
ds4_gpu_tensor *low_tmp,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint64_t out_dim,
const ds4_gpu_tensor *heads,
uint32_t n_tokens,
bool round_low_bf16,
uint64_t full_low_dim,
uint64_t low_offset) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !low || !group_tmp || !low_tmp || !heads || !model_map ||
group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 ||
group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) {
return 0;
}
@autoreleasepool {
const uint64_t low_dim = (uint64_t)n_groups * rank;
if (!full_low_dim) full_low_dim = low_dim;
if ((group_dim % 32u) != 0 || (low_dim % 32u) != 0 || low_dim > UINT32_MAX ||
full_low_dim > UINT32_MAX || (full_low_dim % 32u) || (low_offset % 32u) ||
low_offset > full_low_dim || low_dim > full_low_dim - low_offset) {
fprintf(stderr, "ds4: Metal attention output batch received invalid q8 dimensions\n");
return 0;
}
const uint64_t row_a_bytes = (group_dim / 32u) * 34u;
const uint64_t row_b_bytes = (full_low_dim / 32u) * 34u;
const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes;
const uint64_t out_b_bytes = out_dim * row_b_bytes;
if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset ||
out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) {
fprintf(stderr, "ds4: Metal attention output batch weights are outside the mapped model\n");
return 0;
}
const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float);
const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float);
if (ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(low) < low_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal attention output batch received undersized buffers\n");
return 0;
}
(void)group_tmp;
(void)low_tmp;
const bool use_direct_low =
n_tokens < 32u && getenv("DS4_METAL_DISABLE_ATTN_OUT_LOW_DIRECT") == NULL;
/* The exported TensorOps attention-output kernel is a 64-token tile.
* Keep this on full tiles only; smaller multiples of 32 use the legacy
* path instead of relying on cooperative tensor partial RHS bounds. */
const bool use_mpp_low =
n_tokens >= 32u &&
(n_tokens % DS4_METAL_ATTN_OUT_MPP_TILE_N) == 0 &&
ds4_gpu_use_mpp_attn_out_low_matmul();
const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t);
id<MTLBuffer> group_ids_buffer = nil;
if (!use_direct_low && !use_mpp_low) {
if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) {
group_ids_buffer =
ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids");
if (!group_ids_buffer) {
return 0;
}
} else {
if (!ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer,
&g_attn_out_group_ids_bytes,
ids_bytes,
"ds4_attention_output_group_ids")) {
return 0;
}
group_ids_buffer = g_attn_out_group_ids_buffer;
}
int32_t *ids = (int32_t *)[group_ids_buffer contents];
for (uint32_t t = 0; t < n_tokens; t++) {
for (uint32_t group = 0; group < n_groups; group++) {
ids[(uint64_t)t * n_groups + group] = (int32_t)group;
}
}
}
uint64_t out_a_inner = 0;
id<MTLBuffer> out_a_buf =
ds4_gpu_wrap_model_range(model_map, model_size,
out_a_offset, out_a_bytes,
&out_a_inner);
if (!out_a_buf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
bool ok = true;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) {
ok = false;
}
const bool attn_out_profile =
getenv("DS4_METAL_ATTN_OUT_STAGE_PROFILE") != NULL && g_batch_cb != nil;
if (ok && attn_out_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
ok = false;
} else {
cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) ok = false;
}
}
double attn_out_t0 = attn_out_profile ? ds4_gpu_now_ms() : 0.0;
#define DS4_METAL_PROFILE_ATTN_OUT_STAGE(name) do { \
if (ok && attn_out_profile) { \
if (ds4_gpu_end_commands() == 0) { \
ok = false; \
} else { \
const double now_ms = ds4_gpu_now_ms(); \
fprintf(stderr, \
"ds4: Metal attention output stage tokens=%u %s=%.3f ms\n", \
n_tokens, (name), now_ms - attn_out_t0); \
attn_out_t0 = now_ms; \
if (ds4_gpu_begin_commands() == 0) { \
ok = false; \
} else { \
cb = ds4_gpu_command_buffer(&owned); \
if (!cb || owned) ok = false; \
} \
} \
} \
} while (0)
if (ok) {
/*
* Batched attention-output projections switch from the vector
* kernel to the SIMD matrix kernel once the batch has at least 32
* tokens. This preserves the single-token generation path while
* keeping prefill accumulation stable.
*/
if (use_mpp_low) {
ds4_gpu_mul_mm_id_args mm_args =
ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim,
(uint32_t)rank,
n_groups,
row_a_bytes,
(uint64_t)rank * row_a_bytes,
n_groups,
n_groups,
n_tokens);
/*
* Direct RHS lets MPP read the dense low-rank activation tile
* directly from device memory instead of staging a second
* threadgroup tile. The retained attention-output path is the
* 64-token direct-RHS kernel; the older staged-RHS and 32-token
* variants were not kept as alternate runtime modes.
*/
const char *attn_out_pipeline_name =
"kernel_attn_out_low_q8_0_mpp_direct_rhs_n64";
id<MTLComputePipelineState> mm_pipeline =
ds4_gpu_get_mul_mm_id_pipeline(attn_out_pipeline_name, false);
ok = ds4_gpu_encode_attn_out_low_mpp(cb,
mm_pipeline,
&mm_args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low)) != 0;
if (!ok) {
ds4_gpu_warn_mpp_fallback();
if (ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) {
if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) {
group_ids_buffer =
ds4_gpu_new_transient_buffer(ids_bytes, "attention output group ids");
} else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer,
&g_attn_out_group_ids_bytes,
ids_bytes,
"ds4_attention_output_group_ids")) {
group_ids_buffer = g_attn_out_group_ids_buffer;
}
if (group_ids_buffer) {
int32_t *ids = (int32_t *)[group_ids_buffer contents];
for (uint32_t t = 0; t < n_tokens; t++) {
for (uint32_t group = 0; group < n_groups; group++) {
ids[(uint64_t)t * n_groups + group] = (int32_t)group;
}
}
ds4_gpu_mul_mm_id_map_args map_args =
ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim,
n_groups,
n_groups,
n_groups,
n_tokens);
id<MTLComputePipelineState> map_pipeline =
ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups));
id<MTLComputePipelineState> fallback_pipeline =
ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false);
ok = ds4_gpu_encode_mul_mm_id(cb,
map_pipeline,
fallback_pipeline,
&map_args,
&mm_args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low),
group_ids_buffer,
0) != 0;
}
}
}
} else if (n_tokens >= 32u && ds4_gpu_mul_mm_id_map0_name(n_groups) != NULL) {
ds4_gpu_mul_mm_id_map_args map_args =
ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim,
n_groups,
n_groups,
n_groups,
n_tokens);
ds4_gpu_mul_mm_id_args mm_args =
ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim,
(uint32_t)rank,
n_groups,
row_a_bytes,
(uint64_t)rank * row_a_bytes,
n_groups,
n_groups,
n_tokens);
id<MTLComputePipelineState> map_pipeline =
ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups));
id<MTLComputePipelineState> mm_pipeline =
ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false);
ok = ds4_gpu_encode_mul_mm_id(cb,
map_pipeline,
mm_pipeline,
&map_args,
&mm_args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low),
group_ids_buffer,
0) != 0;
} else if (use_direct_low) {
ds4_gpu_mul_mv_id_args args = {
.nei0 = (int32_t)n_groups,
.nei1 = (int32_t)n_tokens,
.nbi1 = 0,
.ne00 = (int32_t)group_dim,
.ne01 = (int32_t)rank,
.ne02 = (int32_t)n_groups,
.nb00 = 34,
.nb01 = row_a_bytes,
.nb02 = (uint64_t)rank * row_a_bytes,
.ne10 = (int32_t)group_dim,
.ne11 = (int32_t)n_groups,
.ne12 = (int32_t)n_tokens,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = (uint64_t)group_dim * sizeof(float),
.nb12 = (uint64_t)n_groups * group_dim * sizeof(float),
.ne0 = (int32_t)rank,
.ne1 = (int32_t)n_groups,
.nb1 = (uint64_t)rank * sizeof(float),
.nr0 = 2,
};
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4);
ok = ds4_gpu_encode_attn_out_low_q8_direct(cb,
pipeline,
&args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low),
32u * 2u * sizeof(float),
4,
true) != 0;
} else {
ds4_gpu_mul_mv_id_args args = {
.nei0 = (int32_t)n_groups,
.nei1 = (int32_t)n_tokens,
.nbi1 = (uint64_t)n_groups * sizeof(int32_t),
.ne00 = (int32_t)group_dim,
.ne01 = (int32_t)rank,
.ne02 = (int32_t)n_groups,
.nb00 = 34,
.nb01 = row_a_bytes,
.nb02 = (uint64_t)rank * row_a_bytes,
.ne10 = (int32_t)group_dim,
.ne11 = (int32_t)n_groups,
.ne12 = (int32_t)n_tokens,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = (uint64_t)group_dim * sizeof(float),
.nb12 = (uint64_t)n_groups * group_dim * sizeof(float),
.ne0 = (int32_t)rank,
.ne1 = (int32_t)n_groups,
.nb1 = (uint64_t)rank * sizeof(float),
.nr0 = 2,
};
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4);
ok = ds4_gpu_encode_mul_mv_id(cb,
pipeline,
&args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low),
group_ids_buffer,
0,
32u * 2u * sizeof(float),
4,
true) != 0;
}
}
DS4_METAL_PROFILE_ATTN_OUT_STAGE("low_proj");
if (ok && round_low_bf16)
ok = ds4_gpu_dsv41_quantize(low, (uint32_t)low_dim, n_tokens, DS4_V41_BF16);
if (ok && full_low_dim != low_dim) {
const uint32_t matrix_rows = ds4_gpu_mpp_available() ? n_tokens / 32u * 32u : 0;
if (matrix_rows) {
ok = ds4_gpu_matmul_q8_0_kslice_rows_impl(out, model_map, model_size,
out_b_offset, full_low_dim, out_dim, low_offset, low_dim,
low, matrix_rows, low_dim, 0) != 0;
}
/* Older GPUs use the exact matvec reduction for every row.
* Dispatch them together instead of encoding one kernel per token. */
if (ok && matrix_rows < n_tokens) {
const uint32_t rows = n_tokens - matrix_rows;
ds4_gpu_tensor *tail = ds4_gpu_tensor_view(out,
(uint64_t)matrix_rows * out_dim * sizeof(float),
(uint64_t)rows * out_dim * sizeof(float));
ok = tail && ds4_gpu_matmul_q8_0_kslice_exact_rows(tail, model_map, model_size,
out_b_offset, full_low_dim, low_offset, low_dim, out_dim,
low, (uint64_t)matrix_rows * low_dim, rows);
ds4_gpu_tensor_free(tail);
}
} else if (ok) {
ok = ds4_gpu_matmul_q8_0_tensor(out, model_map, model_size,
out_b_offset,
low_dim, out_dim, low, n_tokens) != 0;
}
DS4_METAL_PROFILE_ATTN_OUT_STAGE("out_proj");
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
#undef DS4_METAL_PROFILE_ATTN_OUT_STAGE
return ok ? 1 : 0;
}
}
int ds4_gpu_attention_output_q8_batch_tensor(
ds4_gpu_tensor *out, ds4_gpu_tensor *low,
ds4_gpu_tensor *group_tmp, ds4_gpu_tensor *low_tmp,
const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint64_t out_b_offset,
uint64_t group_dim, uint64_t rank, uint32_t n_groups, uint64_t out_dim,
const ds4_gpu_tensor *heads, uint32_t n_tokens) {
return ds4_gpu_attention_output_q8_batch_impl(out, low, group_tmp, low_tmp,
model_map, model_size, out_a_offset, out_b_offset, group_dim, rank,
n_groups, out_dim, heads, n_tokens, false, 0, 0);
}
int ds4_gpu_dsv41_attention_output_batch(
ds4_gpu_tensor *out, ds4_gpu_tensor *low,
const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint64_t out_b_offset,
const ds4_gpu_tensor *heads, uint32_t n_tokens) {
return ds4_gpu_attention_output_q8_batch_impl(out, low, low, low,
model_map, model_size, out_a_offset, out_b_offset, 4096, 1024,
8, 5120, heads, n_tokens, true, 0, 0);
}
int ds4_gpu_dsv41_attention_output_tp_batch(
ds4_gpu_tensor *out, ds4_gpu_tensor *low,
const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint64_t out_b_offset,
const ds4_gpu_tensor *heads, uint32_t n_tokens, uint32_t tp_rank) {
const uint64_t shard_bytes = UINT64_C(4) * 1024 * (4096 / 32 * 34);
if (tp_rank > 1 || out_a_offset > model_size ||
shard_bytes * 2u > model_size - out_a_offset) return 0;
return ds4_gpu_attention_output_q8_batch_impl(out, low, low, low,
model_map, model_size, out_a_offset + tp_rank * shard_bytes, out_b_offset,
4096, 1024, 4, 5120, heads, n_tokens, true, 8192, tp_rank * 4096u);
}
int ds4_gpu_attention_output_q4_K_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
ds4_gpu_tensor *group_tmp,
ds4_gpu_tensor *low_tmp,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint32_t out_b_type,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint64_t out_dim,
const ds4_gpu_tensor *heads,
uint32_t n_tokens) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !low || !heads || !model_map ||
group_dim == 0 || rank == 0 || n_groups == 0 || out_dim == 0 || n_tokens == 0 ||
group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) {
return 0;
}
if (n_tokens < 32u) return 0;
@autoreleasepool {
const uint64_t low_dim = (uint64_t)n_groups * rank;
if ((group_dim % 256u) != 0 || (low_dim % 256u) != 0 || low_dim > UINT32_MAX) {
return 0;
}
uint64_t row_a_bytes = 0;
uint64_t row_b_bytes = 0;
if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K, (uint32_t)group_dim, &row_a_bytes) ||
!ds4_gpu_quant_row_bytes(out_b_type, (uint32_t)low_dim, &row_b_bytes)) {
return 0;
}
const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes;
const uint64_t out_b_bytes = out_dim * row_b_bytes;
if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset ||
out_b_offset > model_size || out_b_bytes > model_size - out_b_offset) {
fprintf(stderr, "ds4: Metal Q4 attention output batch weights are outside the mapped model\n");
return 0;
}
const uint64_t heads_bytes = (uint64_t)n_tokens * n_groups * group_dim * sizeof(float);
const uint64_t low_bytes = (uint64_t)n_tokens * low_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float);
if (ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(low) < low_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal Q4 attention output batch received undersized buffers\n");
return 0;
}
(void)group_tmp;
(void)low_tmp;
const uint64_t padded_n_tokens_u64 =
((uint64_t)n_tokens + DS4_METAL_ATTN_OUT_MPP_TILE_N - 1u) /
DS4_METAL_ATTN_OUT_MPP_TILE_N * DS4_METAL_ATTN_OUT_MPP_TILE_N;
const uint32_t padded_n_tokens =
padded_n_tokens_u64 <= UINT32_MAX ? (uint32_t)padded_n_tokens_u64 : n_tokens;
/*
* TensorOps requires complete 64-row tiles. For a substantial final
* prefill chunk, compute harmless extra rows when the caller's
* workspace already covers them; the following projection consumes
* only the real rows. Small tails stay on the legacy kernel because
* padding overhead outweighs the faster arithmetic there.
*/
const bool use_mpp_padding =
n_tokens >= 256u &&
padded_n_tokens > n_tokens &&
ds4_gpu_tensor_bytes(heads) >=
(uint64_t)padded_n_tokens * n_groups * group_dim * sizeof(float) &&
ds4_gpu_tensor_bytes(low) >=
(uint64_t)padded_n_tokens * low_dim * sizeof(float);
bool use_mpp_low =
((n_tokens % DS4_METAL_ATTN_OUT_MPP_TILE_N) == 0 ||
use_mpp_padding) &&
ds4_gpu_use_mpp_attn_out_low_matmul();
id<MTLComputePipelineState> mpp_low_pipeline = nil;
if (use_mpp_low) {
mpp_low_pipeline = ds4_gpu_get_mul_mm_id_pipeline(
"kernel_attn_out_low_q4_K_mpp_direct_rhs_n64", false);
if (!mpp_low_pipeline) {
ds4_gpu_warn_mpp_fallback();
use_mpp_low = false;
}
}
const NSUInteger ids_bytes = (NSUInteger)n_tokens * (NSUInteger)n_groups * sizeof(int32_t);
id<MTLBuffer> group_ids_buffer = nil;
if (!use_mpp_low) {
if (getenv("DS4_METAL_DISABLE_ATTN_OUT_IDS_CACHE") != NULL) {
group_ids_buffer =
ds4_gpu_new_transient_buffer(ids_bytes, "attention output Q4 group ids");
} else if (ds4_gpu_ensure_scratch_buffer(&g_attn_out_group_ids_buffer,
&g_attn_out_group_ids_bytes,
ids_bytes,
"ds4_attention_output_group_ids")) {
group_ids_buffer = g_attn_out_group_ids_buffer;
}
if (!group_ids_buffer) return 0;
int32_t *ids = (int32_t *)[group_ids_buffer contents];
for (uint32_t t = 0; t < n_tokens; t++) {
for (uint32_t group = 0; group < n_groups; group++) {
ids[(uint64_t)t * n_groups + group] = (int32_t)group;
}
}
}
uint64_t out_a_inner = 0;
id<MTLBuffer> out_a_buf =
ds4_gpu_wrap_model_range(model_map, model_size,
out_a_offset, out_a_bytes,
&out_a_inner);
if (!out_a_buf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
bool ok = true;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) ok = false;
if (ok) {
ds4_gpu_mul_mm_id_args mm_args =
ds4_gpu_make_mul_mm_id_args((uint32_t)group_dim,
(uint32_t)rank,
n_groups,
row_a_bytes,
(uint64_t)rank * row_a_bytes,
n_groups,
n_groups,
use_mpp_low && use_mpp_padding
? padded_n_tokens
: n_tokens);
if (use_mpp_low) {
ok = ds4_gpu_encode_attn_out_low_mpp(
cb,
mpp_low_pipeline,
&mm_args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low)) != 0;
} else {
ds4_gpu_mul_mm_id_map_args map_args =
ds4_gpu_make_mul_mm_id_map_args((uint32_t)group_dim,
n_groups,
n_groups,
n_groups,
n_tokens);
id<MTLComputePipelineState> map_pipeline =
ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_groups));
id<MTLComputePipelineState> mm_pipeline =
ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false);
ok = ds4_gpu_encode_mul_mm_id(cb,
map_pipeline,
mm_pipeline,
&map_args,
&mm_args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low),
group_ids_buffer,
0) != 0;
}
}
if (ok) {
ok = ds4_gpu_matmul_quant_tensor(out,
model_map,
model_size,
out_b_offset,
out_b_type,
low_dim,
out_dim,
low,
n_tokens) != 0;
}
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
return ok ? 1 : 0;
}
}
int ds4_gpu_attention_output_q8_batch_f16_tensor(
ds4_gpu_tensor *out_h,
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
uint64_t out_dim,
const ds4_gpu_tensor *heads,
uint32_t n_tokens) {
(void)out_h; (void)low; (void)model_map; (void)model_size;
(void)out_a_offset; (void)out_b_offset; (void)group_dim; (void)rank;
(void)n_groups; (void)out_dim; (void)heads; (void)n_tokens;
return 0;
}
static int ds4_gpu_matmul_q8_0_kslice_exact_rows(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t full_in_dim,
uint64_t k_off,
uint64_t k_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t x_elem_off,
uint32_t n_rows) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if ((full_in_dim & 31u) != 0 || (k_off & 31u) != 0 || (k_cnt & 31u) != 0 ||
!n_rows || n_rows > INT32_MAX || k_cnt == 0 || k_off > full_in_dim ||
k_cnt > full_in_dim - k_off ||
full_in_dim > INT32_MAX || out_dim > INT32_MAX ||
x_elem_off > UINT64_MAX / sizeof(float) - (uint64_t)n_rows * k_cnt) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < (x_elem_off + (uint64_t)n_rows * k_cnt) * sizeof(float) ||
ds4_gpu_tensor_bytes(out) < (uint64_t)n_rows * out_dim * sizeof(float)) {
fprintf(stderr, "ds4: Metal Q8_0 kslice matmul received undersized buffers\n");
return 0;
}
const uint64_t row_bytes = (full_in_dim / 32u) * 34u;
const uint64_t weight_bytes = out_dim * row_bytes;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal Q8_0 kslice weights are outside the mapped model\n");
return 0;
}
uint64_t inner = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
weight_offset, weight_bytes, &inner);
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
/* Same matvec kernel as the full projection: ne00 bounds the k loop
* while nb01/nb02 keep the full-row stride, so each row reads only
* the owned k window. */
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(full_in_dim, out_dim);
mv_args.ne00 = (int32_t)k_cnt;
mv_args.ne10 = (int32_t)k_cnt;
mv_args.ne11 = (int32_t)n_rows;
mv_args.ne1 = (int32_t)n_rows;
mv_args.nb11 = k_cnt * sizeof(float);
mv_args.nb12 = (uint64_t)n_rows * k_cnt * sizeof(float);
mv_args.nb13 = mv_args.nb12;
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
if (out_dim > 65536u) mv_dispatch.nsg = 8;
/* Reuse the input across more output rows without changing the K
* traversal or SIMD reduction used by TP attention partials. */
if (n_rows >= 32u && full_in_dim == 8192u && k_cnt == 4096u &&
out_dim == 5120u && mv_dispatch.nsg == 2 && ds4_gpu_tp_world_is_two() &&
ds4_gpu_device_name_contains("M3 Ultra") &&
!getenv("DS4_METAL_DISABLE_V41_PREFILL_Q8_ROWS8")) {
mv_dispatch.function_name = "kernel_mul_mv_q8_0_f32_rows8";
mv_dispatch.nr0 = 8;
mv_dispatch.smem = 32u * 8u * sizeof(float);
}
mv_args.nr0 = mv_dispatch.nr0;
/* TP partial producer: publish the gate's checked flag from the
* same kernel (last-arriving threadgroup) when requested. */
uint32_t fold_slot = 0, fold_value = 0;
const int fold = n_rows == 1u && !owned &&
ds4_gpu_tp_flag_fold_take(out, out_dim * sizeof(float), &fold_slot, &fold_value);
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(
fold ? "kernel_dsv4_mul_mv_q8_0_f32_tp_flag_checked"
: mv_dispatch.function_name,
mv_dispatch.nsg);
if (!pipeline) return 0;
const NSUInteger ntg =
((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)(inner + (k_off / 32u) * 34u) atIndex:1];
[enc setBuffer:xbuf
offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float))
atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
if (fold) {
const uint32_t ntg32 = (uint32_t)ntg;
[enc setBuffer:g_tp_slab_buffer
offset:(NSUInteger)(g_tp_slab_buffer_off + g_tp_gpu_flags_off + (uint64_t)fold_slot * 4u)
atIndex:4];
[enc setBuffer:g_tp_check_buffer offset:(NSUInteger)fold_slot * 4u atIndex:5];
[enc setBytes:&fold_value length:sizeof(fold_value) atIndex:6];
[enc setBuffer:g_tp_fold_ctl offset:(NSUInteger)fold_slot * 8u atIndex:7];
[enc setBytes:&ntg32 length:sizeof(ntg32) atIndex:8];
[enc setThreadgroupMemoryLength:16 atIndex:1];
}
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(ntg, n_rows, 1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (fold) g_tp_flag_prepublished_seq = g_tp_seq + 1u;
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 kslice matvec")) {
return 0;
}
return 1;
}
}
int ds4_gpu_matmul_q8_0_kslice_tensor(
ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint64_t full_in_dim, uint64_t k_off,
uint64_t k_cnt, uint64_t out_dim, const ds4_gpu_tensor *x,
uint64_t x_elem_off) {
return ds4_gpu_matmul_q8_0_kslice_exact_rows(out, model_map, model_size,
weight_offset, full_in_dim, k_off, k_cnt, out_dim, x, x_elem_off, 1);
}
static int ds4_gpu_matmul_q8_0_kslice_rows_impl(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t full_in_dim,
uint64_t out_dim,
uint64_t k_off,
uint64_t k_cnt,
const ds4_gpu_tensor *x,
uint64_t n_rows,
uint64_t x_stride,
uint64_t x_off) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!ds4_gpu_mpp_available() ||
!out || !x || !model_map ||
(full_in_dim % 64u) != 0 || (out_dim % 64u) != 0 ||
(k_off % 32u) != 0 || (k_cnt % 64u) != 0 ||
k_cnt == 0 || k_off > full_in_dim || k_cnt > full_in_dim - k_off ||
x_stride > UINT32_MAX || x_off > x_stride || k_cnt > x_stride - x_off ||
(n_rows % 32u) != 0 || n_rows == 0 ||
full_in_dim > UINT32_MAX || out_dim > UINT32_MAX ||
n_rows > UINT32_MAX) {
return 0;
}
@autoreleasepool {
const uint64_t row_bytes = (full_in_dim / 32u) * 34u;
const uint64_t weight_bytes = out_dim * row_bytes;
const uint64_t x_bytes = n_rows * x_stride * sizeof(float);
const uint64_t out_bytes = n_rows * out_dim * sizeof(float);
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
weight_offset > model_size ||
weight_bytes > model_size - weight_offset) {
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&inner_offset);
if (!wbuf) return 0;
uint64_t tile_rows = 32u;
if ((n_rows % 128u) == 0) tile_rows = 128u;
else if ((n_rows % 64u) == 0) tile_rows = 64u;
const char *fn = tile_rows == 128u
? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n128"
: (tile_rows == 64u
? "kernel_mul_mm_q8_0_f32_nax_direct_rhs_n64"
: "kernel_mul_mm_q8_0_f32_nax_direct_rhs");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mm_pipeline(fn, false, false);
if (!pipeline) return 0;
ds4_gpu_mul_mm_args args =
ds4_gpu_make_mm_args(k_cnt, out_dim, n_rows, row_bytes);
args.nb01 = row_bytes;
args.nb02 = row_bytes * out_dim;
args.nb03 = row_bytes * out_dim;
args.nb11 = x_stride * sizeof(float);
args.nb12 = n_rows * x_stride * sizeof(float);
args.nb13 = args.nb12;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf
offset:(NSUInteger)(inner_offset + (k_off / 32u) * 34u)
atIndex:1];
[enc setBuffer:xbuf
offset:(NSUInteger)(ds4_gpu_tensor_offset(x) +
x_off * sizeof(float))
atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:2u * 64u * 32u * sizeof(uint16_t)
atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_rows / tile_rows),
(NSUInteger)(out_dim / 64u),
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb,
owned,
"Q8_0 NAX kslice rows matmul");
}
}
int ds4_gpu_matmul_q8_0_kslice_rows_tensor(
ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint64_t full_in_dim, uint64_t out_dim,
uint64_t k_off, uint64_t k_cnt, const ds4_gpu_tensor *x,
uint64_t n_rows) {
return ds4_gpu_matmul_q8_0_kslice_rows_impl(out, model_map, model_size,
weight_offset, full_in_dim, out_dim, k_off, k_cnt, x, n_rows,
full_in_dim, k_off);
}
int ds4_gpu_matmul_quant_kslice_tensor(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint64_t full_in_dim,
uint64_t k_off,
uint64_t k_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *x,
uint64_t x_elem_off) {
if (weight_type == DS4_METAL_TENSOR_Q8_0) {
return ds4_gpu_matmul_q8_0_kslice_tensor(out,
model_map,
model_size,
weight_offset,
full_in_dim,
k_off,
k_cnt,
out_dim,
x,
x_elem_off);
}
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !x || !model_map ||
full_in_dim == 0 || k_cnt == 0 || out_dim == 0 ||
k_off + k_cnt > full_in_dim ||
full_in_dim > UINT32_MAX || k_cnt > UINT32_MAX ||
out_dim > UINT32_MAX) {
return 0;
}
uint64_t block_elems = 0;
uint64_t block_bytes = 0;
if (weight_type == DS4_METAL_TENSOR_Q4_K) {
block_elems = 256u;
block_bytes = 144u;
} else if (weight_type == DS4_METAL_TENSOR_Q4_0) {
block_elems = 32u;
block_bytes = 18u;
} else {
fprintf(stderr, "ds4: Metal quant kslice received unsupported type %u\n", weight_type);
return 0;
}
if ((full_in_dim % block_elems) != 0 ||
(k_off % block_elems) != 0 ||
(k_cnt % block_elems) != 0) {
fprintf(stderr, "ds4: Metal quant kslice dimensions are not block aligned\n");
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < (x_elem_off + k_cnt) * sizeof(float) ||
ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) {
fprintf(stderr, "ds4: Metal quant kslice matmul received undersized buffers\n");
return 0;
}
const uint64_t row_bytes = (full_in_dim / block_elems) * block_bytes;
const uint64_t weight_bytes = out_dim * row_bytes;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal quant kslice weights are outside the mapped model\n");
return 0;
}
uint64_t inner = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
weight_offset, weight_bytes, &inner);
if (!wbuf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (weight_type == DS4_METAL_TENSOR_Q4_K &&
(k_cnt % 256u) == 0 &&
getenv("DS4_METAL_DISABLE_Q4_MV_CLASSIC") == NULL) {
const int16_t nsg = 2;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_ext_pipeline("kernel_mul_mv_q4_K_dense_f32", nsg, 8);
if (pipeline) {
ds4_gpu_q8_0_matvec_args args = {
.ne00 = (int32_t)k_cnt,
.ne01 = (int32_t)out_dim,
.ne02 = 1,
.nb00 = 1,
.nb01 = row_bytes,
.nb02 = row_bytes * out_dim,
.nb03 = row_bytes * out_dim,
.ne10 = (int32_t)k_cnt,
.ne11 = 1,
.ne12 = 1,
.nb10 = sizeof(float),
.nb11 = k_cnt * sizeof(float),
.nb12 = k_cnt * sizeof(float),
.nb13 = k_cnt * sizeof(float),
.ne0 = (int32_t)out_dim,
.ne1 = 1,
.nr0 = 2,
.r2 = 1,
.r3 = 1,
};
const uint64_t rows_ptg = (uint64_t)nsg * 2u;
const uint64_t w_skip = (k_off / block_elems) * block_bytes;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)(inner + w_skip) atIndex:1];
[enc setBuffer:xbuf
offset:(NSUInteger)(ds4_gpu_tensor_offset(x) + x_elem_off * sizeof(float))
atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:32 atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + rows_ptg - 1u) / rows_ptg,
1,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q4_K kslice matvec")) return 0;
return 1;
}
}
fprintf(stderr, "ds4: Metal quant kslice has no kernel for type %u\n",
weight_type);
return 0;
}
}
int ds4_gpu_attention_output_q8_tp_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t out_b_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups_total,
uint32_t group0,
uint32_t group_cnt,
uint64_t out_dim,
const ds4_gpu_tensor *heads) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !low || !heads || !model_map ||
group_dim == 0 || rank == 0 || group_cnt == 0 ||
group0 + group_cnt > n_groups_total ||
(group_dim % 32u) != 0 || ((rank * group_cnt) % 32u) != 0 ||
group_dim > UINT32_MAX || rank > UINT32_MAX || out_dim > UINT32_MAX) {
return 0;
}
@autoreleasepool {
const uint64_t low_dim_total = (uint64_t)n_groups_total * rank;
const uint64_t row_a_bytes = (group_dim / 32u) * 34u;
const uint64_t a_group_bytes = rank * row_a_bytes;
const uint64_t out_a_bytes = (uint64_t)n_groups_total * a_group_bytes;
if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) {
fprintf(stderr, "ds4: Metal TP attention output weights are outside the mapped model\n");
return 0;
}
/* The heads buffer holds only the owned groups, compact at its
* base (the head slice keeps q/attention output halves packed). */
if (ds4_gpu_tensor_bytes(heads) < (uint64_t)group_cnt * group_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(low) < (uint64_t)group_cnt * rank * sizeof(float) ||
ds4_gpu_tensor_bytes(out) < out_dim * sizeof(float)) {
fprintf(stderr, "ds4: Metal TP attention output received undersized buffers\n");
return 0;
}
uint64_t out_a_inner = 0;
id<MTLBuffer> out_a_buf =
ds4_gpu_wrap_model_range(model_map, model_size,
out_a_offset, out_a_bytes, &out_a_inner);
if (!out_a_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) return 0;
/* Low projection for the owned groups only: identical dispatch to
* the single-node direct path with the weight base, heads input and
* group count shifted to the slice. The owned low half lands
* compactly at low[0 .. group_cnt*rank). */
ds4_gpu_mul_mv_id_args args = {
.nei0 = (int32_t)group_cnt,
.nei1 = 1,
.nbi1 = 0,
.ne00 = (int32_t)group_dim,
.ne01 = (int32_t)rank,
.ne02 = (int32_t)group_cnt,
.nb00 = 34,
.nb01 = row_a_bytes,
.nb02 = a_group_bytes,
.ne10 = (int32_t)group_dim,
.ne11 = (int32_t)group_cnt,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = (uint64_t)group_dim * sizeof(float),
.nb12 = (uint64_t)group_cnt * group_dim * sizeof(float),
.ne0 = (int32_t)rank,
.ne1 = (int32_t)group_cnt,
.nb1 = (uint64_t)rank * sizeof(float),
.nr0 = 2,
};
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4);
int ok = ds4_gpu_encode_attn_out_low_q8_direct(cb,
pipeline,
&args,
out_a_buf,
(NSUInteger)(out_a_inner + (uint64_t)group0 * a_group_bytes),
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low),
32u * 2u * sizeof(float),
4,
true);
if (!ok) return 0;
/* Expand projection over the owned k window only; the result is this
* rank's partial attention block output. */
return ds4_gpu_matmul_q8_0_kslice_tensor(out, model_map, model_size,
out_b_offset,
low_dim_total,
(uint64_t)group0 * rank,
(uint64_t)group_cnt * rank,
out_dim, low, 0);
}
}
int ds4_gpu_attention_output_low_q8_tensor(
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups,
const ds4_gpu_tensor *heads) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!low || !heads || !model_map || group_dim == 0 || rank == 0 ||
n_groups == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) {
return 0;
}
@autoreleasepool {
const uint64_t low_dim = (uint64_t)n_groups * rank;
if ((group_dim % 32u) != 0 || low_dim > UINT32_MAX) {
fprintf(stderr, "ds4: Metal attention output low received invalid q8 dimensions\n");
return 0;
}
const uint64_t row_a_bytes = (group_dim / 32u) * 34u;
const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes;
if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset) {
fprintf(stderr, "ds4: Metal attention output low weights are outside the mapped model\n");
return 0;
}
const uint64_t heads_bytes = (uint64_t)n_groups * group_dim * sizeof(float);
const uint64_t low_bytes = low_dim * sizeof(float);
if (ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(low) < low_bytes) {
fprintf(stderr, "ds4: Metal attention output low received undersized buffers\n");
return 0;
}
uint64_t out_a_inner = 0;
id<MTLBuffer> out_a_buf =
ds4_gpu_wrap_model_range(model_map, model_size,
out_a_offset, out_a_bytes,
&out_a_inner);
if (!out_a_buf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
bool ok = true;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) {
ok = false;
}
if (ok) {
ds4_gpu_mul_mv_id_args args = {
.nei0 = (int32_t)n_groups,
.nei1 = 1,
.nbi1 = 0,
.ne00 = (int32_t)group_dim,
.ne01 = (int32_t)rank,
.ne02 = (int32_t)n_groups,
.nb00 = 34,
.nb01 = row_a_bytes,
.nb02 = (uint64_t)rank * row_a_bytes,
.ne10 = (int32_t)group_dim,
.ne11 = (int32_t)n_groups,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = (uint64_t)group_dim * sizeof(float),
.nb12 = (uint64_t)n_groups * group_dim * sizeof(float),
.ne0 = (int32_t)rank,
.ne1 = (int32_t)n_groups,
.nb1 = (uint64_t)rank * sizeof(float),
.nr0 = 2,
};
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32", 4);
ok = ds4_gpu_encode_attn_out_low_q8_direct(cb,
pipeline,
&args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low),
32u * 2u * sizeof(float),
4,
true) != 0;
}
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
return ok ? 1 : 0;
}
}
int ds4_gpu_attention_output_low_q8_rows_exact_tensor(
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t n_groups_total,
uint32_t group0,
uint32_t group_cnt,
const ds4_gpu_tensor *heads,
uint32_t n_rows) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!low || !heads || !model_map || group_dim == 0 || rank == 0 ||
n_groups_total == 0 || group_cnt == 0 || group0 > n_groups_total ||
group_cnt > n_groups_total - group0 || n_rows == 0 ||
(uint64_t)n_rows * group_cnt > 65535u || (group_dim % 32u) != 0 ||
group_dim > UINT32_MAX || rank > UINT32_MAX) {
return 0;
}
@autoreleasepool {
const uint64_t row_a_bytes = (group_dim / 32u) * 34u;
const uint64_t a_offset = out_a_offset + (uint64_t)group0 * rank * row_a_bytes;
const uint64_t out_a_bytes = (uint64_t)group_cnt * rank * row_a_bytes;
if (a_offset > model_size || out_a_bytes > model_size - a_offset ||
ds4_gpu_tensor_bytes(heads) <
(uint64_t)n_rows * n_groups_total * group_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(low) < (uint64_t)n_rows * group_cnt * rank * sizeof(float)) {
return 0;
}
uint64_t out_a_inner = 0;
id<MTLBuffer> out_a_buf = ds4_gpu_wrap_model_range(model_map, model_size,
a_offset, out_a_bytes, &out_a_inner);
if (!out_a_buf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
bool ok = cb && !owned;
if (ok) {
/* the rows kernels take every token of one group per threadgroup */
const bool rows_kernel = n_rows >= 2u && n_rows <= 8u;
ds4_gpu_mul_mv_id_args args = {
.nei0 = (int32_t)group_cnt,
.nei1 = rows_kernel ? 1 : (int32_t)n_rows,
.nbi1 = 0,
.ne00 = (int32_t)group_dim,
.ne01 = (int32_t)rank,
.ne02 = (int32_t)group_cnt,
.nb00 = 34,
.nb01 = row_a_bytes,
.nb02 = (uint64_t)rank * row_a_bytes,
.ne10 = (int32_t)group_dim,
.ne11 = (int32_t)group_cnt,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = (uint64_t)group_dim * sizeof(float),
.nb12 = (uint64_t)n_groups_total * group_dim * sizeof(float),
.ne0 = (int32_t)rank,
.ne1 = (int32_t)group_cnt,
.nb1 = (uint64_t)rank * sizeof(float),
.nr0 = 2,
};
const NSUInteger groups = rows_kernel ? n_rows : 1u;
const int mma_off = ds4_gpu_v41_rows_mma_off();
const bool mma_kernel = !mma_off && n_rows >= 3u && n_rows <= 8u &&
(group_dim % 256u) == 0 && (rank % 8u) == 0;
if (mma_kernel) { args.nr0 = 8; args.ne12 = (int32_t)n_rows; }
id<MTLComputePipelineState> pipeline = mma_kernel ?
ds4_gpu_get_pipeline("kernel_dsv4_attn_out_low_q8_0_f32_mma_rows8") :
ds4_gpu_get_mul_mv_pipeline(rows_kernel ?
"kernel_dsv4_attn_out_low_q8_0_f32_rows" : "kernel_dsv4_attn_out_low_q8_0_f32", 4);
ok = pipeline && ds4_gpu_encode_attn_out_low_q8_direct(cb, pipeline, &args,
out_a_buf, (NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads) + (NSUInteger)group0 * group_dim * sizeof(float),
ds4_gpu_tensor_buffer(low), ds4_gpu_tensor_offset(low),
mma_kernel ? (8u * 2u * 8u * 32u + 8u * 64u) * sizeof(float) : 32u * 2u * sizeof(float) * groups,
mma_kernel ? 8u : 4u * groups, true) != 0;
}
if (!had_batch) ok = ds4_gpu_end_commands() != 0 && ok;
return ok ? 1 : 0;
}
}
int ds4_gpu_attention_output_low_q4_K_slice_tensor(
ds4_gpu_tensor *low,
const void *model_map,
uint64_t model_size,
uint64_t out_a_offset,
uint64_t group_dim,
uint64_t rank,
uint32_t group0,
uint32_t group_cnt,
const ds4_gpu_tensor *heads) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!low || !heads || !model_map || group_dim == 0 || rank == 0 ||
group_cnt == 0 || group_dim > UINT32_MAX || rank > UINT32_MAX) {
return 0;
}
@autoreleasepool {
if ((group_dim % 256u) != 0) {
fprintf(stderr, "ds4: Metal attention output low received invalid Q4_K dimensions\n");
return 0;
}
uint64_t row_a_bytes = 0;
if (!ds4_gpu_quant_row_bytes(DS4_METAL_TENSOR_Q4_K,
(uint32_t)group_dim,
&row_a_bytes)) {
return 0;
}
if (rank > UINT64_MAX / row_a_bytes) return 0;
const uint64_t group_weight_bytes = rank * row_a_bytes;
if (group0 > UINT64_MAX / group_weight_bytes ||
group_cnt > UINT64_MAX / group_weight_bytes) {
return 0;
}
const uint64_t group_skip = (uint64_t)group0 * group_weight_bytes;
const uint64_t out_a_bytes = (uint64_t)group_cnt * group_weight_bytes;
if (out_a_offset > UINT64_MAX - group_skip ||
out_a_offset + group_skip > model_size ||
out_a_bytes > model_size - (out_a_offset + group_skip)) {
fprintf(stderr, "ds4: Metal Q4 attention output low weights are outside the mapped model\n");
return 0;
}
const uint64_t heads_bytes = (uint64_t)group_cnt * group_dim * sizeof(float);
const uint64_t low_bytes = (uint64_t)group_cnt * rank * sizeof(float);
if (ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(low) < low_bytes) {
fprintf(stderr, "ds4: Metal Q4 attention output low received undersized buffers\n");
return 0;
}
uint64_t out_a_inner = 0;
id<MTLBuffer> out_a_buf =
ds4_gpu_wrap_model_range(model_map, model_size,
out_a_offset + group_skip,
out_a_bytes,
&out_a_inner);
if (!out_a_buf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
bool ok = true;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || owned) {
ok = false;
}
if (ok) {
ds4_gpu_mul_mv_id_args args = {
.nei0 = (int32_t)group_cnt,
.nei1 = 1,
.nbi1 = 0,
.ne00 = (int32_t)group_dim,
.ne01 = (int32_t)rank,
.ne02 = (int32_t)group_cnt,
.nb00 = 1,
.nb01 = row_a_bytes,
.nb02 = group_weight_bytes,
.ne10 = (int32_t)group_dim,
.ne11 = (int32_t)group_cnt,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = (uint64_t)group_dim * sizeof(float),
.nb12 = (uint64_t)group_cnt * group_dim * sizeof(float),
.ne0 = (int32_t)rank,
.ne1 = (int32_t)group_cnt,
.nb1 = (uint64_t)rank * sizeof(float),
.nr0 = 2,
};
const NSUInteger nsg = 2;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q4_K_f32", (int16_t)nsg);
ok = ds4_gpu_encode_attn_out_low_q8_direct(cb,
pipeline,
&args,
out_a_buf,
(NSUInteger)out_a_inner,
ds4_gpu_tensor_buffer(heads),
ds4_gpu_tensor_offset(heads),
ds4_gpu_tensor_buffer(low),
ds4_gpu_tensor_offset(low),
32u,
nsg,
false) != 0;
}
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
return ok ? 1 : 0;
}
}
static NSUInteger ds4_gpu_align_up_ns(NSUInteger value, NSUInteger align) {
return (value + align - 1u) & ~(align - 1u);
}
static int ds4_gpu_encode_cpy_f32_f32_3d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t cols,
uint32_t rows,
uint32_t planes,
uint64_t src_row_stride,
uint64_t src_plane_stride,
uint64_t dst_row_stride,
uint64_t dst_plane_stride) {
if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0;
ds4_gpu_cpy_args args = {
.nk0 = (int64_t)cols,
.ne00 = (int64_t)cols,
.ne01 = (int64_t)rows,
.ne02 = (int64_t)planes,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = src_row_stride,
.nb02 = src_plane_stride,
.nb03 = (uint64_t)planes * src_plane_stride,
.ne0 = (int64_t)cols,
.ne1 = (int64_t)rows,
.ne2 = (int64_t)planes,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = dst_row_stride,
.nb2 = dst_plane_stride,
.nb3 = (uint64_t)planes * dst_plane_stride,
};
const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline);
const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_cpy_f32_f32_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_cpy_f32_f32_3d_src_strided(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t cols,
uint32_t rows,
uint32_t planes,
uint64_t src_col_stride,
uint64_t src_row_stride,
uint64_t src_plane_stride,
uint64_t dst_row_stride,
uint64_t dst_plane_stride) {
if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0;
ds4_gpu_cpy_args args = {
.nk0 = (int64_t)cols,
.ne00 = (int64_t)cols,
.ne01 = (int64_t)rows,
.ne02 = (int64_t)planes,
.ne03 = 1,
.nb00 = src_col_stride,
.nb01 = src_row_stride,
.nb02 = src_plane_stride,
.nb03 = (uint64_t)planes * src_plane_stride,
.ne0 = (int64_t)cols,
.ne1 = (int64_t)rows,
.ne2 = (int64_t)planes,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = dst_row_stride,
.nb2 = dst_plane_stride,
.nb3 = (uint64_t)planes * dst_plane_stride,
};
const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f32_pipeline);
const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_cpy_f32_f32_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_cpy_f32_f16_1d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t n) {
if (!cb || !src || !dst || n == 0) return 0;
const int use_contiguous =
ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0;
id<MTLComputePipelineState> pipeline = use_contiguous
? g_cpy_contig_f32_f16_pipeline
: g_cpy_f32_f16_pipeline;
const NSUInteger work_items = use_contiguous
? ((NSUInteger)n + 3u) / 4u
: (NSUInteger)n;
const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline);
const NSUInteger groups = (work_items + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
if (use_contiguous) {
[enc setBytes:&n length:sizeof(n) atIndex:0];
} else {
ds4_gpu_cpy_args args =
ds4_gpu_make_cpy_1d_args(n, sizeof(float), sizeof(uint16_t));
[enc setBytes:&args length:sizeof(args) atIndex:0];
}
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_cpy_f32_f16_2d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t cols,
uint32_t rows,
uint64_t src_row_stride,
uint64_t dst_row_stride) {
if (!cb || !src || !dst || cols == 0 || rows == 0) return 0;
ds4_gpu_cpy_args args = {
.nk0 = (int64_t)cols,
.ne00 = (int64_t)cols,
.ne01 = (int64_t)rows,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = src_row_stride,
.nb02 = (uint64_t)rows * src_row_stride,
.nb03 = (uint64_t)rows * src_row_stride,
.ne0 = (int64_t)cols,
.ne1 = (int64_t)rows,
.ne2 = 1,
.ne3 = 1,
.nb0 = sizeof(uint16_t),
.nb1 = dst_row_stride,
.nb2 = (uint64_t)rows * dst_row_stride,
.nb3 = (uint64_t)rows * dst_row_stride,
};
const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline);
const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_cpy_f32_f16_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_cpy_f32_f16_3d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t cols,
uint32_t rows,
uint32_t planes,
uint64_t src_row_stride,
uint64_t src_plane_stride,
uint64_t dst_row_stride,
uint64_t dst_plane_stride) {
if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0;
ds4_gpu_cpy_args args = {
.nk0 = (int64_t)cols,
.ne00 = (int64_t)cols,
.ne01 = (int64_t)rows,
.ne02 = (int64_t)planes,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = src_row_stride,
.nb02 = src_plane_stride,
.nb03 = (uint64_t)planes * src_plane_stride,
.ne0 = (int64_t)cols,
.ne1 = (int64_t)rows,
.ne2 = (int64_t)planes,
.ne3 = 1,
.nb0 = sizeof(uint16_t),
.nb1 = dst_row_stride,
.nb2 = dst_plane_stride,
.nb3 = (uint64_t)planes * dst_plane_stride,
};
const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f32_f16_pipeline);
const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_cpy_f32_f16_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_cpy_f16_f16_3d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t cols,
uint32_t rows,
uint32_t planes,
uint64_t src_row_stride,
uint64_t src_plane_stride,
uint64_t dst_row_stride,
uint64_t dst_plane_stride) {
if (!cb || !src || !dst || cols == 0 || rows == 0 || planes == 0) return 0;
ds4_gpu_cpy_args args = {
.nk0 = (int64_t)cols,
.ne00 = (int64_t)cols,
.ne01 = (int64_t)rows,
.ne02 = (int64_t)planes,
.ne03 = 1,
.nb00 = sizeof(uint16_t),
.nb01 = src_row_stride,
.nb02 = src_plane_stride,
.nb03 = (uint64_t)planes * src_plane_stride,
.ne0 = (int64_t)cols,
.ne1 = (int64_t)rows,
.ne2 = (int64_t)planes,
.ne3 = 1,
.nb0 = sizeof(uint16_t),
.nb1 = dst_row_stride,
.nb2 = dst_plane_stride,
.nb3 = (uint64_t)planes * dst_plane_stride,
};
const NSUInteger nth = ds4_gpu_cpy_threads(cols, g_cpy_f16_f16_pipeline);
const NSUInteger col_groups = ((NSUInteger)cols + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_cpy_f16_f16_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(col_groups * rows, planes, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_cpy_f16_f32_1d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t n) {
if (!cb || !src || !dst || n == 0) return 0;
const int use_contiguous =
ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") <= 0;
id<MTLComputePipelineState> pipeline = use_contiguous
? g_cpy_contig_f16_f32_pipeline
: g_cpy_f16_f32_pipeline;
const NSUInteger work_items = use_contiguous
? ((NSUInteger)n + 3u) / 4u
: (NSUInteger)n;
const NSUInteger nth = ds4_gpu_cpy_threads((uint32_t)work_items, pipeline);
const NSUInteger groups = (work_items + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
if (use_contiguous) {
[enc setBytes:&n length:sizeof(n) atIndex:0];
} else {
ds4_gpu_cpy_args args =
ds4_gpu_make_cpy_1d_args(n, sizeof(uint16_t), sizeof(float));
[enc setBytes:&args length:sizeof(args) atIndex:0];
}
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_copy_to_f16_1d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
bool src_is_f16,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t n) {
if (!cb || !src || !dst) return 0;
if (n == 0) return 1;
if (!src_is_f16) {
return ds4_gpu_encode_cpy_f32_f16_1d(cb, src, src_off, dst, dst_off, n);
}
if (ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") <= 0) {
const NSUInteger work_items = ((NSUInteger)n + 3u) / 4u;
const NSUInteger nth = ds4_gpu_cpy_threads(
(uint32_t)work_items,
g_cpy_contig_f16_f16_pipeline);
const NSUInteger groups = (work_items + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_cpy_contig_f16_f16_pipeline];
[enc setBytes:&n length:sizeof(n) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
if (g_batch_cb && cb == g_batch_cb) ds4_gpu_close_batch_encoder();
id<MTLBlitCommandEncoder> blit = ds4_gpu_blit_encoder(cb, "f16_copy", 1u);
if (!blit) return 0;
[blit copyFromBuffer:src
sourceOffset:src_off
toBuffer:dst
destinationOffset:dst_off
size:(NSUInteger)n * sizeof(uint16_t)];
[blit endEncoding];
return 1;
}
static int ds4_gpu_encode_copy_raw_ring_to_f16(
id<MTLCommandBuffer> cb,
id<MTLBuffer> raw,
NSUInteger raw_offset,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_raw,
uint32_t head_dim,
id<MTLBuffer> dst,
NSUInteger dst_offset) {
if (!cb || !raw || !dst || raw_cap == 0 || raw_start >= raw_cap ||
n_raw == 0 || n_raw > raw_cap || head_dim == 0) {
return 0;
}
const NSUInteger raw_row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger dst_row_bytes = (NSUInteger)head_dim * sizeof(uint16_t);
const uint32_t tail_rows = raw_cap - raw_start < n_raw
? raw_cap - raw_start
: n_raw;
const uint32_t head_rows = n_raw - tail_rows;
const uint64_t tail_count = (uint64_t)tail_rows * head_dim;
const uint64_t head_count = (uint64_t)head_rows * head_dim;
const uint64_t raw_inner = (uint64_t)raw_start * raw_row_bytes;
const uint64_t dst_inner = (uint64_t)tail_rows * dst_row_bytes;
if (tail_count > UINT32_MAX || head_count > UINT32_MAX ||
raw_inner > NSUIntegerMax - raw_offset ||
dst_inner > NSUIntegerMax - dst_offset) {
return 0;
}
if (tail_rows &&
!ds4_gpu_encode_cpy_f32_f16_1d(
cb,
raw,
raw_offset + (NSUInteger)raw_inner,
dst,
dst_offset,
(uint32_t)tail_count)) {
return 0;
}
if (head_rows &&
!ds4_gpu_encode_cpy_f32_f16_1d(
cb,
raw,
raw_offset,
dst,
dst_offset + (NSUInteger)dst_inner,
(uint32_t)head_count)) {
return 0;
}
return 1;
}
static int ds4_gpu_encode_flash_kv_stage_f16(
id<MTLCommandBuffer> cb,
id<MTLBuffer> raw,
NSUInteger raw_offset,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_raw,
id<MTLBuffer> comp,
NSUInteger comp_offset,
bool comp_is_f16,
id<MTLBuffer> ids,
NSUInteger ids_offset,
uint32_t n_comp,
uint32_t head_dim,
id<MTLBuffer> dst,
NSUInteger dst_offset,
id<MTLBuffer> mask,
NSUInteger mask_offset,
id<MTLBuffer> pad,
NSUInteger pad_offset,
bool fuse_pad,
bool shared_pad,
bool *did_fuse_pad) {
if (did_fuse_pad) *did_fuse_pad = false;
if (!cb || !raw || !comp || !dst || raw_cap == 0 ||
raw_start >= raw_cap || n_raw == 0 || n_raw > raw_cap ||
n_comp == 0 || head_dim == 0) {
return 0;
}
const bool require =
getenv("DS4_METAL_REQUIRE_GATHERED_KV_STAGE") != NULL;
const bool supported_shape =
(comp_is_f16 || ids != nil) && head_dim == 512u && raw_cap <= UINT32_MAX / 128u;
const uint64_t row_vecs64 = 128u;
const uint64_t comp_count64 = (uint64_t)n_comp * head_dim;
const uint64_t dst_comp_inner64 =
(uint64_t)n_raw * head_dim * sizeof(uint16_t);
if (comp_count64 > UINT32_MAX ||
dst_comp_inner64 > NSUIntegerMax - dst_offset) {
return 0;
}
const uint64_t total_vecs64 =
((uint64_t)n_raw + n_comp) * row_vecs64;
const bool valid_grid =
total_vecs64 != 0 && total_vecs64 <= UINT32_MAX;
id<MTLComputePipelineState> stage_pipeline = ids ?
ds4_gpu_get_pipeline("kernel_dsv41_flash_kv_stage") : g_flash_kv_stage_f16_pipeline;
const bool eligible =
supported_shape && valid_grid && !g_quality_mode &&
getenv("DS4_METAL_DISABLE_GATHERED_KV_STAGE") == NULL &&
stage_pipeline != nil &&
(ds4_gpu_device_name_contains("M3") ||
ds4_gpu_device_name_contains("M5"));
const bool component_disabled = eligible &&
(ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F32_F16_COPY") > 0 ||
ds4_gpu_env_bool("DS4_METAL_DISABLE_CONTIG_F16_F16_COPY") > 0);
const bool use_fusion = eligible && !component_disabled;
const bool use_pad_fusion =
use_fusion && fuse_pad && mask != nil && pad != nil &&
getenv("DS4_METAL_DISABLE_GATHERED_KV_PAD_FUSION") == NULL;
if (require && supported_shape && !use_fusion) {
fprintf(stderr,
"ds4: required Metal gathered KV staging kernel was not selected\n");
return 0;
}
if (use_fusion) {
ds4_gpu_flash_kv_stage_f16_args args = {
.raw_cap = raw_cap,
.raw_start = raw_start,
.n_raw = n_raw,
.n_comp = n_comp,
.pad_rows = use_pad_fusion ? 32u : 0u,
.shared_pad = use_pad_fusion && shared_pad ? 1u : 0u,
};
const NSUInteger total_vecs = (NSUInteger)total_vecs64 +
(use_pad_fusion ? 32u * 128u + 32u : 0u);
NSUInteger nth =
stage_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > total_vecs) nth = total_vecs;
if (nth == 0) return 0;
const NSUInteger groups = (total_vecs + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
if (g_kv_task.pending) {
/* Safety net: a deferred kv task nothing consumed runs before the
* staging reads the cache row. */
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_kv_norm_task_flush()) return 0;
enc = ds4_gpu_compute_encoder(cb);
}
[enc setComputePipelineState:stage_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:raw offset:raw_offset atIndex:1];
[enc setBuffer:comp offset:comp_offset atIndex:2];
[enc setBuffer:dst offset:dst_offset atIndex:3];
[enc setBuffer:(use_pad_fusion ? mask : dst)
offset:(use_pad_fusion ? mask_offset : dst_offset)
atIndex:4];
[enc setBuffer:(use_pad_fusion ? pad : dst)
offset:(use_pad_fusion ? pad_offset : dst_offset)
atIndex:5];
if (ids) [enc setBuffer:ids offset:ids_offset atIndex:6];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (did_fuse_pad) *did_fuse_pad = use_pad_fusion;
return 1;
}
if (ids) return 0; /* gathered rows have no separate copy path */
if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb,
raw,
raw_offset,
raw_cap,
raw_start,
n_raw,
head_dim,
dst,
dst_offset)) {
return 0;
}
return ds4_gpu_encode_copy_to_f16_1d(
cb,
comp,
comp_offset,
comp_is_f16,
dst,
dst_offset + (NSUInteger)dst_comp_inner64,
(uint32_t)comp_count64);
}
int ds4_gpu_flash_kv_stage_f16_tensor(
ds4_gpu_tensor *dst,
const ds4_gpu_tensor *raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_raw,
const ds4_gpu_tensor *comp,
uint32_t comp_is_f16,
uint32_t n_comp,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!dst || !raw || !comp || raw_cap == 0 || raw_start >= raw_cap ||
n_raw == 0 || n_raw > raw_cap || n_comp == 0 ||
comp_is_f16 == 0 || head_dim != 512u) {
return 0;
}
@autoreleasepool {
const uint64_t raw_bytes =
(uint64_t)raw_cap * head_dim * sizeof(float);
const uint64_t comp_bytes =
(uint64_t)n_comp * head_dim *
(comp_is_f16 ? sizeof(uint16_t) : sizeof(float));
const uint64_t dst_bytes =
((uint64_t)n_raw + n_comp) * head_dim * sizeof(uint16_t);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw);
id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp);
id<MTLBuffer> dstbuf = ds4_gpu_tensor_buffer(dst);
if (!rawbuf || !compbuf || !dstbuf ||
ds4_gpu_tensor_bytes(raw) < raw_bytes ||
ds4_gpu_tensor_bytes(comp) < comp_bytes ||
ds4_gpu_tensor_bytes(dst) < dst_bytes) {
fprintf(stderr,
"ds4: Metal gathered KV staging received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_kv_stage_f16(
cb,
rawbuf,
ds4_gpu_tensor_offset(raw),
raw_cap,
raw_start,
n_raw,
compbuf,
ds4_gpu_tensor_offset(comp),
comp_is_f16 != 0,
nil,
0,
n_comp,
head_dim,
dstbuf,
ds4_gpu_tensor_offset(dst),
nil,
0,
nil,
0,
false,
false,
NULL)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(
cb, owned, "gathered KV staging")) {
return 0;
}
}
return 1;
}
static int ds4_gpu_encode_fill_f16_1d(
id<MTLCommandBuffer> cb,
id<MTLBuffer> buf,
NSUInteger offset,
uint32_t n,
float value) {
if (!cb || !buf || n == 0) return 0;
ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(n, 1, 0, 0.0f, 0.0f);
args.val = value;
NSUInteger nth = (NSUInteger)n;
const NSUInteger max_threads = g_unary_fill_f16_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth > 256u) nth = 256u;
if (nth == 0) nth = 1u;
const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_unary_fill_f16_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:buf offset:offset atIndex:1];
[enc setBuffer:buf offset:offset atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_flash_attention_raw_heads(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_head,
uint32_t head_dim) {
if (head_dim != 512 || n_head == 0 || n_raw == 0 || raw_cap < n_raw) {
return 0;
}
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
const uint64_t heads_bytes = q_bytes;
if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
ds4_gpu_tensor_bytes(heads) < heads_bytes) {
fprintf(stderr, "ds4: Metal DS4 FlashAttention received undersized buffers\n");
return 0;
}
const uint32_t ncpsg = 32;
const uint32_t nwg = 32;
const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_raw, nwg, ncpsg);
const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_raw * sizeof(uint16_t);
const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16;
const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 +
(NSUInteger)ncpsg * sizeof(uint16_t);
const NSUInteger nrows = (NSUInteger)n_head;
const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) +
nrows * (2u * (NSUInteger)nwg) * sizeof(float);
const bool use_persistent_zero_mask =
ds4_gpu_device_is_pre_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_PRE_M5_DECODE_RAW_ZERO_ATTN_MASK") == NULL &&
!g_quality_mode &&
!g_ssd_streaming_mode &&
g_tp_split_world == 1;
id<MTLBuffer> mask_buffer = nil;
if (use_persistent_zero_mask &&
ds4_gpu_ensure_zero_attention_mask(mask_bytes)) {
mask_buffer = g_flash_attn_zero_mask_buffer;
} else {
mask_buffer =
ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
if (mask_buffer) memset([mask_buffer contents], 0, mask_bytes);
}
if (!mask_buffer ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_bytes,
"ds4_flash_attn_kv_f16") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_flash_attn_pad") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer,
&g_flash_attn_tmp_bytes,
tmp_bytes,
"ds4_flash_attn_tmp")) {
return 0;
}
id<MTLComputePipelineState> pad_pipeline = nil;
if ((n_raw % ncpsg) != 0) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
id<MTLComputePipelineState> vec_pipeline =
ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512",
true, true, false, false, (n_raw % ncpsg) != 0,
false,
(int32_t)head_dim,
(int32_t)head_dim,
(int32_t)nsg,
(int32_t)nwg);
id<MTLComputePipelineState> reduce_pipeline =
ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg);
if (!vec_pipeline || !reduce_pipeline) return 0;
if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb,
rawbuf,
ds4_gpu_tensor_offset(raw_kv),
raw_cap,
raw_start,
n_raw,
head_dim,
g_flash_attn_kv_buffer,
0)) {
return 0;
}
if ((n_raw % ncpsg) != 0) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)n_raw,
.ne_12_2 = 1,
.ne_12_3 = 1,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_raw * row_bytes_f16,
.nb13 = (uint64_t)n_raw * row_bytes_f16,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_raw * row_bytes_f16,
.nb23 = (uint64_t)n_raw * row_bytes_f16,
.ne31 = 1,
.ne32 = 1,
.ne33 = 1,
.nb31 = mask_bytes,
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
ds4_gpu_flash_attn_vec_args vec_args = {
.ne01 = 1,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * row_bytes,
.nb02 = row_bytes,
.nb03 = (uint64_t)n_head * row_bytes,
.ne11 = (int32_t)n_raw,
.ne_12_2 = 1,
.ne_12_3 = 1,
.ns10 = (int32_t)head_dim,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_raw * row_bytes_f16,
.nb13 = (uint64_t)n_raw * row_bytes_f16,
.ns20 = (int32_t)head_dim,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_raw * row_bytes_f16,
.nb23 = (uint64_t)n_raw * row_bytes_f16,
.ne31 = 1,
.ne32 = 1,
.ne33 = 1,
.nb31 = mask_bytes,
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = 1,
.ne3 = 1,
.scale = 1.0f / sqrtf((float)head_dim),
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) +
4u * ncpsg +
2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg;
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:vec_pipeline];
[enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
[enc setBuffer:mask_buffer offset:0 atIndex:4];
[enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
ds4_gpu_flash_attn_reduce_args reduce_args = {
.nrows = (int32_t)nrows,
};
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:reduce_pipeline];
[enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
/* Rectangular causal/window mask for raw prefill: q covers rows
* [q_row0, q_row0 + n_q) of an n_kv-token block whose keys all live in
* rows [0, n_kv). The square prefill case is q_row0 == 0, n_q == n_kv. */
static void ds4_gpu_fill_raw_prefill_mask(
uint16_t *mask,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_kv,
uint32_t window) {
const uint16_t neg_inf_half = 0xfc00u;
for (uint32_t q = 0; q < n_q; q++) {
const uint32_t qpos = q_row0 + q;
uint16_t *row = mask + (uint64_t)q * n_kv;
for (uint32_t k = 0; k < n_kv; k++) {
const bool causal = k <= qpos;
const bool in_window = window == 0 || qpos - k < window;
row[k] = causal && in_window ? 0u : neg_inf_half;
}
}
}
static void ds4_gpu_fill_glm_prefill_mask(
uint16_t *mask,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len) {
const uint16_t neg_inf_half = 0xfc00u;
for (uint32_t q = 0; q < n_tokens; q++) {
const uint32_t qpos = pos0 + q;
uint16_t *row = mask + (uint64_t)q * cache_len;
for (uint32_t k = 0; k < cache_len; k++) {
row[k] = k <= qpos ? 0u : neg_inf_half;
}
}
}
static id<MTLBuffer> ds4_gpu_glm_prefill_mask_buffer(
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
NSUInteger mask_bytes) {
const int same_shape =
g_glm_flash_attn_mask_valid &&
g_glm_flash_attn_mask_buffer &&
g_glm_flash_attn_mask_bytes >= mask_bytes &&
g_glm_flash_attn_mask_pos0 == pos0 &&
g_glm_flash_attn_mask_tokens == n_tokens &&
g_glm_flash_attn_mask_cache_len == cache_len;
if (same_shape) return g_glm_flash_attn_mask_buffer;
if (g_glm_flash_attn_mask_buffer) {
[g_transient_buffers addObject:g_glm_flash_attn_mask_buffer];
g_glm_flash_attn_mask_buffer = nil;
}
g_glm_flash_attn_mask_bytes = 0;
g_glm_flash_attn_mask_valid = 0;
if (!ds4_gpu_ensure_scratch_buffer(&g_glm_flash_attn_mask_buffer,
&g_glm_flash_attn_mask_bytes,
mask_bytes,
"ds4_glm_flash_attn_mask")) {
return nil;
}
ds4_gpu_fill_glm_prefill_mask((uint16_t *)[g_glm_flash_attn_mask_buffer contents],
pos0,
n_tokens,
cache_len);
g_glm_flash_attn_mask_pos0 = pos0;
g_glm_flash_attn_mask_tokens = n_tokens;
g_glm_flash_attn_mask_cache_len = cache_len;
g_glm_flash_attn_mask_valid = 1;
return g_glm_flash_attn_mask_buffer;
}
static void ds4_gpu_fill_raw_decode_batch_mask(
uint16_t *mask,
uint32_t n_tokens,
uint32_t n_raw,
uint32_t pos0,
uint32_t window) {
const uint16_t neg_inf_half = 0xfc00u;
const uint32_t last_pos = pos0 + n_tokens - 1u;
/* The caller has already copied the SWA ring into logical order when it
* wraps, so key row k represents first_raw_pos + k. */
const uint32_t first_raw_pos = last_pos + 1u - n_raw;
for (uint32_t q = 0; q < n_tokens; q++) {
const uint32_t qpos = pos0 + q;
uint16_t *row = mask + (uint64_t)q * n_raw;
for (uint32_t k = 0; k < n_raw; k++) {
const uint32_t kpos = first_raw_pos + k;
const bool causal = kpos <= qpos;
const bool in_window = causal && (window == 0 || qpos - kpos < window);
row[k] = causal && in_window ? 0u : neg_inf_half;
}
}
}
static void ds4_gpu_fill_mixed_decode_batch_mask(
uint16_t *mask,
uint32_t n_tokens,
uint32_t n_raw,
uint32_t n_comp,
uint32_t pos0,
uint32_t window,
uint32_t ratio) {
const uint16_t neg_inf_half = 0xfc00u;
const uint32_t n_keys = n_raw + n_comp;
const uint32_t last_pos = pos0 + n_tokens - 1u;
/* Raw keys are laid out by logical position; compressed keys follow them. */
const uint32_t first_raw_pos = last_pos + 1u - n_raw;
for (uint32_t q = 0; q < n_tokens; q++) {
const uint32_t qpos = pos0 + q;
uint16_t *row = mask + (uint64_t)q * n_keys;
for (uint32_t k = 0; k < n_raw; k++) {
const uint32_t kpos = first_raw_pos + k;
const bool causal = kpos <= qpos;
const bool in_window = causal && (window == 0 || qpos - kpos < window);
row[k] = causal && in_window ? 0u : neg_inf_half;
}
const uint32_t n_visible = (qpos + 1u) / ratio;
for (uint32_t c = 0; c < n_comp; c++) {
row[n_raw + c] = c < n_visible ? 0u : neg_inf_half;
}
}
}
static bool ds4_gpu_fill_visual_mixed_batch_mask(
uint16_t *mask,
const int32_t *tokens,
uint32_t vocab_size,
uint32_t n_tokens,
uint32_t n_raw,
uint32_t n_comp,
uint32_t pos0,
uint32_t window,
uint32_t ratio) {
if (!mask || !tokens || vocab_size == 0 || n_tokens == 0 || n_raw == 0) {
return false;
}
const uint16_t neg_inf_half = 0xfc00u;
const uint32_t n_keys = n_raw + n_comp;
const uint32_t first_raw_pos = pos0 + n_tokens - n_raw;
uint32_t *bounds = malloc((size_t)n_tokens * 2u * sizeof(*bounds));
if (!bounds || !ds4_deepseek4_attention_bounds(
(const int *)tokens, n_tokens, vocab_size,
pos0, n_raw, window, bounds)) {
free(bounds);
return false;
}
for (uint32_t q = 0; q < n_tokens; q++) {
const uint32_t qpos = pos0 + q;
uint16_t *row = mask + (uint64_t)q * n_keys;
for (uint32_t k = 0; k < n_raw; k++) {
const uint32_t kpos = first_raw_pos + k;
row[k] = kpos >= bounds[2u * q] &&
kpos <= bounds[2u * q + 1u]
? 0u : neg_inf_half;
}
const uint32_t n_visible = ratio ? (qpos + 1u) / ratio : 0u;
for (uint32_t c = 0; c < n_comp; c++) {
row[n_raw + c] = c < n_visible ? 0u : neg_inf_half;
}
}
free(bounds);
return true;
}
/* Rectangular causal/window + compressed-key visibility mask: q covers rows
* [q_row0, q_row0 + n_q) of an n_tokens-token chunk whose raw keys all stay
* resident, followed by n_comp compressed keys. The square prefill case is
* q_row0 == 0, n_q == n_tokens. */
static void ds4_gpu_fill_static_mixed_prefill_mask(
uint16_t *mask,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio) {
const uint16_t neg_inf_half = 0xfc00u;
const uint32_t n_keys = n_tokens + n_comp;
for (uint32_t q = 0; q < n_q; q++) {
const uint32_t qpos = q_row0 + q;
uint16_t *row = mask + (uint64_t)q * n_keys;
for (uint32_t k = 0; k < n_tokens; k++) {
const bool causal = k <= qpos;
const bool in_window = window == 0 || qpos - k < window;
row[k] = causal && in_window ? 0u : neg_inf_half;
}
const uint32_t n_visible = (qpos + 1u) / ratio;
for (uint32_t c = 0; c < n_comp; c++) {
row[n_tokens + c] = c < n_visible ? 0u : neg_inf_half;
}
}
}
/* Static-mixed prefill FlashAttention over a rectangular problem: q holds
* n_q query rows for token positions [q_row0, q_row0 + n_q) of the chunk,
* while the keys stay full (all n_tokens raw rows plus n_comp compressed
* rows). The classic square prefill is q_row0 == 0, n_q == n_tokens. */
static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long(
id<MTLCommandBuffer> __strong *cbp,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!cbp || !*cbp) return 0;
id<MTLCommandBuffer> cb = *cbp;
if (head_dim != 512 || n_head == 0 || n_q == 0 || n_tokens == 0 || ratio == 0) {
return 0;
}
const uint32_t n_keys = n_tokens + n_comp;
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf;
id<MTLBuffer> maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf;
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
(comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
const uint64_t comp_mask_bytes = use_comp_mask
? (uint64_t)n_comp * (q_row0 + n_q) * sizeof(float) : 0u;
if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
(n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) ||
(use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) ||
ds4_gpu_tensor_bytes(heads) < q_bytes) {
fprintf(stderr, "ds4: Metal prefill static mixed DS4 non-vector FlashAttention received undersized buffers\n");
return 0;
}
const uint32_t nqptg = 8;
const uint32_t ncpsg = 64;
const uint32_t nsg = head_dim >= 512 ? 8u : 4u;
const bool has_kvpad = (n_keys % ncpsg) != 0;
const bool bc_mask = (n_q % nqptg) != 0;
const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_q * sizeof(uint16_t);
const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16;
const NSUInteger pad_bytes = has_kvpad
? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t))
: 1u;
const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg;
const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg;
const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);
const uint32_t mask_cache_kind = ratio == 4u
? DS4_GPU_PREFILL_MASK_CACHE_RATIO4
: (ratio == 128u ? DS4_GPU_PREFILL_MASK_CACHE_RATIO128 : 0u);
bool mask_cache_created = false;
ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache =
use_comp_mask == 0u && mask_cache_kind != 0u &&
q_row0 == 0u && n_q == n_tokens
? ds4_gpu_get_zero_prefix_prefill_mask_cache(mask_cache_kind,
n_tokens,
n_comp,
n_keys,
window,
ratio,
nqptg,
ncpsg,
has_kvpad,
bc_mask,
mask_bytes,
blk_bytes,
&mask_cache_created)
: NULL;
id<MTLBuffer> mask_buffer = mask_cache
? mask_cache->mask
: ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
if (mask_cache && mask_cache_created) {
ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents],
0u,
n_tokens,
n_tokens,
n_comp,
window,
ratio);
mask_cache->valid = true;
}
if (!mask_buffer ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_bytes,
"ds4_flash_attn_kv_f16") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_flash_attn_pad") ||
(!mask_cache &&
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
&g_flash_attn_blk_bytes,
blk_bytes,
"ds4_flash_attn_blk"))) {
return 0;
}
id<MTLBuffer> blk_buffer = mask_cache
? mask_cache->blk
: g_flash_attn_blk_buffer;
const bool flash_stage_profile =
getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil;
double flash_stage_t0 = 0.0;
if (flash_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
int profile_owned = 0;
cb = ds4_gpu_command_buffer(&profile_owned);
if (!cb || profile_owned) return 0;
*cbp = cb;
flash_stage_t0 = ds4_gpu_now_ms();
}
#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \
if (flash_stage_profile) { \
if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \
"static_mixed_nonvec", (name), n_q, n_comp, n_keys, \
n_head, head_dim, window, ratio, &flash_stage_t0)) { \
return 0; \
} \
cb = *cbp; \
} \
} while (0)
if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
rawbuf,
ds4_gpu_tensor_offset(raw_kv),
g_flash_attn_kv_buffer,
0,
n_tokens * head_dim)) {
return 0;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw");
if (n_comp &&
!ds4_gpu_encode_copy_to_f16_1d(cb,
compbuf,
ds4_gpu_tensor_offset(comp_kv),
comp_kv_f16 != 0,
g_flash_attn_kv_buffer,
(NSUInteger)n_tokens * row_bytes_f16,
n_comp * head_dim)) {
return 0;
}
if (n_comp) {
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp");
}
if (!mask_cache) {
ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents],
q_row0,
n_q,
n_tokens,
n_comp,
window,
ratio);
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill");
if (use_comp_mask && n_comp != 0) {
if (!ds4_gpu_encode_cpy_f32_f16_2d(cb,
maskbuf,
ds4_gpu_tensor_offset(comp_mask) +
(NSUInteger)((uint64_t)q_row0 * n_comp * sizeof(float)),
mask_buffer,
(NSUInteger)n_tokens * sizeof(uint16_t),
n_comp,
n_q,
(uint64_t)n_comp * sizeof(float),
(uint64_t)n_keys * sizeof(uint16_t))) {
return 0;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy");
}
id<MTLComputePipelineState> pad_pipeline = nil;
if (has_kvpad) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
id<MTLComputePipelineState> blk_pipeline =
ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
id<MTLComputePipelineState> attn_pipeline =
ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512",
true, true, false, false, has_kvpad, bc_mask,
(int32_t)head_dim,
(int32_t)head_dim,
(int32_t)nsg);
if (!blk_pipeline || !attn_pipeline) return 0;
if (has_kvpad) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)n_keys,
.ne_12_2 = 1,
.ne_12_3 = 1,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_keys * row_bytes_f16,
.nb13 = (uint64_t)n_keys * row_bytes_f16,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_keys * row_bytes_f16,
.nb23 = (uint64_t)n_keys * row_bytes_f16,
.ne31 = (int32_t)n_q,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_keys * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
}
ds4_gpu_flash_attn_blk_args blk_args = {
.ne01 = (int32_t)n_q,
.ne30 = (int32_t)n_keys,
.ne31 = (int32_t)n_q,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_keys * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = nil;
if (!mask_cache || !mask_cache->blk_ready) {
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:blk_pipeline];
[enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
[enc setBuffer:mask_buffer offset:0 atIndex:1];
[enc setBuffer:blk_buffer offset:0 atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (mask_cache) mask_cache->blk_ready = true;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map");
ds4_gpu_flash_attn_vec_args args = {
.ne01 = (int32_t)n_q,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * row_bytes,
.nb02 = row_bytes,
.nb03 = (uint64_t)n_q * n_head * row_bytes,
.ne11 = (int32_t)n_keys,
.ne_12_2 = 1,
.ne_12_3 = 1,
.ns10 = (int32_t)head_dim,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_keys * row_bytes_f16,
.nb13 = (uint64_t)n_keys * row_bytes_f16,
.ns20 = (int32_t)head_dim,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_keys * row_bytes_f16,
.nb23 = (uint64_t)n_keys * row_bytes_f16,
.ne31 = (int32_t)n_q,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_keys * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = (int32_t)n_q,
.ne3 = 1,
.scale = 1.0f / sqrtf((float)head_dim),
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u);
const NSUInteger shared_elems = (NSUInteger)nqptg *
((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:attn_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
[enc setBuffer:mask_buffer offset:0 atIndex:4];
[enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:blk_buffer offset:0 atIndex:7];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention");
#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE
return 1;
}
static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec(
id<MTLCommandBuffer> __strong *cbp,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!cbp || !*cbp) return 0;
id<MTLCommandBuffer> cb = *cbp;
if (head_dim != 512 || n_head == 0 || n_tokens == 0 || ratio == 0) {
return 0;
}
const uint32_t n_keys = n_tokens + n_comp;
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf;
id<MTLBuffer> maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf;
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
(comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u;
if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
(n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) ||
(use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) ||
ds4_gpu_tensor_bytes(heads) < q_bytes) {
fprintf(stderr, "ds4: Metal prefill static mixed DS4 FlashAttention received undersized buffers\n");
return 0;
}
const uint32_t ncpsg = 32;
const uint32_t nwg = 32;
const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg);
const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t);
const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16;
const bool has_kvpad = (n_keys % ncpsg) != 0;
const NSUInteger pad_bytes = has_kvpad
? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t))
: 1u;
const NSUInteger nrows = (NSUInteger)n_tokens * n_head;
const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) +
nrows * (2u * (NSUInteger)nwg) * sizeof(float);
id<MTLBuffer> mask_buffer =
ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
if (!mask_buffer ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_bytes,
"ds4_flash_attn_kv") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_flash_attn_pad") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer,
&g_flash_attn_tmp_bytes,
tmp_bytes,
"ds4_flash_attn_tmp")) {
return 0;
}
const bool flash_stage_profile =
getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil;
double flash_stage_t0 = 0.0;
if (flash_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
int profile_owned = 0;
cb = ds4_gpu_command_buffer(&profile_owned);
if (!cb || profile_owned) return 0;
*cbp = cb;
flash_stage_t0 = ds4_gpu_now_ms();
}
#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \
if (flash_stage_profile) { \
if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \
"static_mixed_vec", (name), n_tokens, n_comp, n_keys, \
n_head, head_dim, window, ratio, &flash_stage_t0)) { \
return 0; \
} \
cb = *cbp; \
} \
} while (0)
if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
rawbuf,
ds4_gpu_tensor_offset(raw_kv),
g_flash_attn_kv_buffer,
0,
n_tokens * head_dim)) {
return 0;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw");
if (n_comp) {
if (!ds4_gpu_encode_copy_to_f16_1d(cb,
compbuf,
ds4_gpu_tensor_offset(comp_kv),
comp_kv_f16 != 0,
g_flash_attn_kv_buffer,
(NSUInteger)n_tokens * row_bytes_f16,
n_comp * head_dim)) {
return 0;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_comp");
}
ds4_gpu_fill_static_mixed_prefill_mask((uint16_t *)[mask_buffer contents],
0,
n_tokens,
n_tokens,
n_comp,
window,
ratio);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill");
if (use_comp_mask && n_comp != 0) {
if (!ds4_gpu_encode_cpy_f32_f16_2d(cb,
maskbuf,
ds4_gpu_tensor_offset(comp_mask),
mask_buffer,
(NSUInteger)n_tokens * sizeof(uint16_t),
n_comp,
n_tokens,
(uint64_t)n_comp * sizeof(float),
(uint64_t)n_keys * sizeof(uint16_t))) {
return 0;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_comp_copy");
}
id<MTLComputePipelineState> pad_pipeline = nil;
id<MTLComputeCommandEncoder> enc = nil;
if (has_kvpad) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
id<MTLComputePipelineState> vec_pipeline =
ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512",
true, true, false, false, has_kvpad,
false,
(int32_t)head_dim,
(int32_t)head_dim,
(int32_t)nsg,
(int32_t)nwg);
id<MTLComputePipelineState> reduce_pipeline =
ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg);
if (!vec_pipeline || !reduce_pipeline) return 0;
if (has_kvpad) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)n_keys,
.ne_12_2 = 1,
.ne_12_3 = 1,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_keys * row_bytes_f16,
.nb13 = (uint64_t)n_keys * row_bytes_f16,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_keys * row_bytes_f16,
.nb23 = (uint64_t)n_keys * row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_keys * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
}
ds4_gpu_flash_attn_vec_args vec_args = {
.ne01 = (int32_t)n_tokens,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * row_bytes,
.nb02 = row_bytes,
.nb03 = (uint64_t)n_tokens * n_head * row_bytes,
.ne11 = (int32_t)n_keys,
.ne_12_2 = 1,
.ne_12_3 = 1,
.ns10 = (int32_t)head_dim,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_keys * row_bytes_f16,
.nb13 = (uint64_t)n_keys * row_bytes_f16,
.ns20 = (int32_t)head_dim,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_keys * row_bytes_f16,
.nb23 = (uint64_t)n_keys * row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_keys * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = (int32_t)n_tokens,
.ne3 = 1,
.scale = 1.0f / sqrtf((float)head_dim),
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) +
4u * ncpsg +
2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg;
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:vec_pipeline];
[enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
[enc setBuffer:mask_buffer offset:0 atIndex:4];
[enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec");
ds4_gpu_flash_attn_reduce_args reduce_args = {
.nrows = (int32_t)nrows,
};
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:reduce_pipeline];
[enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce");
#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE
return 1;
}
static int ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(
id<MTLCommandBuffer> __strong *cbp,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
/* The vector sibling below only handles the small square case; every
* rectangular (TP row-split) problem goes through the long path. */
if (n_tokens >= 20 || q_row0 != 0 || n_q != n_tokens) {
return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec_long(cbp,
heads,
sinks_buf,
sinks_offset,
q,
raw_kv,
comp_kv,
comp_kv_f16,
comp_mask,
use_comp_mask,
q_row0,
n_q,
n_tokens,
n_comp,
window,
ratio,
n_head,
head_dim);
}
return ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_vec(cbp,
heads,
sinks_buf,
sinks_offset,
q,
raw_kv,
comp_kv,
comp_kv_f16,
comp_mask,
use_comp_mask,
n_tokens,
n_comp,
window,
ratio,
n_head,
head_dim);
}
/* Raw prefill FlashAttention over a rectangular problem: q holds n_q query
* rows that correspond to token positions [q_row0, q_row0 + n_q) of the
* chunk, raw_kv holds all n_kv key rows, and heads receives one output row
* per query row. The classic square prefill is q_row0 == 0, n_q == n_kv. */
static int ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(
id<MTLCommandBuffer> __strong *cbp,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_kv,
uint32_t window,
uint32_t n_head,
uint32_t head_dim,
bool raw_kv_f16,
bool has_sinks,
float scale) {
if (!cbp || !*cbp) return 0;
id<MTLCommandBuffer> cb = *cbp;
if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) {
return 0;
}
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
const uint64_t q_bytes = (uint64_t)n_q * n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)n_kv * head_dim *
(raw_kv_f16 ? sizeof(uint16_t) : sizeof(float));
if (!qbuf || !rawbuf || !headsbuf || (has_sinks && !sinks_buf) ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
ds4_gpu_tensor_bytes(heads) < q_bytes) {
fprintf(stderr, "ds4: Metal prefill shared-KV FlashAttention received undersized buffers\n");
return 0;
}
const uint32_t nqptg = 8;
const uint32_t ncpsg = 64;
const uint32_t nsg = head_dim >= 512 ? 8u : 4u;
const bool has_kvpad = (n_kv % ncpsg) != 0;
const bool bc_mask = (n_q % nqptg) != 0;
const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_q * (NSUInteger)n_kv * sizeof(uint16_t);
const NSUInteger kv_bytes = (NSUInteger)n_kv * row_bytes_f16;
const NSUInteger pad_bytes = has_kvpad
? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_q * sizeof(uint16_t))
: 1u;
const NSUInteger nblk0 = ((NSUInteger)n_kv + ncpsg - 1u) / ncpsg;
const NSUInteger nblk1 = ((NSUInteger)n_q + nqptg - 1u) / nqptg;
const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);
bool mask_cache_created = false;
ds4_gpu_zero_prefix_prefill_mask_cache_entry *mask_cache =
q_row0 == 0u && n_q == n_kv
? ds4_gpu_get_zero_prefix_prefill_mask_cache(
DS4_GPU_PREFILL_MASK_CACHE_RAW,
n_kv,
0u,
n_kv,
window,
0u,
nqptg,
ncpsg,
has_kvpad,
bc_mask,
mask_bytes,
blk_bytes,
&mask_cache_created)
: NULL;
id<MTLBuffer> mask_buffer = mask_cache
? mask_cache->mask
: ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
if (mask_cache && mask_cache_created) {
ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents],
0u, n_kv, n_kv, window);
mask_cache->valid = true;
}
if (!mask_buffer ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_bytes,
"ds4_flash_attn_kv_f16") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_flash_attn_pad") ||
(!mask_cache &&
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
&g_flash_attn_blk_bytes,
blk_bytes,
"ds4_flash_attn_blk"))) {
return 0;
}
id<MTLBuffer> blk_buffer = mask_cache
? mask_cache->blk
: g_flash_attn_blk_buffer;
const bool flash_stage_profile =
getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil;
double flash_stage_t0 = 0.0;
if (flash_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
int profile_owned = 0;
cb = ds4_gpu_command_buffer(&profile_owned);
if (!cb || profile_owned) return 0;
*cbp = cb;
flash_stage_t0 = ds4_gpu_now_ms();
}
#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \
if (flash_stage_profile) { \
if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \
"raw_nonvec", (name), n_q, 0, n_kv, \
n_head, head_dim, window, 0, &flash_stage_t0)) { \
return 0; \
} \
cb = *cbp; \
} \
} while (0)
if (!mask_cache) {
ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents],
q_row0, n_q, n_kv, window);
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill");
id<MTLComputePipelineState> pad_pipeline = nil;
if (has_kvpad) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
id<MTLComputePipelineState> blk_pipeline =
ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
id<MTLComputePipelineState> attn_pipeline =
ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512",
true, has_sinks, false, false, has_kvpad, bc_mask,
(int32_t)head_dim,
(int32_t)head_dim,
(int32_t)nsg);
if (!blk_pipeline || !attn_pipeline) return 0;
if (!ds4_gpu_encode_copy_to_f16_1d(cb,
rawbuf,
ds4_gpu_tensor_offset(raw_kv),
raw_kv_f16,
g_flash_attn_kv_buffer,
0,
n_kv * head_dim)) {
return 0;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw");
if (has_kvpad) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)n_kv,
.ne_12_2 = 1,
.ne_12_3 = 1,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_kv * row_bytes_f16,
.nb13 = (uint64_t)n_kv * row_bytes_f16,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_kv * row_bytes_f16,
.nb23 = (uint64_t)n_kv * row_bytes_f16,
.ne31 = (int32_t)n_q,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_kv * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
}
ds4_gpu_flash_attn_blk_args blk_args = {
.ne01 = (int32_t)n_q,
.ne30 = (int32_t)n_kv,
.ne31 = (int32_t)n_q,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_kv * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = nil;
if (!mask_cache || !mask_cache->blk_ready) {
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:blk_pipeline];
[enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
[enc setBuffer:mask_buffer offset:0 atIndex:1];
[enc setBuffer:blk_buffer offset:0 atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (mask_cache) mask_cache->blk_ready = true;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("block_map");
ds4_gpu_flash_attn_vec_args args = {
.ne01 = (int32_t)n_q,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * row_bytes,
.nb02 = row_bytes,
.nb03 = (uint64_t)n_q * n_head * row_bytes,
.ne11 = (int32_t)n_kv,
.ne_12_2 = 1,
.ne_12_3 = 1,
.ns10 = (int32_t)head_dim,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_kv * row_bytes_f16,
.nb13 = (uint64_t)n_kv * row_bytes_f16,
.ns20 = (int32_t)head_dim,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_kv * row_bytes_f16,
.nb23 = (uint64_t)n_kv * row_bytes_f16,
.ne31 = (int32_t)n_q,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_kv * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = (int32_t)n_q,
.ne3 = 1,
.scale = scale,
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u);
const NSUInteger shared_elems = (NSUInteger)nqptg *
((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:attn_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
[enc setBuffer:mask_buffer offset:0 atIndex:4];
[enc setBuffer:has_sinks ? sinks_buf : qbuf
offset:has_sinks ? sinks_offset : ds4_gpu_tensor_offset(q)
atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:blk_buffer offset:0 atIndex:7];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention");
#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE
return 1;
}
static int ds4_gpu_encode_flash_attention_prefill_raw_heads(
id<MTLCommandBuffer> __strong *cbp,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_kv,
uint32_t window,
uint32_t n_head,
uint32_t head_dim) {
if (!cbp || !*cbp) return 0;
id<MTLCommandBuffer> cb = *cbp;
if (head_dim != 512 || n_head == 0 || n_q == 0 || n_kv == 0) {
return 0;
}
/* The vector sibling below only handles the small square case; every
* rectangular (TP row-split) problem goes through the non-vector path. */
if (n_kv >= 20 || q_row0 != 0 || n_q != n_kv) {
return ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(cbp,
heads,
sinks_buf,
sinks_offset,
q,
raw_kv,
q_row0,
n_q,
n_kv,
window,
n_head,
head_dim,
false,
true,
1.0f / sqrtf((float)head_dim));
}
const uint32_t n_tokens = n_q;
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)n_tokens * head_dim * sizeof(float);
if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
ds4_gpu_tensor_bytes(heads) < q_bytes) {
fprintf(stderr, "ds4: Metal prefill raw DS4 FlashAttention received undersized buffers\n");
return 0;
}
const uint32_t ncpsg = 32;
const uint32_t nwg = 32;
const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_tokens, nwg, ncpsg);
const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)n_tokens * sizeof(uint16_t);
const NSUInteger kv_f16_offset = 0;
const NSUInteger kv_f16_bytes = (NSUInteger)n_tokens * row_bytes_f16;
const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 +
(NSUInteger)ncpsg * (NSUInteger)n_tokens * sizeof(uint16_t);
const NSUInteger nrows = (NSUInteger)n_tokens * n_head;
const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) +
nrows * (2u * (NSUInteger)nwg) * sizeof(float);
id<MTLBuffer> mask_buffer =
ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
if (!mask_buffer ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_flash_attn_pad") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_f16_bytes,
"ds4_flash_attn_kv_f16") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer,
&g_flash_attn_tmp_bytes,
tmp_bytes,
"ds4_flash_attn_tmp")) {
return 0;
}
const bool flash_stage_profile =
getenv("DS4_METAL_FLASH_ATTN_STAGE_PROFILE") != NULL && g_batch_cb != nil;
double flash_stage_t0 = 0.0;
if (flash_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
int profile_owned = 0;
cb = ds4_gpu_command_buffer(&profile_owned);
if (!cb || profile_owned) return 0;
*cbp = cb;
flash_stage_t0 = ds4_gpu_now_ms();
}
#define DS4_METAL_PROFILE_FLASH_ATTN_STAGE(name) do { \
if (flash_stage_profile) { \
if (!ds4_gpu_flash_attn_stage_profile_boundary(cbp, \
"raw_vec", (name), n_tokens, 0, n_tokens, \
n_head, head_dim, window, 0, &flash_stage_t0)) { \
return 0; \
} \
cb = *cbp; \
} \
} while (0)
ds4_gpu_fill_raw_prefill_mask((uint16_t *)[mask_buffer contents], 0, n_tokens, n_tokens, window);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("mask_fill");
id<MTLComputePipelineState> pad_pipeline = nil;
if ((n_tokens % ncpsg) != 0) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
id<MTLComputePipelineState> vec_pipeline =
ds4_gpu_get_flash_attn_vec_pipeline("kernel_flash_attn_ext_vec_f16_dk512_dv512",
true, true, false, false, true,
false,
(int32_t)head_dim,
(int32_t)head_dim,
(int32_t)nsg,
(int32_t)nwg);
id<MTLComputePipelineState> reduce_pipeline =
ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim, (int32_t)nwg);
if (!vec_pipeline || !reduce_pipeline) return 0;
if (!ds4_gpu_encode_cpy_f32_f16_1d(cb,
rawbuf,
ds4_gpu_tensor_offset(raw_kv),
g_flash_attn_kv_buffer,
kv_f16_offset,
n_tokens * head_dim)) {
return 0;
}
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("copy_raw");
if ((n_tokens % ncpsg) != 0) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)n_tokens,
.ne_12_2 = 1,
.ne_12_3 = 1,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_tokens * row_bytes_f16,
.nb13 = (uint64_t)n_tokens * row_bytes_f16,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_tokens * row_bytes_f16,
.nb23 = (uint64_t)n_tokens * row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_tokens * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2];
[enc setBuffer:mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("pad");
}
ds4_gpu_flash_attn_vec_args vec_args = {
.ne01 = (int32_t)n_tokens,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * row_bytes,
.nb02 = row_bytes,
.nb03 = (uint64_t)n_tokens * n_head * row_bytes,
.ne11 = (int32_t)n_tokens,
.ne_12_2 = 1,
.ne_12_3 = 1,
.ns10 = (int32_t)head_dim,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_tokens * row_bytes_f16,
.nb13 = (uint64_t)n_tokens * row_bytes_f16,
.ns20 = (int32_t)head_dim,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_tokens * row_bytes_f16,
.nb23 = (uint64_t)n_tokens * row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_tokens * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = (int32_t)n_tokens,
.ne3 = 1,
.scale = 1.0f / sqrtf((float)head_dim),
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) +
4u * ncpsg +
2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg;
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:vec_pipeline];
[enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:kv_f16_offset atIndex:3];
[enc setBuffer:mask_buffer offset:0 atIndex:4];
[enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_head, nwg)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_vec");
ds4_gpu_flash_attn_reduce_args reduce_args = {
.nrows = (int32_t)nrows,
};
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:reduce_pipeline];
[enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_FLASH_ATTN_STAGE("attention_reduce");
#undef DS4_METAL_PROFILE_FLASH_ATTN_STAGE
return 1;
}
static int ds4_gpu_encode_flash_attention_gathered_heads(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_ids,
uint32_t n_comp,
const ds4_gpu_tensor *comp_mask,
uint32_t use_mask,
uint32_t n_head,
uint32_t head_dim) {
const uint32_t n_keys = n_raw + n_comp;
if (head_dim != 512 || n_head == 0 || n_raw == 0 || n_keys == 0 ||
raw_cap < n_raw || n_keys < n_raw) {
return 0;
}
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : nil;
id<MTLBuffer> idsbuf = comp_ids ? ds4_gpu_tensor_buffer(comp_ids) : nil;
if (comp_ids && (!idsbuf || ds4_gpu_tensor_bytes(comp_ids) < (uint64_t)n_comp * sizeof(int32_t))) return 0;
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
id<MTLBuffer> maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : nil;
const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
(comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u;
if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
(n_comp && !compbuf) ||
(use_mask && !maskbuf) ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
(n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) ||
ds4_gpu_tensor_bytes(heads) < q_bytes ||
(use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes)) {
fprintf(stderr, "ds4: Metal gathered DS4 FlashAttention received undersized buffers\n");
return 0;
}
const uint32_t ncpsg = 32;
/* Split-K work groups per head. 32 keys per group is the smallest
* useful unit; at short contexts 32 groups meant 4 keys each plus a
* 32-way reduce. DS4_METAL_FLASH_NWG=N forces a value (0 = adaptive). */
static int nwg_env = -1;
if (nwg_env < 0) {
const char *e = getenv("DS4_METAL_FLASH_NWG");
nwg_env = e ? atoi(e) : 0;
}
uint32_t nwg = 32;
if (nwg_env > 0) {
nwg = (uint32_t)nwg_env;
} else {
/* 32 always: 8 groups at 128 keys saved 1 us in attention but made the
* later concurrent FFN block ~50 us slower (unexplained, measured);
* DS4_METAL_FLASH_NWG overrides for experiments. */
nwg = 32u;
}
if (nwg > 32u) nwg = 32u;
if (nwg < 1u) nwg = 1u;
const uint32_t nsg = ds4_gpu_flash_attn_vec_nsg(n_keys, nwg, ncpsg);
const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_keys * sizeof(uint16_t);
const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16;
const NSUInteger pad_bytes = 2u * (NSUInteger)ncpsg * row_bytes_f16 +
(NSUInteger)ncpsg * sizeof(uint16_t);
const NSUInteger nrows = (NSUInteger)n_head;
const NSUInteger tmp_bytes = nrows * (NSUInteger)head_dim * (NSUInteger)nwg * sizeof(float) +
nrows * (2u * (NSUInteger)nwg) * sizeof(float);
const bool packed_shape =
ds4_gpu_ported_m5_decode_feature_enabled(
"DS4_METAL_DISABLE_PRE_M5_FLASH_ATTN_PACKED32_REDUCE",
NULL) &&
!g_quality_mode && use_mask == 0u && comp_kv_f16 != 0u && n_comp != 0u &&
n_head == 64u && head_dim == 512u && nsg == 1u && nwg == 32u &&
n_keys <= 1024u && g_decode_attn_rope_fuse != 0 &&
g_decode_attn_rope_args.head_dim == 512 &&
g_decode_attn_rope_args.n_dims == 64 &&
g_decode_attn_rope_args.row_bytes == 2048 &&
g_decode_attn_rope_args.inverse != 0;
const bool m5_persistent_zero_mask =
ds4_gpu_device_is_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_M5_PERSISTENT_ZERO_ATTN_MASK") == NULL &&
(!packed_shape ||
getenv("DS4_METAL_DISABLE_M5_PACKED_ZERO_MASK") == NULL);
const bool use_persistent_zero_mask =
use_mask == 0u &&
(ds4_gpu_device_name_contains("M3") ||
m5_persistent_zero_mask) &&
getenv("DS4_METAL_DISABLE_PERSISTENT_ZERO_ATTN_MASK") == NULL;
if (!(use_persistent_zero_mask
? ds4_gpu_ensure_zero_attention_mask(mask_bytes)
: ds4_gpu_ensure_scratch_buffer(&g_flash_attn_mask_buffer,
&g_flash_attn_mask_bytes,
mask_bytes,
"ds4_flash_attn_mask")) ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_bytes,
"ds4_flash_attn_kv") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_flash_attn_pad") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_tmp_buffer,
&g_flash_attn_tmp_bytes,
tmp_bytes,
"ds4_flash_attn_tmp")) {
return 0;
}
id<MTLBuffer> flash_mask_buffer = use_persistent_zero_mask
? g_flash_attn_zero_mask_buffer
: g_flash_attn_mask_buffer;
const bool has_kvpad = (n_keys % ncpsg) != 0;
const bool use_shared_kvpad =
has_kvpad &&
ds4_gpu_device_name_contains("M3") &&
getenv("DS4_METAL_DISABLE_SHARED_KV_PAD") == NULL;
const NSUInteger packed_threads = 8u * 32u;
const NSUInteger packed_shared_bytes =
512u * sizeof(uint16_t) +
8u * 128u * sizeof(uint16_t) +
(32u * 32u + 2u * 32u + 32u) * sizeof(float) +
32u * 33u * 4u * sizeof(float);
const bool packed_requested =
packed_shape &&
getenv("DS4_METAL_DISABLE_M5_FLASH_ATTN_PACKED32_REDUCE") == NULL;
id<MTLComputePipelineState> packed_pipeline = nil;
if (packed_requested) {
packed_pipeline = ds4_gpu_get_flash_attn_vec_pipeline(
"kernel_dsv4_flash_attn_vec_packed32_reduce_rope_f16_dk512_dv512",
true, true, false, false, has_kvpad, use_shared_kvpad,
(int32_t)head_dim, (int32_t)head_dim, 1, 32);
}
const NSUInteger max_tgmem = [g_device maxThreadgroupMemoryLength];
const bool use_packed =
packed_pipeline != nil &&
packed_pipeline.threadExecutionWidth == 32u &&
packed_pipeline.maxTotalThreadsPerThreadgroup >= packed_threads &&
(max_tgmem == 0u || max_tgmem >= packed_shared_bytes);
if (packed_requested &&
getenv("DS4_METAL_TRACE_M5_FLASH_ATTN_PACKED32_REDUCE") != NULL) {
fprintf(stderr,
"ds4: packed FA use=%d max_threads=%lu tew=%lu tgmem=%lu need=%lu\n",
use_packed ? 1 : 0,
(unsigned long)(packed_pipeline ?
packed_pipeline.maxTotalThreadsPerThreadgroup : 0u),
(unsigned long)(packed_pipeline ?
packed_pipeline.threadExecutionWidth : 0u),
(unsigned long)max_tgmem,
(unsigned long)packed_shared_bytes);
}
id<MTLComputePipelineState> vec_pipeline = nil;
id<MTLComputePipelineState> reduce_pipeline = nil;
if (!use_packed) {
vec_pipeline = ds4_gpu_get_flash_attn_vec_pipeline(
"kernel_flash_attn_ext_vec_f16_dk512_dv512",
true, true, false, false, has_kvpad, use_shared_kvpad,
(int32_t)head_dim, (int32_t)head_dim,
(int32_t)nsg, (int32_t)nwg);
reduce_pipeline =
ds4_gpu_get_flash_attn_reduce_pipeline((int32_t)head_dim,
(int32_t)nwg);
if (!vec_pipeline || !reduce_pipeline) return 0;
}
if (!use_persistent_zero_mask &&
!ds4_gpu_encode_fill_f16_1d(cb, flash_mask_buffer, 0, n_keys, 0.0f)) {
return 0;
}
if (use_mask && n_comp &&
!ds4_gpu_encode_cpy_f32_f16_1d(cb,
maskbuf,
ds4_gpu_tensor_offset(comp_mask),
flash_mask_buffer,
(NSUInteger)n_raw * sizeof(uint16_t),
n_comp)) {
return 0;
}
bool pad_fused = false;
if (!ds4_gpu_encode_flash_kv_stage_f16(
cb,
rawbuf,
ds4_gpu_tensor_offset(raw_kv),
raw_cap,
raw_start,
n_raw,
compbuf,
ds4_gpu_tensor_offset(comp_kv),
comp_kv_f16 != 0,
idsbuf,
comp_ids ? ds4_gpu_tensor_offset(comp_ids) : 0,
n_comp,
head_dim,
g_flash_attn_kv_buffer,
0,
flash_mask_buffer,
0,
g_flash_attn_pad_buffer,
0,
has_kvpad,
use_shared_kvpad,
&pad_fused)) {
return 0;
}
id<MTLComputePipelineState> pad_pipeline = nil;
if (has_kvpad && !pad_fused) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
if (has_kvpad && !pad_fused) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)n_keys,
.ne_12_2 = 1,
.ne_12_3 = 1,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_keys * row_bytes_f16,
.nb13 = (uint64_t)n_keys * row_bytes_f16,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_keys * row_bytes_f16,
.nb23 = (uint64_t)n_keys * row_bytes_f16,
.ne31 = 1,
.ne32 = 1,
.ne33 = 1,
.nb31 = mask_bytes,
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:flash_mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
ds4_gpu_flash_attn_vec_args vec_args = {
.ne01 = 1,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * row_bytes,
.nb02 = row_bytes,
.nb03 = (uint64_t)n_head * row_bytes,
.ne11 = (int32_t)n_keys,
.ne_12_2 = 1,
.ne_12_3 = 1,
.ns10 = (int32_t)head_dim,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_keys * row_bytes_f16,
.nb13 = (uint64_t)n_keys * row_bytes_f16,
.ns20 = (int32_t)head_dim,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_keys * row_bytes_f16,
.nb23 = (uint64_t)n_keys * row_bytes_f16,
.ne31 = 1,
.ne32 = 1,
.ne33 = 1,
.nb31 = mask_bytes,
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = 1,
.ne3 = 1,
.scale = 1.0f / sqrtf((float)head_dim),
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger shared_elems = (ds4_gpu_align_up_ns(head_dim, 128u) +
4u * ncpsg +
2u * ds4_gpu_align_up_ns(head_dim, 128u)) * nsg;
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
if (use_packed) {
id<MTLComputeCommandEncoder> packed_enc = ds4_gpu_compute_encoder(cb);
[packed_enc setComputePipelineState:packed_pipeline];
[packed_enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
[packed_enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[packed_enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[packed_enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
[packed_enc setBuffer:flash_mask_buffer offset:0 atIndex:4];
[packed_enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
[packed_enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[packed_enc setBuffer:headsbuf
offset:ds4_gpu_tensor_offset(heads) atIndex:7];
[packed_enc setBytes:&g_decode_attn_rope_args
length:sizeof(g_decode_attn_rope_args) atIndex:8];
[packed_enc setThreadgroupMemoryLength:packed_shared_bytes atIndex:0];
[packed_enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(packed_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, packed_enc);
g_decode_attn_rope_fuse = 0;
g_decode_attn_rope_fuse_used = 1;
return 1;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:vec_pipeline];
[enc setBytes:&vec_args length:sizeof(vec_args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
[enc setBuffer:flash_mask_buffer offset:0 atIndex:4];
[enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:7];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, n_head, nwg)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
ds4_gpu_flash_attn_reduce_args reduce_args = {
.nrows = (int32_t)nrows,
};
/* When ds4.c defers the inverse RoPE tail, the reduce threadgroup that owns
* a head's whole row rotates it in place, removing a dispatch per layer. */
const int fuse_rope = g_decode_attn_rope_fuse;
g_decode_attn_rope_fuse = 0;
if (fuse_rope) g_decode_attn_rope_fuse_used = 1;
id<MTLComputePipelineState> reduce_pso = reduce_pipeline;
if (fuse_rope) {
reduce_pso = ds4_gpu_get_flash_attn_reduce_rope_pipeline((int32_t)head_dim,
(int32_t)nwg);
if (!reduce_pso) return 0;
}
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:reduce_pso];
[enc setBytes:&reduce_args length:sizeof(reduce_args) atIndex:0];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:2];
if (fuse_rope) {
[enc setBytes:&g_decode_attn_rope_args
length:sizeof(g_decode_attn_rope_args) atIndex:3];
}
[enc dispatchThreadgroups:MTLSizeMake(nrows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32u * nwg, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_flash_attention_decode_raw_batch_heads(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t window,
uint32_t n_head,
uint32_t head_dim,
bool noncausal,
const int32_t *visual_tokens,
uint32_t vocab_size) {
if (head_dim != 512 || n_head == 0 || n_tokens == 0 ||
n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) {
return 0;
}
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
if (!qbuf || !rawbuf || !headsbuf || !sinks_buf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
ds4_gpu_tensor_bytes(heads) < q_bytes) {
fprintf(stderr, "ds4: Metal decode raw batch FlashAttention received undersized buffers\n");
return 0;
}
const uint32_t nqptg = 8;
const uint32_t ncpsg = 64;
const uint32_t nsg = head_dim >= 512 ? 8u : 4u;
const bool has_kvpad = (n_raw % ncpsg) != 0;
const bool bc_mask = (n_tokens % nqptg) != 0;
const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_raw * (NSUInteger)n_tokens * sizeof(uint16_t);
const NSUInteger kv_bytes = (NSUInteger)n_raw * row_bytes_f16;
const NSUInteger pad_bytes = has_kvpad
? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t))
: 1u;
const NSUInteger nblk0 = ((NSUInteger)n_raw + ncpsg - 1u) / ncpsg;
const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg;
const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);
id<MTLBuffer> mask_buffer =
ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
if (!mask_buffer ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_bytes,
"ds4_flash_attn_kv_f16") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_flash_attn_pad") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
&g_flash_attn_blk_bytes,
blk_bytes,
"ds4_flash_attn_blk")) {
return 0;
}
if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb,
rawbuf,
ds4_gpu_tensor_offset(raw_kv),
raw_cap,
raw_start,
n_raw,
head_dim,
g_flash_attn_kv_buffer,
0)) {
return 0;
}
if (visual_tokens) {
if (!ds4_gpu_fill_visual_mixed_batch_mask(
(uint16_t *)[mask_buffer contents],
visual_tokens, vocab_size, n_tokens, n_raw, 0u,
pos0, window, 0u)) return 0;
} else if (noncausal) {
memset([mask_buffer contents], 0, mask_bytes);
} else {
ds4_gpu_fill_raw_decode_batch_mask((uint16_t *)[mask_buffer contents],
n_tokens,
n_raw,
pos0,
window);
}
id<MTLComputePipelineState> pad_pipeline = nil;
if (has_kvpad) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
id<MTLComputePipelineState> blk_pipeline =
ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
id<MTLComputePipelineState> attn_pipeline =
ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512",
true, true, false, false, has_kvpad, bc_mask,
(int32_t)head_dim,
(int32_t)head_dim,
(int32_t)nsg);
if (!blk_pipeline || !attn_pipeline) return 0;
if (has_kvpad) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)n_raw,
.ne_12_2 = 1,
.ne_12_3 = 1,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_raw * row_bytes_f16,
.nb13 = (uint64_t)n_raw * row_bytes_f16,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_raw * row_bytes_f16,
.nb23 = (uint64_t)n_raw * row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_raw * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
ds4_gpu_flash_attn_blk_args blk_args = {
.ne01 = (int32_t)n_tokens,
.ne30 = (int32_t)n_raw,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_raw * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:blk_pipeline];
[enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
[enc setBuffer:mask_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
ds4_gpu_flash_attn_vec_args args = {
.ne01 = (int32_t)n_tokens,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * row_bytes,
.nb02 = row_bytes,
.nb03 = (uint64_t)n_tokens * n_head * row_bytes,
.ne11 = (int32_t)n_raw,
.ne_12_2 = 1,
.ne_12_3 = 1,
.ns10 = (int32_t)head_dim,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_raw * row_bytes_f16,
.nb13 = (uint64_t)n_raw * row_bytes_f16,
.ns20 = (int32_t)head_dim,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_raw * row_bytes_f16,
.nb23 = (uint64_t)n_raw * row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_raw * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = (int32_t)n_tokens,
.ne3 = 1,
.scale = 1.0f / sqrtf((float)head_dim),
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u);
const NSUInteger shared_elems = (NSUInteger)nqptg *
((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:attn_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
[enc setBuffer:mask_buffer offset:0 atIndex:4];
[enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *heads,
id<MTLBuffer> sinks_buf,
NSUInteger sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim,
const int32_t *visual_tokens,
uint32_t vocab_size) {
if (n_comp == 0) {
return ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb,
heads,
sinks_buf,
sinks_offset,
q,
raw_kv,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
window,
n_head,
head_dim,
false,
visual_tokens,
vocab_size);
}
if (head_dim != 512 || n_head == 0 || n_tokens == 0 ||
n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
ratio == 0 || !comp_kv || (use_comp_mask && !comp_mask)) {
return 0;
}
const uint32_t n_keys = n_raw + n_comp;
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_kv);
id<MTLBuffer> maskbuf = use_comp_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf;
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
(comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
const uint64_t comp_mask_bytes = use_comp_mask ? (uint64_t)n_comp * n_tokens * sizeof(float) : 0u;
if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf || !sinks_buf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
ds4_gpu_tensor_bytes(comp_kv) < comp_bytes ||
(use_comp_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) ||
ds4_gpu_tensor_bytes(heads) < q_bytes) {
fprintf(stderr, "ds4: Metal decode mixed batch FlashAttention received undersized buffers\n");
return 0;
}
const uint32_t nqptg = 8;
const uint32_t ncpsg = 64;
const uint32_t nsg = head_dim >= 512 ? 8u : 4u;
const bool has_kvpad = (n_keys % ncpsg) != 0;
const bool bc_mask = (n_tokens % nqptg) != 0;
const NSUInteger row_bytes = (NSUInteger)head_dim * sizeof(float);
const NSUInteger row_bytes_f16 = (NSUInteger)head_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_keys * (NSUInteger)n_tokens * sizeof(uint16_t);
const NSUInteger kv_bytes = (NSUInteger)n_keys * row_bytes_f16;
const NSUInteger pad_bytes = has_kvpad
? (NSUInteger)ncpsg * (2u * row_bytes_f16 + (NSUInteger)n_tokens * sizeof(uint16_t))
: 1u;
const NSUInteger nblk0 = ((NSUInteger)n_keys + ncpsg - 1u) / ncpsg;
const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg;
const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);
id<MTLBuffer> mask_buffer =
ds4_gpu_new_transient_buffer(mask_bytes, "ds4_flash_attn_mask");
if (!mask_buffer ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_bytes,
"ds4_flash_attn_kv_f16") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_flash_attn_pad") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
&g_flash_attn_blk_bytes,
blk_bytes,
"ds4_flash_attn_blk")) {
return 0;
}
if (!ds4_gpu_encode_copy_raw_ring_to_f16(cb,
rawbuf,
ds4_gpu_tensor_offset(raw_kv),
raw_cap,
raw_start,
n_raw,
head_dim,
g_flash_attn_kv_buffer,
0) ||
!ds4_gpu_encode_copy_to_f16_1d(cb,
compbuf,
ds4_gpu_tensor_offset(comp_kv),
comp_kv_f16 != 0,
g_flash_attn_kv_buffer,
(NSUInteger)n_raw * row_bytes_f16,
n_comp * head_dim)) {
return 0;
}
if (visual_tokens) {
if (!ds4_gpu_fill_visual_mixed_batch_mask(
(uint16_t *)[mask_buffer contents],
visual_tokens, vocab_size, n_tokens, n_raw, n_comp,
pos0, window, ratio)) return 0;
} else {
ds4_gpu_fill_mixed_decode_batch_mask((uint16_t *)[mask_buffer contents],
n_tokens,
n_raw,
n_comp,
pos0,
window,
ratio);
}
if (use_comp_mask) {
if (!ds4_gpu_encode_cpy_f32_f16_2d(cb,
maskbuf,
ds4_gpu_tensor_offset(comp_mask),
mask_buffer,
(NSUInteger)n_raw * sizeof(uint16_t),
n_comp,
n_tokens,
(uint64_t)n_comp * sizeof(float),
(uint64_t)n_keys * sizeof(uint16_t))) {
return 0;
}
}
id<MTLComputePipelineState> pad_pipeline = nil;
if (has_kvpad) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
id<MTLComputePipelineState> blk_pipeline =
ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
id<MTLComputePipelineState> attn_pipeline =
ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk512_dv512",
true, true, false, false, has_kvpad, bc_mask,
(int32_t)head_dim,
(int32_t)head_dim,
(int32_t)nsg);
if (!blk_pipeline || !attn_pipeline) return 0;
if (has_kvpad) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)n_keys,
.ne_12_2 = 1,
.ne_12_3 = 1,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_keys * row_bytes_f16,
.nb13 = (uint64_t)n_keys * row_bytes_f16,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_keys * row_bytes_f16,
.nb23 = (uint64_t)n_keys * row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_keys * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
ds4_gpu_flash_attn_blk_args blk_args = {
.ne01 = (int32_t)n_tokens,
.ne30 = (int32_t)n_keys,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_keys * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:blk_pipeline];
[enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
[enc setBuffer:mask_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
ds4_gpu_flash_attn_vec_args args = {
.ne01 = (int32_t)n_tokens,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * row_bytes,
.nb02 = row_bytes,
.nb03 = (uint64_t)n_tokens * n_head * row_bytes,
.ne11 = (int32_t)n_keys,
.ne_12_2 = 1,
.ne_12_3 = 1,
.ns10 = (int32_t)head_dim,
.nb11 = row_bytes_f16,
.nb12 = (uint64_t)n_keys * row_bytes_f16,
.nb13 = (uint64_t)n_keys * row_bytes_f16,
.ns20 = (int32_t)head_dim,
.nb21 = row_bytes_f16,
.nb22 = (uint64_t)n_keys * row_bytes_f16,
.nb23 = (uint64_t)n_keys * row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)n_keys * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = (int32_t)n_tokens,
.ne3 = 1,
.scale = 1.0f / sqrtf((float)head_dim),
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger padded_v = ds4_gpu_align_up_ns(head_dim, 64u);
const NSUInteger shared_elems = (NSUInteger)nqptg *
((NSUInteger)head_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:attn_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:0 atIndex:3];
[enc setBuffer:mask_buffer offset:0 atIndex:4];
[enc setBuffer:sinks_buf offset:sinks_offset atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
int ds4_gpu_attention_prefill_raw_heads_range_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_kv,
uint32_t window,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_kv == 0) return 0;
@autoreleasepool {
if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
sinks_offset,
(uint64_t)n_head * sizeof(float),
&sinks_inner);
if (!sinks_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_prefill_raw_heads(&cb,
heads,
sinks_buf,
(NSUInteger)sinks_inner,
q,
raw_kv,
q_row0,
n_q,
n_kv,
window,
n_head,
head_dim)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill raw attention heads")) return 0;
}
return 1;
}
int ds4_gpu_attention_prefill_raw_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t window,
uint32_t n_head,
uint32_t head_dim) {
return ds4_gpu_attention_prefill_raw_heads_range_tensor(heads,
model_map,
model_size,
sinks_offset,
q,
raw_kv,
0,
n_tokens,
n_tokens,
window,
n_head,
head_dim);
}
int ds4_gpu_glm_attention_dense_compact_lora_causal_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_kv,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!lora_out || !qk_low || !kv_lora_cache ||
n_q == 0 || n_kv == 0 || n_kv > cache_cap ||
q_row0 >= n_kv || n_q > n_kv - q_row0 ||
n_head == 0 || kv_lora_dim != 512u || qk_dim == 0) {
return 0;
}
@autoreleasepool {
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_prefill_raw_heads_nonvec(
&cb,
lora_out,
nil,
0,
qk_low,
kv_lora_cache,
q_row0,
n_q,
n_kv,
0,
n_head,
kv_lora_dim,
cache_f16,
false,
1.0f / sqrtf((float)qk_dim))) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(
cb, owned, "GLM dense compact prefill attention")) {
return 0;
}
}
return 1;
}
int ds4_gpu_attention_decode_raw_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t window,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 ||
n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) {
return 0;
}
@autoreleasepool {
if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
sinks_offset,
(uint64_t)n_head * sizeof(float),
&sinks_inner);
if (!sinks_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb,
heads,
sinks_buf,
(NSUInteger)sinks_inner,
q,
raw_kv,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
window,
n_head,
head_dim,
false,
NULL,
0)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode raw batch attention heads")) return 0;
}
return 1;
}
int ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_tokens,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !raw_kv || !model_map ||
n_tokens == 0 || n_raw == 0 || raw_cap < n_raw ||
raw_start >= raw_cap || n_head == 0 || head_dim == 0) {
return 0;
}
@autoreleasepool {
const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float);
if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal noncausal attention sinks range is outside the mapped model\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
sinks_offset,
sink_bytes,
&sinks_inner);
if (!sinks_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb,
heads,
sinks_buf,
(NSUInteger)sinks_inner,
q,
raw_kv,
n_tokens,
0,
n_raw,
raw_cap,
raw_start,
0,
n_head,
head_dim,
true,
NULL,
0)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph noncausal raw batch attention heads")) return 0;
}
return 1;
}
int ds4_gpu_attention_decode_mixed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !raw_kv || !model_map || n_tokens == 0 ||
n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
ratio == 0 || (n_comp != 0 && !comp_kv) ||
(use_comp_mask != 0 && !comp_mask)) {
return 0;
}
@autoreleasepool {
if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
sinks_offset,
(uint64_t)n_head * sizeof(float),
&sinks_inner);
if (!sinks_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(cb,
heads,
sinks_buf,
(NSUInteger)sinks_inner,
q,
raw_kv,
comp_kv,
comp_kv_f16,
comp_mask,
use_comp_mask,
n_tokens,
pos0,
n_raw,
raw_cap,
raw_start,
n_comp,
window,
ratio,
n_head,
head_dim,
NULL,
0)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph decode mixed batch attention heads")) return 0;
}
return 1;
}
int ds4_gpu_attention_visual_mixed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t use_comp_mask,
const int32_t *tokens,
uint32_t vocab_size,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !raw_kv || !model_map || !tokens || vocab_size == 0 ||
n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
(n_comp != 0 && (!comp_kv || ratio == 0)) ||
(use_comp_mask != 0 && !comp_mask)) return 0;
@autoreleasepool {
const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float);
if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal visual attention sinks are outside the mapped model\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(
model_map, model_size, sinks_offset, sink_bytes, &sinks_inner);
if (!sinks_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb || !ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(
cb, heads, sinks_buf, (NSUInteger)sinks_inner,
q, raw_kv, comp_kv, comp_kv_f16,
comp_mask, use_comp_mask,
n_tokens, pos0, n_raw, raw_cap, raw_start,
n_comp, window, ratio, n_head, head_dim,
tokens, vocab_size)) return 0;
if (!ds4_gpu_finish_command_buffer(
cb, owned, "graph visual mixed batch attention heads")) return 0;
}
return 1;
}
/* A DSpark verify pass runs; its rows take decode's attention path. */
static int g_v41_verify_rows;
void ds4_gpu_dsv41_verify_rows(int on) { g_v41_verify_rows = on; }
int ds4_gpu_attention_indexed_mixed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *topk,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
uint32_t n_comp,
uint32_t top_k,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !model_map || !q || !raw_kv || !comp_kv || !topk ||
n_tokens == 0 || n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
n_comp == 0 || top_k == 0 || top_k > n_comp || (top_k & (top_k - 1u)) != 0 ||
ratio == 0 || n_head == 0 || head_dim != 512) {
return 0;
}
@autoreleasepool {
if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal indexed attention sinks range is outside the mapped model\n");
return 0;
}
const uint64_t row_bytes = (uint64_t)head_dim * sizeof(float);
const uint64_t row_bytes_f16 = (uint64_t)head_dim * sizeof(uint16_t);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * row_bytes;
const uint64_t raw_bytes = (uint64_t)raw_cap * row_bytes;
const uint64_t comp_bytes = (uint64_t)n_comp * (comp_kv_f16 ? row_bytes_f16 : row_bytes);
const uint64_t topk_bytes = (uint64_t)top_k * n_tokens * sizeof(int32_t);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> compbuf = ds4_gpu_tensor_buffer(comp_kv);
id<MTLBuffer> topkbuf = ds4_gpu_tensor_buffer(topk);
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
if (!qbuf || !rawbuf || !compbuf || !topkbuf || !headsbuf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
ds4_gpu_tensor_bytes(comp_kv) < comp_bytes ||
ds4_gpu_tensor_bytes(topk) < topk_bytes ||
ds4_gpu_tensor_bytes(heads) < q_bytes) {
fprintf(stderr, "ds4: Metal indexed mixed attention received undersized buffers\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
sinks_offset,
(uint64_t)n_head * sizeof(float),
&sinks_inner);
if (!sinks_buf) return 0;
id<MTLComputePipelineState> sort_pipeline =
ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline,
"kernel_dsv4_sort_i32_rows_asc");
const bool decode_one_token = n_tokens == 1u;
/* A few rows take decode's split path and score order: one token's
* work per split keeps the GPU busy and the rows match decode steps. */
const bool decode_rows = (n_tokens == 1u || (n_tokens <= 8u && g_v41_verify_rows)) && !g_quality_mode;
const bool prefill_dual_heads =
!decode_one_token && !g_quality_mode && ds4_gpu_mpp_available() &&
(n_head == 64u || n_head == 32u) &&
top_k == 512u && window == 128u && head_dim == 512u;
const uint32_t decode_splits = decode_rows ? 12u : 1u;
const bool split_decode = decode_splits > 1u;
id<MTLComputePipelineState> attn_pipeline =
split_decode ?
ds4_gpu_hot_pipeline(
g_dsv4_indexed_attention_heads8_split_pipeline,
"kernel_dsv4_indexed_mixed_attention_heads8_split") :
decode_one_token ?
ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_rb16_pipeline,
"kernel_dsv4_indexed_mixed_attention_heads8_rb16") :
prefill_dual_heads ?
ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads16_dual_pipeline,
"kernel_dsv4_indexed_mixed_attention_heads16_dual") :
ds4_gpu_hot_pipeline(g_dsv4_indexed_attention_heads8_pipeline,
"kernel_dsv4_indexed_mixed_attention_heads8");
id<MTLComputePipelineState> split_reduce_pipeline = split_decode ?
ds4_gpu_hot_pipeline(
g_dsv4_indexed_attention_heads8_split_reduce_pipeline,
"kernel_dsv4_indexed_mixed_attention_heads8_split_reduce") :
nil;
if (!sort_pipeline || !attn_pipeline ||
(split_decode && !split_reduce_pipeline)) {
return 0;
}
if ((NSUInteger)top_k > sort_pipeline.maxTotalThreadsPerThreadgroup) {
fprintf(stderr, "ds4: Metal indexed attention top-k exceeds sort threadgroup limit\n");
return 0;
}
/*
* Fast decode attends to the same full top-k compressed rows but keeps
* them in score order, avoiding a chronological sort dispatch.
* --quality restores the sorted order for stricter reproducibility.
*/
const bool skip_decode_sort = decode_rows;
if (!skip_decode_sort &&
!ds4_gpu_ensure_scratch_buffer(&g_indexed_topk_buffer,
&g_indexed_topk_bytes,
(NSUInteger)topk_bytes,
"ds4_indexed_topk_sorted")) {
return 0;
}
ds4_gpu_dsv4_topk_mask_args sort_args = {
.ne00 = (int64_t)top_k,
.ne01 = (int64_t)n_tokens,
.nb00 = sizeof(int32_t),
.nb01 = (uint64_t)top_k * sizeof(int32_t),
.ne0 = (int64_t)top_k,
.ne1 = (int64_t)n_tokens,
.nb0 = sizeof(int32_t),
.nb1 = (uint64_t)top_k * sizeof(int32_t),
};
ds4_gpu_dsv4_indexed_attention_args attn_args = {
.n_tokens = n_tokens,
.n_head = n_head,
.n_raw = n_raw,
.raw_cap = raw_cap,
.raw_start = raw_start,
.n_comp = n_comp,
.top_k = top_k,
.pos0 = pos0,
.window = window,
.ratio = ratio,
.comp_kv_f16 = comp_kv_f16 ? 1u : 0u,
.n_splits = decode_splits,
.q_token_stride = (uint64_t)n_head * row_bytes,
.q_head_stride = row_bytes,
.raw_row_stride = row_bytes,
.comp_row_stride = comp_kv_f16 ? row_bytes_f16 : row_bytes,
.topk_token_stride = (uint64_t)top_k * sizeof(int32_t),
.dst_token_stride = (uint64_t)n_head * row_bytes,
.dst_head_stride = row_bytes,
.scale = 1.0f / sqrtf((float)head_dim),
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = nil;
if (!skip_decode_sort) {
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:sort_pipeline];
[enc setBytes:&sort_args length:sizeof(sort_args) atIndex:0];
[enc setBuffer:topkbuf offset:ds4_gpu_tensor_offset(topk) atIndex:1];
[enc setBuffer:g_indexed_topk_buffer offset:0 atIndex:2];
[enc setThreadgroupMemoryLength:(NSUInteger)top_k * sizeof(int32_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(top_k, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
if (split_decode) {
const uint64_t nrows = (uint64_t)n_tokens * n_head;
const uint64_t partial_bytes =
nrows * head_dim * decode_splits * sizeof(float);
const uint64_t stats_bytes =
nrows * decode_splits * 2u * sizeof(float);
if (partial_bytes > NSUIntegerMax - stats_bytes ||
!ds4_gpu_ensure_scratch_buffer(
&g_flash_attn_tmp_buffer,
&g_flash_attn_tmp_bytes,
(NSUInteger)(partial_bytes + stats_bytes),
"ds4_dsv4_indexed_attention_split_tmp")) {
return 0;
}
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:attn_pipeline];
[enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2];
[enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3];
[enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer
offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0
atIndex:4];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:5];
[enc setThreadgroupMemoryLength:16u * 128u * sizeof(uint16_t) * 4u
atIndex:0];
[enc dispatchThreadgroups:
MTLSizeMake((NSUInteger)n_tokens,
((NSUInteger)n_head + 7u) / 8u,
decode_splits)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:split_reduce_pipeline];
[enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0];
[enc setBuffer:g_flash_attn_tmp_buffer offset:0 atIndex:1];
[enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:2];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)nrows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
} else {
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:attn_pipeline];
[enc setBytes:&attn_args length:sizeof(attn_args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_kv) atIndex:2];
[enc setBuffer:compbuf offset:ds4_gpu_tensor_offset(comp_kv) atIndex:3];
[enc setBuffer:skip_decode_sort ? topkbuf : g_indexed_topk_buffer
offset:skip_decode_sort ? ds4_gpu_tensor_offset(topk) : 0
atIndex:4];
[enc setBuffer:sinks_buf offset:(NSUInteger)sinks_inner atIndex:5];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:6];
[enc setThreadgroupMemoryLength:(decode_one_token ? 16u : 1u) *
128u * 4u * sizeof(uint16_t)
atIndex:0];
[enc dispatchThreadgroups:
MTLSizeMake((NSUInteger)n_tokens,
((NSUInteger)n_head + (prefill_dual_heads ? 15u : 7u)) /
(prefill_dual_heads ? 16u : 8u),
1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph indexed mixed attention heads")) return 0;
}
return 1;
}
int ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t q_row0,
uint32_t n_q,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !raw_kv || !model_map || n_q == 0 || n_tokens == 0 ||
ratio == 0 || (n_comp != 0 && !comp_kv)) {
return 0;
}
@autoreleasepool {
if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
sinks_offset,
(uint64_t)n_head * sizeof(float),
&sinks_inner);
if (!sinks_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb,
heads,
sinks_buf,
(NSUInteger)sinks_inner,
q,
raw_kv,
comp_kv,
comp_kv_f16,
NULL,
0,
q_row0,
n_q,
n_tokens,
n_comp,
window,
ratio,
n_head,
head_dim)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill static mixed attention heads")) return 0;
}
return 1;
}
int ds4_gpu_attention_prefill_static_mixed_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
return ds4_gpu_attention_prefill_static_mixed_heads_range_tensor(heads,
model_map,
model_size,
sinks_offset,
q,
raw_kv,
comp_kv,
comp_kv_f16,
0,
n_tokens,
n_tokens,
n_comp,
window,
ratio,
n_head,
head_dim);
}
int ds4_gpu_attention_prefill_masked_mixed_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
const ds4_gpu_tensor *comp_mask,
uint32_t n_tokens,
uint32_t n_comp,
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !raw_kv || !comp_kv || !comp_mask || !model_map ||
n_tokens == 0 || n_comp == 0 || ratio == 0) {
return 0;
}
@autoreleasepool {
if (sinks_offset > model_size || (uint64_t)n_head * sizeof(float) > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal attention sinks range is outside the mapped model\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size,
sinks_offset,
(uint64_t)n_head * sizeof(float),
&sinks_inner);
if (!sinks_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_prefill_static_mixed_heads_nonvec(&cb,
heads,
sinks_buf,
(NSUInteger)sinks_inner,
q,
raw_kv,
comp_kv,
comp_kv_f16,
comp_mask,
1,
0,
n_tokens,
n_tokens,
n_comp,
window,
ratio,
n_head,
head_dim)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph prefill masked mixed attention heads")) return 0;
}
return 1;
}
int ds4_gpu_attention_decode_heads_tensor(
ds4_gpu_tensor *heads,
const void *model_map,
uint64_t model_size,
uint64_t sinks_offset,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv,
uint32_t n_raw,
uint32_t raw_cap,
uint32_t raw_start,
const ds4_gpu_tensor *comp_kv,
uint32_t comp_kv_f16,
uint32_t n_comp,
const ds4_gpu_tensor *comp_mask,
uint32_t use_mask,
uint32_t n_head,
uint32_t head_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !model_map || !q || !raw_kv ||
n_raw == 0 || n_head == 0 || head_dim == 0 ||
raw_cap < n_raw || raw_start >= raw_cap ||
n_raw > UINT32_MAX - n_comp || n_raw + n_comp > 8192u ||
(n_comp != 0 && !comp_kv) ||
(use_mask != 0 && !comp_mask)) {
return 0;
}
@autoreleasepool {
const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
const uint64_t raw_bytes = (uint64_t)raw_cap * head_dim * sizeof(float);
const uint64_t comp_bytes = (uint64_t)n_comp * head_dim *
(comp_kv_f16 ? sizeof(uint16_t) : sizeof(float));
const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float);
if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) {
fprintf(stderr, "ds4: Metal graph attention heads sink range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_kv);
id<MTLBuffer> compbuf = n_comp ? ds4_gpu_tensor_buffer(comp_kv) : rawbuf;
id<MTLBuffer> maskbuf = use_mask ? ds4_gpu_tensor_buffer(comp_mask) : rawbuf;
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
const uint64_t comp_mask_bytes = use_mask ? (uint64_t)n_comp * sizeof(float) : 0u;
if (!qbuf || !rawbuf || !compbuf || !maskbuf || !headsbuf ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(raw_kv) < raw_bytes ||
(n_comp && ds4_gpu_tensor_bytes(comp_kv) < comp_bytes) ||
(use_mask && ds4_gpu_tensor_bytes(comp_mask) < comp_mask_bytes) ||
ds4_gpu_tensor_bytes(heads) < q_bytes) {
fprintf(stderr, "ds4: Metal graph attention heads received undersized buffers\n");
return 0;
}
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, sinks_offset, sink_bytes, &sinks_inner);
if (!sinks_buf) return 0;
if (n_comp == 0) {
/* The raw-only path stages through the same kernel when it
* stages at all; a pending deferred kv task that nothing
* consumes must run before the attention reads the cache. */
if (g_kv_task.pending && !ds4_gpu_kv_norm_task_flush()) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_raw_heads(cb,
heads,
sinks_buf,
(NSUInteger)sinks_inner,
q,
raw_kv,
n_raw,
raw_cap,
raw_start,
n_head,
head_dim)) {
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph raw attention heads")) return 0;
return 1;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_gathered_heads(cb,
heads,
sinks_buf,
(NSUInteger)sinks_inner,
q,
raw_kv,
n_raw,
raw_cap,
raw_start,
comp_kv,
comp_kv_f16,
NULL,
n_comp,
comp_mask,
use_mask,
n_head,
head_dim)) {
return 0;
}
if (g_kv_task.pending) {
static int warned;
if (!warned) {
warned = 1;
fprintf(stderr, "ds4: deferred kv norm task was not consumed by the attention path\n");
}
if (!ds4_gpu_kv_norm_task_flush()) return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "graph attention heads")) return 0;
}
return 1;
}
/* DeepSeek V4.1 decode attention over the raw ring and the compressed rows
* selected by `ids`, staged by one kernel (0: the staging kernel does not
* apply; the caller gathers into `selected` and runs the plain entry). */
int ds4_gpu_dsv41_attention_decode(ds4_gpu_tensor *heads, const void *model_map, uint64_t model_size,
uint64_t sinks_offset, const ds4_gpu_tensor *q,
const ds4_gpu_tensor *raw_kv, uint32_t n_raw, uint32_t raw_cap,
uint32_t raw_start, const ds4_gpu_tensor *comp_cache,
const ds4_gpu_tensor *ids, ds4_gpu_tensor *selected,
uint32_t n_comp, uint32_t attended, uint32_t n_head, uint32_t head_dim) {
(void)selected;
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !model_map || !q || !raw_kv || !comp_cache || !ids || !n_raw || !attended ||
attended > n_comp || raw_cap < n_raw || raw_start >= raw_cap || head_dim != 512u ||
n_raw + attended > 8192u || g_quality_mode ||
getenv("DS4_METAL_DISABLE_GATHERED_KV_STAGE") != NULL ||
!(ds4_gpu_device_name_contains("M3") || ds4_gpu_device_name_contains("M5")) ||
ds4_gpu_tensor_bytes(comp_cache) < (uint64_t)n_comp * head_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(ids) < (uint64_t)attended * sizeof(int32_t) ||
ds4_gpu_tensor_bytes(q) < (uint64_t)n_head * head_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(heads) < (uint64_t)n_head * head_dim * sizeof(float)) return 0;
@autoreleasepool {
const uint64_t sink_bytes = (uint64_t)n_head * sizeof(float);
if (sinks_offset > model_size || sink_bytes > model_size - sinks_offset) return 0;
uint64_t sinks_inner = 0;
id<MTLBuffer> sinks_buf = ds4_gpu_wrap_model_range(model_map, model_size, sinks_offset,
sink_bytes, &sinks_inner);
if (!sinks_buf) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!ds4_gpu_encode_flash_attention_gathered_heads(cb, heads, sinks_buf, (NSUInteger)sinks_inner,
q, raw_kv, n_raw, raw_cap, raw_start, comp_cache, 0, ids, attended, NULL, 0,
n_head, head_dim)) return 0;
if (g_kv_task.pending && !ds4_gpu_kv_norm_task_flush()) return 0;
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 gathered attention");
}
}
int ds4_gpu_swiglu_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *gate,
const ds4_gpu_tensor *up,
uint32_t n,
float clamp,
float weight) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !gate || !up || n == 0) return 0;
if (!isfinite(clamp) || clamp < 0.0f || !isfinite(weight)) return 0;
@autoreleasepool {
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t bytes = (uint64_t)n * sizeof(float);
if (!gatebuf || !upbuf || !outbuf ||
ds4_gpu_tensor_bytes(gate) < bytes ||
ds4_gpu_tensor_bytes(up) < bytes ||
ds4_gpu_tensor_bytes(out) < bytes) {
fprintf(stderr, "ds4: Metal SwiGLU received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glu_args args = {
.ne00 = (int32_t)n,
.nb01 = (uint64_t)n * sizeof(float),
.ne10 = (int32_t)n,
.nb11 = (uint64_t)n * sizeof(float),
.ne0 = (int32_t)n,
.nb1 = (uint64_t)n * sizeof(float),
.i00 = 0,
.i10 = 0,
.alpha = weight,
.limit = clamp,
};
NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > (NSUInteger)n) nth = (NSUInteger)n;
if (nth == 0u) nth = 1u;
const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_swiglu_flat_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:1];
[enc setBuffer:upbuf offset:ds4_gpu_tensor_offset(up) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "SwiGLU")) return 0;
}
return 1;
}
int ds4_gpu_add_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
uint32_t n) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !a || !b || n == 0) return 0;
@autoreleasepool {
id<MTLBuffer> abuf = ds4_gpu_tensor_buffer(a);
id<MTLBuffer> bbuf = ds4_gpu_tensor_buffer(b);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t bytes = (uint64_t)n * sizeof(float);
if (!abuf || !bbuf || !outbuf ||
ds4_gpu_tensor_bytes(a) < bytes ||
ds4_gpu_tensor_bytes(b) < bytes ||
ds4_gpu_tensor_bytes(out) < bytes) {
fprintf(stderr, "ds4: Metal tensor add received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_add_flat_args args = { .n = n };
NSUInteger nth = g_add2_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > (NSUInteger)n) nth = (NSUInteger)n;
if (nth == 0u) nth = 1u;
const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_add2_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1];
[enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add")) return 0;
}
return 1;
}
int ds4_gpu_add3_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *a,
const ds4_gpu_tensor *b,
const ds4_gpu_tensor *c,
uint32_t n) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !a || !b || !c || n == 0) return 0;
@autoreleasepool {
id<MTLBuffer> abuf = ds4_gpu_tensor_buffer(a);
id<MTLBuffer> bbuf = ds4_gpu_tensor_buffer(b);
id<MTLBuffer> cbuf = ds4_gpu_tensor_buffer(c);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t bytes = (uint64_t)n * sizeof(float);
if (!abuf || !bbuf || !cbuf || !outbuf ||
ds4_gpu_tensor_bytes(a) < bytes ||
ds4_gpu_tensor_bytes(b) < bytes ||
ds4_gpu_tensor_bytes(c) < bytes ||
ds4_gpu_tensor_bytes(out) < bytes) {
fprintf(stderr, "ds4: Metal tensor add3 received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_add_flat_args args = { .n = n };
NSUInteger nth = g_add3_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > (NSUInteger)n) nth = (NSUInteger)n;
if (nth == 0u) nth = 1u;
const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_add3_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:abuf offset:ds4_gpu_tensor_offset(a) atIndex:1];
[enc setBuffer:bbuf offset:ds4_gpu_tensor_offset(b) atIndex:2];
[enc setBuffer:cbuf offset:ds4_gpu_tensor_offset(c) atIndex:3];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "tensor add3")) return 0;
}
return 1;
}
typedef struct {
uint32_t width;
uint32_t rows;
uint32_t layer;
uint32_t n_threads;
float scale;
} ds4_gpu_directional_steering_project_args;
int ds4_gpu_directional_steering_project_tensor(
ds4_gpu_tensor *x,
const ds4_gpu_tensor *directions,
uint32_t layer,
uint32_t width,
uint32_t rows,
float scale) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!x || !directions || width == 0 || rows == 0 || scale == 0.0f) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_dsv4_directional_steering_project_f32");
if (!pipeline) return 0;
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> dbuf = ds4_gpu_tensor_buffer(directions);
const uint64_t x_bytes = (uint64_t)width * rows * sizeof(float);
const uint64_t dir_bytes = (uint64_t)(layer + 1u) * width * sizeof(float);
if (!xbuf || !dbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(directions) < dir_bytes) {
fprintf(stderr, "ds4: Metal directional steering received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
while (nth > width && nth > 1u) nth >>= 1;
if (nth == 0) nth = 1;
ds4_gpu_directional_steering_project_args args = {
.width = width,
.rows = rows,
.layer = layer,
.n_threads = (uint32_t)nth,
.scale = scale,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:dbuf offset:ds4_gpu_tensor_offset(directions) atIndex:2];
[enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "directional steering")) return 0;
}
return 1;
}
static NSUInteger ds4_gpu_bin_threads(uint32_t width, id<MTLComputePipelineState> pipeline) {
NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup;
if (nth_max > 256u) nth_max = 256u;
NSUInteger nth = 1u;
while (2u * nth < (NSUInteger)width && nth < nth_max) nth *= 2u;
return nth ? nth : 1u;
}
static int ds4_gpu_encode_unary_f32_rows(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t width,
uint32_t rows,
int c4,
float min,
float max) {
if (!cb || !pipeline || !src || !dst || width == 0 || rows == 0) return 0;
if (c4 && (width & 3u) != 0) return 0;
ds4_gpu_unary_args args = ds4_gpu_make_unary_rows_args(width, rows, c4, 0.0f, 0.0f);
args.min = min;
args.max = max;
NSUInteger nth_max = pipeline.maxTotalThreadsPerThreadgroup;
if (nth_max > 256u) nth_max = 256u;
NSUInteger nth = (NSUInteger)args.ne00;
if (nth > nth_max) nth = nth_max;
if (nth == 0) nth = 1u;
const NSUInteger nk0 = ((NSUInteger)args.ne00 + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nk0 * (NSUInteger)args.ne01,
(NSUInteger)args.ne02,
(NSUInteger)args.ne03)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_bin_f32_rows(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_bin_args *args,
id<MTLBuffer> a,
NSUInteger a_off,
id<MTLBuffer> b,
NSUInteger b_off,
id<MTLBuffer> out,
NSUInteger out_off) {
if (!cb || !pipeline || !args || !a || !b || !out || args->ne0 <= 0 || args->ne1 <= 0) {
return 0;
}
const NSUInteger nth = ds4_gpu_bin_threads((uint32_t)args->ne0, pipeline);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:a offset:a_off atIndex:1];
[enc setBuffer:b offset:b_off atIndex:2];
[enc setBuffer:out offset:out_off atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)args->ne1,
(NSUInteger)args->ne2,
(NSUInteger)args->ne3)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static ds4_gpu_bin_args ds4_gpu_make_bin_rowwise_scalar_args(uint32_t width, uint32_t rows) {
const uint64_t lhs_row_bytes = (uint64_t)width * sizeof(float);
const uint64_t rhs_row_bytes = sizeof(float);
return (ds4_gpu_bin_args) {
.ne00 = (int32_t)width,
.ne01 = (int32_t)rows,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = lhs_row_bytes,
.nb02 = (uint64_t)rows * lhs_row_bytes,
.nb03 = (uint64_t)rows * lhs_row_bytes,
.ne10 = 1,
.ne11 = (int32_t)rows,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = rhs_row_bytes,
.nb12 = (uint64_t)rows * rhs_row_bytes,
.nb13 = (uint64_t)rows * rhs_row_bytes,
.ne0 = (int32_t)width,
.ne1 = (int32_t)rows,
.ne2 = 1,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = lhs_row_bytes,
.nb2 = (uint64_t)rows * lhs_row_bytes,
.nb3 = (uint64_t)rows * lhs_row_bytes,
.offs = 0,
.o1 = { 0 },
};
}
static ds4_gpu_mul_mv_id_args ds4_gpu_make_mul_mv_id_args(
uint32_t src0_cols,
uint32_t src0_rows,
uint32_t src0_experts,
uint64_t src0_row_bytes,
uint64_t src0_expert_bytes,
uint32_t src1_expert_rows,
uint32_t selected_experts,
uint32_t n_tokens,
uint32_t nr0) {
const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float);
const uint64_t src0_blocks = src0_cols / 256u;
const uint64_t src0_block_bytes = src0_blocks ? src0_row_bytes / src0_blocks : 1u;
return (ds4_gpu_mul_mv_id_args) {
.nei0 = (int32_t)selected_experts,
.nei1 = (int32_t)n_tokens,
.nbi1 = (uint64_t)selected_experts * sizeof(int32_t),
.ne00 = (int32_t)src0_cols,
.ne01 = (int32_t)src0_rows,
.ne02 = (int32_t)src0_experts,
.nb00 = src0_block_bytes,
.nb01 = src0_row_bytes,
.nb02 = src0_expert_bytes,
.ne10 = (int32_t)src0_cols,
.ne11 = (int32_t)src1_expert_rows,
.ne12 = (int32_t)n_tokens,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = src1_row_bytes,
.nb12 = (uint64_t)src1_expert_rows * src1_row_bytes,
.ne0 = (int32_t)src0_rows,
.ne1 = (int32_t)selected_experts,
.nb1 = (uint64_t)src0_rows * sizeof(float),
.nr0 = (int32_t)nr0,
};
}
static ds4_gpu_mul_mm_id_map_args ds4_gpu_make_mul_mm_id_map_args(
uint32_t src0_cols,
uint32_t src0_experts,
uint32_t src1_expert_rows,
uint32_t selected_experts,
uint32_t n_tokens) {
const uint64_t src1_row_bytes = (uint64_t)src0_cols * sizeof(float);
return (ds4_gpu_mul_mm_id_map_args) {
.ne02 = (int32_t)src0_experts,
.ne10 = (int32_t)src0_cols,
.ne11 = (int32_t)src1_expert_rows,
.nb11 = src1_row_bytes,
.nb12 = (uint64_t)src1_expert_rows * src1_row_bytes,
.ne21 = (int32_t)n_tokens,
.ne20 = (int32_t)selected_experts,
.nb21 = (uint64_t)selected_experts * sizeof(int32_t),
};
}
static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args(
uint32_t src0_cols,
uint32_t src0_rows,
uint32_t src0_experts,
uint64_t src0_row_bytes,
uint64_t src0_expert_bytes,
uint32_t src1_expert_rows,
uint32_t selected_experts,
uint32_t n_tokens) {
return ds4_gpu_make_mul_mm_id_args_src1_size(src0_cols,
src0_rows,
src0_experts,
src0_row_bytes,
src0_expert_bytes,
src1_expert_rows,
selected_experts,
n_tokens,
sizeof(float));
}
static ds4_gpu_mul_mm_id_args ds4_gpu_make_mul_mm_id_args_src1_size(
uint32_t src0_cols,
uint32_t src0_rows,
uint32_t src0_experts,
uint64_t src0_row_bytes,
uint64_t src0_expert_bytes,
uint32_t src1_expert_rows,
uint32_t selected_experts,
uint32_t n_tokens,
uint32_t src1_elem_size) {
const uint64_t src1_row_bytes = (uint64_t)src0_cols * src1_elem_size;
return (ds4_gpu_mul_mm_id_args) {
.ne00 = (int32_t)src0_cols,
.ne02 = (int32_t)src0_experts,
.nb01 = src0_row_bytes,
.nb02 = src0_expert_bytes,
.nb03 = (uint64_t)src0_experts * src0_expert_bytes,
.ne11 = (int32_t)src1_expert_rows,
.nb10 = src1_elem_size,
.nb11 = src1_row_bytes,
.nb12 = (uint64_t)src1_expert_rows * src1_row_bytes,
.nb13 = (uint64_t)n_tokens * (uint64_t)src1_expert_rows * src1_row_bytes,
.ne20 = (int32_t)selected_experts,
.ne21 = (int32_t)n_tokens,
.ne0 = (int32_t)src0_rows,
.ne1 = (int32_t)selected_experts,
.r2 = 1,
.r3 = 1,
};
}
static uint32_t ds4_gpu_routed_mv_nr0(uint32_t type) {
switch (type) {
case DS4_METAL_TENSOR_Q8_0: return 2;
case DS4_METAL_TENSOR_Q8_K: return 2;
case DS4_METAL_TENSOR_Q4_K: return 2;
case DS4_METAL_TENSOR_MXFP4: return 2;
case DS4_METAL_TENSOR_Q2_K:
case DS4_METAL_TENSOR_IQ2_XXS: return 4;
default: return 0;
}
}
static const char *ds4_gpu_metal_tensor_type_name(uint32_t type) {
switch (type) {
case DS4_METAL_TENSOR_IQ2_XXS: return "iq2_xxs";
case DS4_METAL_TENSOR_Q2_K: return "q2_k";
case DS4_METAL_TENSOR_Q4_K: return "q4_k";
case DS4_METAL_TENSOR_Q5_K: return "q5_k";
case DS4_METAL_TENSOR_Q6_K: return "q6_k";
case DS4_METAL_TENSOR_MXFP4: return "mxfp4";
default: return "unknown";
}
}
static const char *ds4_gpu_trim_env_value(const char *env, size_t *len_out) {
if (len_out) *len_out = 0;
if (!env) return NULL;
while (isspace((unsigned char)*env)) env++;
size_t n = strlen(env);
while (n > 0 && isspace((unsigned char)env[n - 1])) n--;
if (len_out) *len_out = n;
return env;
}
static bool ds4_gpu_profile_layer_value_match(const char *env, uint32_t layer_index) {
size_t env_len = 0;
env = ds4_gpu_trim_env_value(env, &env_len);
if (!env || env_len == 0) return true;
if (ds4_gpu_env_value_eq(env, env_len, "all")) return true;
const char *p = env;
const char *end_env = env + env_len;
while (p < end_env) {
while (p < end_env && (*p == ' ' || *p == '\t' || *p == ',')) p++;
if (p >= end_env) break;
char *end = NULL;
const unsigned long first = strtoul(p, &end, 10);
if (end == p || end > end_env || first > UINT32_MAX) return false;
unsigned long last = first;
p = end;
if (p < end_env && *p == '-') {
p++;
last = strtoul(p, &end, 10);
if (end == p || end > end_env || last > UINT32_MAX) return false;
p = end;
}
if (first <= layer_index && layer_index <= last) return true;
while (p < end_env && (*p == ' ' || *p == '\t')) p++;
if (p < end_env && *p != ',') return false;
}
return false;
}
static bool ds4_gpu_stage_profile_enabled_for_layer(const char *flag_env_name,
const char *layer_env_name,
uint32_t layer_index) {
size_t flag_len = 0;
const char *flag = ds4_gpu_trim_env_value(getenv(flag_env_name), &flag_len);
if (!flag) return false;
size_t layer_len = 0;
const char *layer = ds4_gpu_trim_env_value(getenv(layer_env_name), &layer_len);
const bool has_layer_filter = layer && layer_len != 0;
if (flag_len != 0) {
if (ds4_gpu_env_value_eq(flag, flag_len, "0") ||
ds4_gpu_env_value_eq(flag, flag_len, "false") ||
ds4_gpu_env_value_eq(flag, flag_len, "no") ||
ds4_gpu_env_value_eq(flag, flag_len, "off")) {
return false;
}
if (!has_layer_filter &&
!ds4_gpu_env_value_eq(flag, flag_len, "1") &&
!ds4_gpu_env_value_eq(flag, flag_len, "true") &&
!ds4_gpu_env_value_eq(flag, flag_len, "yes") &&
!ds4_gpu_env_value_eq(flag, flag_len, "on") &&
!ds4_gpu_env_value_eq(flag, flag_len, "all")) {
return ds4_gpu_profile_layer_value_match(flag, layer_index);
}
}
return ds4_gpu_profile_layer_value_match(layer, layer_index);
}
static NSUInteger ds4_gpu_routed_mv_smem(uint32_t type) {
if (type == DS4_METAL_TENSOR_Q8_0) {
return 32u * 2u * sizeof(float);
}
if (type == DS4_METAL_TENSOR_IQ2_XXS) {
return 256u * sizeof(uint64_t) + 128u * sizeof(uint8_t);
}
if (type == DS4_METAL_TENSOR_MXFP4) {
return 32u * sizeof(float);
}
return 0;
}
static NSUInteger ds4_gpu_routed_mv_nsg(uint32_t type) {
return type == DS4_METAL_TENSOR_Q8_0 ? 4u : 2u;
}
static bool ds4_gpu_routed_mv_rows_per_group_is_nr0(uint32_t type) {
return type == DS4_METAL_TENSOR_Q8_0;
}
static id<MTLComputePipelineState> ds4_gpu_routed_mv_pipeline(uint32_t type) {
switch (type) {
case DS4_METAL_TENSOR_Q8_0:
return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_0_f32", 4);
case DS4_METAL_TENSOR_Q8_K:
return ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_id_q8_K_f32", 2);
case DS4_METAL_TENSOR_IQ2_XXS: return g_moe_mul_mv_id_iq2_xxs_pipeline;
case DS4_METAL_TENSOR_Q2_K: return g_moe_mul_mv_id_q2_k_pipeline;
case DS4_METAL_TENSOR_Q4_K: return g_moe_mul_mv_id_q4_k_pipeline;
case DS4_METAL_TENSOR_MXFP4: return g_moe_mul_mv_id_mxfp4_pipeline;
default: return nil;
}
}
/* TensorOps routed-MoE prefill uses bits 0/1/2 for gate/up/down. Unsupported
* tensor types keep their established kernels. --quality and the global
* Metal4 comparison switch retain the reference path. */
static int ds4_gpu_routed_mm_mpp_mask(void) {
return ds4_gpu_mpp_available() ? 7 : 0;
}
static id<MTLComputePipelineState> ds4_gpu_routed_mm_pipeline(uint32_t type) {
switch (type) {
case DS4_METAL_TENSOR_Q8_0:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f32", false);
case DS4_METAL_TENSOR_Q8_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f32", false);
case DS4_METAL_TENSOR_IQ2_XXS:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f32", false);
case DS4_METAL_TENSOR_Q2_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f32", false);
case DS4_METAL_TENSOR_Q4_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32", false);
case DS4_METAL_TENSOR_Q5_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f32", false);
case DS4_METAL_TENSOR_Q6_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f32", false);
case DS4_METAL_TENSOR_MXFP4:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_mxfp4_f32", false);
default:
return nil;
}
}
static id<MTLComputePipelineState> ds4_gpu_routed_mm_addr_pipeline(uint32_t type) {
switch (type) {
case DS4_METAL_TENSOR_Q2_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f32", false);
case DS4_METAL_TENSOR_Q4_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f32", false);
case DS4_METAL_TENSOR_MXFP4:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_mxfp4_f32", false);
default:
return nil;
}
}
static id<MTLComputePipelineState> ds4_gpu_routed_mm_f16_rhs_pipeline(uint32_t type) {
switch (type) {
case DS4_METAL_TENSOR_Q8_0:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_0_f16", false);
case DS4_METAL_TENSOR_Q8_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q8_K_f16", false);
case DS4_METAL_TENSOR_IQ2_XXS:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f16", false);
case DS4_METAL_TENSOR_Q2_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q2_K_f16", false);
case DS4_METAL_TENSOR_Q4_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f16", false);
case DS4_METAL_TENSOR_Q5_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q5_K_f16", false);
case DS4_METAL_TENSOR_Q6_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q6_K_f16", false);
case DS4_METAL_TENSOR_MXFP4:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_mxfp4_f16", false);
default:
return nil;
}
}
static id<MTLComputePipelineState> ds4_gpu_routed_mm_addr_f16_rhs_pipeline(uint32_t type) {
switch (type) {
case DS4_METAL_TENSOR_Q2_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q2_K_f16", false);
case DS4_METAL_TENSOR_Q4_K:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_q4_K_f16", false);
case DS4_METAL_TENSOR_MXFP4:
return ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_addr_mxfp4_f16", false);
default:
return nil;
}
}
static int ds4_gpu_encode_mul_mv_id(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:src0 offset:src0_off atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:dst offset:dst_off atIndex:3];
[enc setBuffer:ids offset:ids_off atIndex:4];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_attn_out_low_q8_direct(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !src0 || !src1 || !dst ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
return 0;
}
/* Two row conventions in the classic matvec family: Q8_0 k-splits one
* nr0-row group across all nsg simdgroups (cross-simdgroup reduce), so a
* threadgroup covers nr0 rows; Q4_K gives each simdgroup its own nr0
* rows, covering nr0*nsg. Dispatching Q8 with the Q4 stride leaves
* (nsg-1)/nsg of the output rows unwritten. */
const NSUInteger rows_per_group = rows_per_group_is_nr0 ?
(NSUInteger)args->nr0 : (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:src0 offset:src0_off atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:dst offset:dst_off atIndex:3];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_id_pair(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
id<MTLBuffer> src0_a,
NSUInteger src0_a_off,
id<MTLBuffer> src0_b,
NSUInteger src0_b_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !src0_a || !src0_b || !src1 || !dst_a || !dst_b || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:src0_a offset:src0_a_off atIndex:1];
[enc setBuffer:src0_b offset:src0_b_off atIndex:2];
[enc setBuffer:src1 offset:src1_off atIndex:3];
[enc setBuffer:dst_a offset:dst_a_off atIndex:4];
[enc setBuffer:dst_b offset:dst_b_off atIndex:5];
[enc setBuffer:ids offset:ids_off atIndex:6];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_id_pair_swiglu(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
id<MTLBuffer> src0_a,
NSUInteger src0_a_off,
id<MTLBuffer> src0_b,
NSUInteger src0_b_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !act ||
!src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
[enc setBuffer:src0_a offset:src0_a_off atIndex:2];
[enc setBuffer:src0_b offset:src0_b_off atIndex:3];
[enc setBuffer:src1 offset:src1_off atIndex:4];
[enc setBuffer:dst_a offset:dst_a_off atIndex:5];
[enc setBuffer:dst_b offset:dst_b_off atIndex:6];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:7];
[enc setBuffer:ids offset:ids_off atIndex:8];
[enc setBuffer:weights offset:weights_off atIndex:9];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
/* A few rows' pair/SwiGLU with the distinct experts read once: the plan
* dispatch, then one threadgroup per distinct expert with a simdgroup per
* (row, slot) pair (kernel_mul_mv_id_mxfp4_pair_swiglu_dedup_f32). */
static int ds4_gpu_encode_mul_mv_id_pair_swiglu_dedup(
id<MTLCommandBuffer> cb,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
id<MTLBuffer> src0_a, NSUInteger src0_a_off,
id<MTLBuffer> src0_b, NSUInteger src0_b_off,
id<MTLBuffer> src1, NSUInteger src1_off,
id<MTLBuffer> dst_a, NSUInteger dst_a_off,
id<MTLBuffer> dst_b, NSUInteger dst_b_off,
id<MTLBuffer> dst_mid, NSUInteger dst_mid_off,
id<MTLBuffer> ids, NSUInteger ids_off,
id<MTLBuffer> weights, NSUInteger weights_off,
NSUInteger threadgroup_bytes, uint32_t rows) {
static id<MTLComputePipelineState> plan_pipeline, pair_pipeline;
static id<MTLBuffer> plan;
if (!cb || !args || !act || !rows || rows > 8u || args->nei0 <= 0 || args->nei0 > 8 ||
args->nei1 != (int32_t)rows || (uint64_t)args->nei0 * rows > 64u) return 0;
if (!plan_pipeline) plan_pipeline = ds4_gpu_get_pipeline("kernel_dsv41_moe_dedup_plan");
if (!pair_pipeline) pair_pipeline = ds4_gpu_get_pipeline("kernel_mul_mv_id_mxfp4_pair_swiglu_dedup_f32");
if (!plan) plan = [g_device newBufferWithLength:4096 options:MTLResourceStorageModePrivate];
if (!plan_pipeline || !pair_pipeline || !plan) return 0;
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + nr0 - 1u) / nr0;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:plan_pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:ids offset:ids_off atIndex:1];
[enc setBuffer:plan offset:0 atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pair_pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
[enc setBuffer:src0_a offset:src0_a_off atIndex:2];
[enc setBuffer:src0_b offset:src0_b_off atIndex:3];
[enc setBuffer:src1 offset:src1_off atIndex:4];
[enc setBuffer:dst_a offset:dst_a_off atIndex:5];
[enc setBuffer:dst_b offset:dst_b_off atIndex:6];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:7];
[enc setBuffer:ids offset:ids_off atIndex:8];
[enc setBuffer:weights offset:weights_off atIndex:9];
[enc setBuffer:plan offset:0 atIndex:10];
if (threadgroup_bytes != 0) [enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, rows, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
DS4MetalQ4ExpertTable *gate_table,
DS4MetalQ4ExpertTable *up_table,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0,
bool queue_residency) {
if (!cb || !pipeline || !args || !act || !gate_table || !up_table ||
!gate_table.argumentBuffer || !up_table.argumentBuffer ||
!src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
[enc setBuffer:gate_table.argumentBuffer offset:0 atIndex:2];
[enc setBuffer:up_table.argumentBuffer offset:0 atIndex:3];
[enc setBuffer:src1 offset:src1_off atIndex:4];
[enc setBuffer:dst_a offset:dst_a_off atIndex:5];
[enc setBuffer:dst_b offset:dst_b_off atIndex:6];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:7];
[enc setBuffer:ids offset:ids_off atIndex:8];
[enc setBuffer:weights offset:weights_off atIndex:9];
if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) ||
!ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) {
ds4_gpu_end_compute_encoder(cb, enc);
return 0;
}
ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, queue_residency);
ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, queue_residency);
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
DS4MetalQ4ExpertTable *gate_table,
DS4MetalQ4ExpertTable *up_table,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !act || !gate_table || !up_table ||
!gate_table.addressBuffer || !up_table.addressBuffer ||
!src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
[enc setBuffer:gate_table.addressBuffer offset:0 atIndex:2];
[enc setBuffer:up_table.addressBuffer offset:0 atIndex:3];
[enc setBuffer:src1 offset:src1_off atIndex:4];
[enc setBuffer:dst_a offset:dst_a_off atIndex:5];
[enc setBuffer:dst_b offset:dst_b_off atIndex:6];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:7];
[enc setBuffer:ids offset:ids_off atIndex:8];
[enc setBuffer:weights offset:weights_off atIndex:9];
if (!ds4_gpu_bind_q4_expert_table_anchors(enc, gate_table, 10, 6) ||
!ds4_gpu_bind_q4_expert_table_anchors(enc, up_table, 16, 6)) {
ds4_gpu_end_compute_encoder(cb, enc);
return 0;
}
if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) {
ds4_gpu_use_q4_expert_table_resources(cb, enc, gate_table, false);
ds4_gpu_use_q4_expert_table_resources(cb, enc, up_table, false);
}
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_table_q4_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
DS4MetalQ4ExpertTable *table,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool queue_residency) {
if (!cb || !pipeline || !args || !table || !table.argumentBuffer ||
!src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:table.argumentBuffer offset:0 atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:dst offset:dst_off atIndex:3];
[enc setBuffer:ids offset:ids_off atIndex:4];
if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) {
ds4_gpu_end_compute_encoder(cb, enc);
return 0;
}
ds4_gpu_use_q4_expert_table_resources(cb, enc, table, queue_residency);
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_addr_q4_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
DS4MetalQ4ExpertTable *table,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg) {
if (!cb || !pipeline || !args || !table || !table.addressBuffer ||
!src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:table.addressBuffer offset:0 atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:dst offset:dst_off atIndex:3];
[enc setBuffer:ids offset:ids_off atIndex:4];
if (!ds4_gpu_bind_q4_expert_table_anchors(enc, table, 5, 6)) {
ds4_gpu_end_compute_encoder(cb, enc);
return 0;
}
if (getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL) {
ds4_gpu_use_q4_expert_table_resources(cb, enc, table, false);
}
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static uint32_t ds4_gpu_q4_expert_group_size(uint32_t n_total_expert) {
uint32_t group_size = 32;
const char *env = getenv("DS4_METAL_Q4_EXPERT_GROUP_SIZE");
if (env && env[0]) {
char *end = NULL;
unsigned long v = strtoul(env, &end, 10);
if (end != env && *end == '\0' && v > 0 && v <= UINT32_MAX) {
group_size = (uint32_t)v;
}
}
if (group_size == 0) group_size = 1;
if (group_size > n_total_expert) group_size = n_total_expert;
return group_size;
}
static int ds4_gpu_encode_mul_mv_group_q4_pair_swiglu(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
const ds4_gpu_moe_expert_group_args *group,
id<MTLBuffer> src0_a,
NSUInteger src0_a_off,
id<MTLBuffer> src0_b,
NSUInteger src0_b_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !act || !group ||
!src0_a || !src0_b || !src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 <= 0 || args->nei1 <= 0 ||
group->expert_count == 0) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
[enc setBytes:group length:sizeof(*group) atIndex:2];
[enc setBuffer:src0_a offset:src0_a_off atIndex:3];
[enc setBuffer:src0_b offset:src0_b_off atIndex:4];
[enc setBuffer:src1 offset:src1_off atIndex:5];
[enc setBuffer:dst_a offset:dst_a_off atIndex:6];
[enc setBuffer:dst_b offset:dst_b_off atIndex:7];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:8];
[enc setBuffer:ids offset:ids_off atIndex:9];
[enc setBuffer:weights offset:weights_off atIndex:10];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_group_q4_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_moe_expert_group_args *group,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg) {
if (!cb || !pipeline || !args || !group || !src0 || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
group->expert_count == 0) {
return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:group length:sizeof(*group) atIndex:1];
[enc setBuffer:src0 offset:src0_off atIndex:2];
[enc setBuffer:src1 offset:src1_off atIndex:3];
[enc setBuffer:dst offset:dst_off atIndex:4];
[enc setBuffer:ids offset:ids_off atIndex:5];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_id_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> add_in,
NSUInteger add_in_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg) {
if (!cb || !pipeline || !args || !src0 || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 ||
args->nei0 <= 0 || args->nei0 > 8 || args->nei1 <= 0) {
return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:src0 offset:src0_off atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:dst offset:dst_off atIndex:3];
[enc setBuffer:ids offset:ids_off atIndex:4];
[enc setBuffer:(add_in ? add_in : dst) offset:(add_in ? add_in_off : dst_off) atIndex:5];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_q4_gather_slots6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_q4_gather_slots6_args *args,
__unsafe_unretained id<MTLBuffer> src_groups[6],
const NSUInteger src_group_offsets[6],
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> dst,
NSUInteger dst_off) {
if (!cb || !pipeline || !args || !src_groups || !src_group_offsets || !ids || !dst ||
args->expert_bytes == 0 || (args->expert_bytes & 15u) != 0 ||
args->group_size == 0 || args->n_slots == 0 || args->n_slots > 6) {
return 0;
}
for (uint32_t i = 0; i < 6; i++) {
if (!src_groups[i]) return 0;
}
const uint64_t chunks = args->expert_bytes >> 4;
if (chunks == 0 || chunks > NSUIntegerMax) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
for (uint32_t i = 0; i < 6; i++) {
[enc setBuffer:src_groups[i] offset:src_group_offsets[i] atIndex:1 + i];
}
[enc setBuffer:ids offset:ids_off atIndex:7];
[enc setBuffer:dst offset:dst_off atIndex:8];
const NSUInteger threads = 256u;
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)chunks + threads - 1u) / threads,
(NSUInteger)args->n_slots,
1)
threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_slots6_pair_swiglu(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
__unsafe_unretained id<MTLBuffer> src0_a[6],
const NSUInteger src0_a_off[6],
__unsafe_unretained id<MTLBuffer> src0_b[6],
const NSUInteger src0_b_off[6],
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off ||
!src1 || !dst_a || !dst_b || !dst_mid || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) {
return 0;
}
for (uint32_t i = 0; i < 6; i++) {
if (!src0_a[i] || !src0_b[i]) return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
for (uint32_t i = 0; i < 6; i++) {
[enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i];
}
for (uint32_t i = 0; i < 6; i++) {
[enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i];
}
[enc setBuffer:src1 offset:src1_off atIndex:14];
[enc setBuffer:dst_a offset:dst_a_off atIndex:15];
[enc setBuffer:dst_b offset:dst_b_off atIndex:16];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:17];
[enc setBuffer:weights offset:weights_off atIndex:18];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_slots6_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
__unsafe_unretained id<MTLBuffer> src0[6],
const NSUInteger src0_off[6],
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg) {
if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0) {
return 0;
}
for (uint32_t i = 0; i < 6; i++) {
if (!src0[i]) return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
for (uint32_t i = 0; i < 6; i++) {
[enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
}
[enc setBuffer:src1 offset:src1_off atIndex:7];
[enc setBuffer:dst offset:dst_off atIndex:8];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_group6_pair_swiglu(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
__unsafe_unretained id<MTLBuffer> src0_a[6],
const NSUInteger src0_a_off[6],
__unsafe_unretained id<MTLBuffer> src0_b[6],
const NSUInteger src0_b_off[6],
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off ||
!src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 != 384) {
return 0;
}
for (uint32_t i = 0; i < 6; i++) {
if (!src0_a[i] || !src0_b[i]) return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
for (uint32_t i = 0; i < 6; i++) {
[enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i];
}
for (uint32_t i = 0; i < 6; i++) {
[enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:8 + i];
}
[enc setBuffer:src1 offset:src1_off atIndex:14];
[enc setBuffer:dst_a offset:dst_a_off atIndex:15];
[enc setBuffer:dst_b offset:dst_b_off atIndex:16];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:17];
[enc setBuffer:ids offset:ids_off atIndex:18];
[enc setBuffer:weights offset:weights_off atIndex:19];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_group6_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
__unsafe_unretained id<MTLBuffer> src0[6],
const NSUInteger src0_off[6],
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg) {
if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 != 384) {
return 0;
}
for (uint32_t i = 0; i < 6; i++) {
if (!src0[i]) return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
for (uint32_t i = 0; i < 6; i++) {
[enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
}
[enc setBuffer:src1 offset:src1_off atIndex:7];
[enc setBuffer:dst offset:dst_off atIndex:8];
[enc setBuffer:ids offset:ids_off atIndex:9];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
ds4_gpu_stream_expert_cache_entry * const *entries,
uint32_t n_entries,
id<MTLBuffer> gate_addrs,
id<MTLBuffer> up_addrs,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0,
id<MTLBuffer> overflow_gate,
id<MTLBuffer> overflow_up) {
if (!cb || !pipeline || !args || !act || !entries ||
(n_entries == 0 && !overflow_gate) ||
!gate_addrs || !up_addrs ||
!src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 ||
args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
for (uint32_t i = 0; i < n_entries; i++) {
if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0;
}
if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
n_entries,
0)) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
[enc setBuffer:gate_addrs offset:0 atIndex:2];
[enc setBuffer:up_addrs offset:0 atIndex:3];
[enc setBuffer:src1 offset:src1_off atIndex:4];
[enc setBuffer:dst_a offset:dst_a_off atIndex:5];
[enc setBuffer:dst_b offset:dst_b_off atIndex:6];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:7];
[enc setBuffer:ids offset:ids_off atIndex:8];
[enc setBuffer:weights offset:weights_off atIndex:9];
for (uint32_t i = 0; i < n_entries; i++) {
[enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead];
[enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead];
}
/* Overflow experts are addressed straight into the mapped model views
* when a layer's unique selected set exceeds the cache budget. */
if (overflow_gate) [enc useResource:overflow_gate usage:MTLResourceUsageRead];
if (overflow_up) [enc useResource:overflow_up usage:MTLResourceUsageRead];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_addr_iq2(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
ds4_gpu_stream_expert_cache_entry * const *entries,
uint32_t n_entries,
id<MTLBuffer> addrs,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0,
uint32_t resource_kind,
id<MTLBuffer> overflow_resource) {
if (!cb || !pipeline || !args || !entries ||
(n_entries == 0 && !overflow_resource) ||
n_entries > DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT || resource_kind > 2 ||
!addrs || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 ||
args->nei0 <= 0 || args->nei0 > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
for (uint32_t i = 0; i < n_entries; i++) {
if (!entries[i] || !entries[i]->gate_buffer ||
!entries[i]->up_buffer || !entries[i]->down_buffer) return 0;
}
if (n_entries && !ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
n_entries,
0)) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:addrs offset:0 atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:dst offset:dst_off atIndex:3];
[enc setBuffer:ids offset:ids_off atIndex:4];
for (uint32_t i = 0; i < n_entries; i++) {
id<MTLBuffer> b = resource_kind == 0 ? entries[i]->gate_buffer :
resource_kind == 1 ? entries[i]->up_buffer :
entries[i]->down_buffer;
[enc useResource:b usage:MTLResourceUsageRead];
}
if (overflow_resource) [enc useResource:overflow_resource usage:MTLResourceUsageRead];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_addr_q2_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
ds4_gpu_stream_expert_cache_entry * const *entries,
uint32_t n_entries,
id<MTLBuffer> addrs,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
id<MTLBuffer> overflow_down) {
if (!cb || !pipeline || !args || !entries ||
(n_entries == 0 && !overflow_down) ||
!addrs || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
for (uint32_t i = 0; i < n_entries; i++) {
if (!entries[i] || !entries[i]->down_buffer) return 0;
}
if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
n_entries,
0)) {
return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:addrs offset:0 atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:dst offset:dst_off atIndex:3];
[enc setBuffer:ids offset:ids_off atIndex:4];
for (uint32_t i = 0; i < n_entries; i++) {
[enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead];
}
if (overflow_down) [enc useResource:overflow_down usage:MTLResourceUsageRead];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
const ds4_gpu_stream_expert_split_args *split,
ds4_gpu_stream_expert_cache_entry * const entries[6],
id<MTLBuffer> gate_addrs,
id<MTLBuffer> up_addrs,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !act || !split || !entries ||
!gate_addrs || !up_addrs || !src1 || !dst_a || !dst_b || !dst_mid ||
!ids || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
for (uint32_t i = 0; i < 6; i++) {
if ((split->active_mask & (1u << i)) == 0) continue;
if (!entries[i] || !entries[i]->gate_buffer || !entries[i]->up_buffer) return 0;
}
if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
6,
split->active_mask)) {
return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
[enc setBytes:split length:sizeof(*split) atIndex:2];
[enc setBuffer:gate_addrs offset:0 atIndex:3];
[enc setBuffer:up_addrs offset:0 atIndex:4];
[enc setBuffer:src1 offset:src1_off atIndex:5];
[enc setBuffer:dst_a offset:dst_a_off atIndex:6];
[enc setBuffer:dst_b offset:dst_b_off atIndex:7];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:8];
[enc setBuffer:ids offset:ids_off atIndex:9];
[enc setBuffer:weights offset:weights_off atIndex:10];
for (uint32_t i = 0; i < 6; i++) {
if ((split->active_mask & (1u << i)) == 0) continue;
[enc useResource:entries[i]->gate_buffer usage:MTLResourceUsageRead];
[enc useResource:entries[i]->up_buffer usage:MTLResourceUsageRead];
}
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_stream_expert_split_args *split,
ds4_gpu_stream_expert_cache_entry * const entries[6],
id<MTLBuffer> addrs,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg) {
if (!cb || !pipeline || !args || !split || !entries || !addrs ||
!src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 <= 0 || args->ne02 > 384) {
return 0;
}
for (uint32_t i = 0; i < 6; i++) {
if ((split->active_mask & (1u << i)) == 0) continue;
if (!entries[i] || !entries[i]->down_buffer) return 0;
}
if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(entries,
6,
split->active_mask)) {
return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:split length:sizeof(*split) atIndex:1];
[enc setBuffer:addrs offset:0 atIndex:2];
[enc setBuffer:src1 offset:src1_off atIndex:3];
[enc setBuffer:dst offset:dst_off atIndex:4];
[enc setBuffer:ids offset:ids_off atIndex:5];
for (uint32_t i = 0; i < 6; i++) {
if ((split->active_mask & (1u << i)) == 0) continue;
[enc useResource:entries[i]->down_buffer usage:MTLResourceUsageRead];
}
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_group8_pair_swiglu(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act,
__unsafe_unretained id<MTLBuffer> src0_a[8],
const NSUInteger src0_a_off[8],
__unsafe_unretained id<MTLBuffer> src0_b[8],
const NSUInteger src0_b_off[8],
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst_a,
NSUInteger dst_a_off,
id<MTLBuffer> dst_b,
NSUInteger dst_b_off,
id<MTLBuffer> dst_mid,
NSUInteger dst_mid_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !act || !src0_a || !src0_a_off || !src0_b || !src0_b_off ||
!src1 || !dst_a || !dst_b || !dst_mid || !ids || !weights ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 != 384) {
return 0;
}
for (uint32_t i = 0; i < 8; i++) {
if (!src0_a[i] || !src0_b[i]) return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBytes:act length:sizeof(*act) atIndex:1];
for (uint32_t i = 0; i < 8; i++) {
[enc setBuffer:src0_a[i] offset:src0_a_off[i] atIndex:2 + i];
}
for (uint32_t i = 0; i < 8; i++) {
[enc setBuffer:src0_b[i] offset:src0_b_off[i] atIndex:10 + i];
}
[enc setBuffer:src1 offset:src1_off atIndex:18];
[enc setBuffer:dst_a offset:dst_a_off atIndex:19];
[enc setBuffer:dst_b offset:dst_b_off atIndex:20];
[enc setBuffer:dst_mid offset:dst_mid_off atIndex:21];
[enc setBuffer:ids offset:ids_off atIndex:22];
[enc setBuffer:weights offset:weights_off atIndex:23];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_group8_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
__unsafe_unretained id<MTLBuffer> src0[8],
const NSUInteger src0_off[8],
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg) {
if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 != 384) {
return 0;
}
for (uint32_t i = 0; i < 8; i++) {
if (!src0[i]) return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
for (uint32_t i = 0; i < 8; i++) {
[enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
}
[enc setBuffer:src1 offset:src1_off atIndex:9];
[enc setBuffer:dst offset:dst_off atIndex:10];
[enc setBuffer:ids offset:ids_off atIndex:11];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_group24_id(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
__unsafe_unretained id<MTLBuffer> src0[24],
const NSUInteger src0_off[24],
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg,
bool rows_per_group_is_nr0) {
if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 != 384) {
return 0;
}
for (uint32_t i = 0; i < 24; i++) {
if (!src0[i]) return 0;
}
const NSUInteger nr0 = (NSUInteger)args->nr0;
const NSUInteger rows_per_group = rows_per_group_is_nr0 ? nr0 : nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
const NSUInteger pairs = (NSUInteger)args->nei0 * (NSUInteger)args->nei1;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
for (uint32_t i = 0; i < 24; i++) {
[enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
}
[enc setBuffer:src1 offset:src1_off atIndex:25];
[enc setBuffer:dst offset:dst_off atIndex:26];
[enc setBuffer:ids offset:ids_off atIndex:27];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, 1, pairs)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mv_group24_sum6(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mv_id_args *args,
__unsafe_unretained id<MTLBuffer> src0[24],
const NSUInteger src0_off[24],
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off,
NSUInteger threadgroup_bytes,
NSUInteger nsg) {
if (!cb || !pipeline || !args || !src0 || !src0_off || !src1 || !dst || !ids ||
args->ne00 <= 0 || args->ne01 <= 0 || args->nei0 != 6 || args->nei1 <= 0 ||
args->ne02 != 384) {
return 0;
}
for (uint32_t i = 0; i < 24; i++) {
if (!src0[i]) return 0;
}
const NSUInteger rows_per_group = (NSUInteger)args->nr0 * nsg;
const NSUInteger row_groups = ((NSUInteger)args->ne01 + rows_per_group - 1u) / rows_per_group;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
for (uint32_t i = 0; i < 24; i++) {
[enc setBuffer:src0[i] offset:src0_off[i] atIndex:1 + i];
}
[enc setBuffer:src1 offset:src1_off atIndex:25];
[enc setBuffer:dst offset:dst_off atIndex:26];
[enc setBuffer:ids offset:ids_off atIndex:27];
if (threadgroup_bytes != 0) {
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(row_groups, (NSUInteger)args->nei1, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mm_id(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> map_pipeline,
id<MTLComputePipelineState> mm_pipeline,
const ds4_gpu_mul_mm_id_map_args *map_args,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
id<MTLBuffer> ids,
NSUInteger ids_off) {
if (!cb || !map_pipeline || !mm_pipeline || !map_args || !mm_args ||
!src0 || !src1 || !dst || !ids ||
mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
return 0;
}
return ds4_gpu_encode_mul_mm_id_map(cb,
map_pipeline,
map_args,
mm_args,
ids,
ids_off) &&
ds4_gpu_encode_mul_mm_id_mapped(cb,
mm_pipeline,
mm_args,
src0,
src0_off,
src1,
src1_off,
dst,
dst_off);
}
static int ds4_gpu_encode_mul_mm_id_map(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> map_pipeline,
const ds4_gpu_mul_mm_id_map_args *map_args,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> ids,
NSUInteger ids_off) {
if (!cb || !map_pipeline || !map_args || !mm_args || !ids ||
mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
return 0;
}
const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t);
const NSUInteger hids_bytes =
(NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t);
if (tpe_bytes > NSUIntegerMax - hids_bytes) return 0;
const NSUInteger work_offset = (tpe_bytes + hids_bytes + 7u) & ~7u;
const uint64_t pair_rows =
(uint64_t)(uint32_t)mm_args->ne20 * (uint32_t)mm_args->ne21;
const uint64_t work_cap =
(pair_rows + 31u * (uint32_t)mm_args->ne02 + 31u) / 32u;
const NSUInteger work_item_bytes = 2u * sizeof(uint32_t);
if (work_cap > (NSUIntegerMax - 8u) / work_item_bytes ||
work_offset > NSUIntegerMax - 8u -
(NSUInteger)work_cap * work_item_bytes) {
return 0;
}
const NSUInteger total_bytes =
work_offset + 8u + (NSUInteger)work_cap * work_item_bytes;
if (!ds4_gpu_ensure_scratch_buffer(&g_moe_id_map_buffer,
&g_moe_id_map_bytes,
total_bytes,
"ds4_moe_id_map")) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:map_pipeline];
[enc setBytes:map_args length:sizeof(*map_args) atIndex:0];
[enc setBuffer:ids offset:ids_off atIndex:1];
[enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:2];
[enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:3];
[enc setBuffer:g_moe_id_map_buffer offset:work_offset atIndex:4];
[enc setThreadgroupMemoryLength:(NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne20 * sizeof(uint16_t) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake((NSUInteger)mm_args->ne02, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_moe_packed_rhs(
id<MTLCommandBuffer> cb,
const ds4_gpu_mul_mm_id_args *args,
id<MTLBuffer> src, NSUInteger offset, bool f16, NSUInteger limit) {
if (!cb || !args || !src || !g_moe_id_map_buffer ||
args->ne00 <= 0 || args->ne20 <= 0 || args->ne21 <= 0 || args->ne02 <= 0)
return 0;
const uint64_t rows = (uint64_t)args->ne20*args->ne21;
const uint64_t tiles = (rows + 31u*(uint64_t)args->ne02 + 31u)/32u;
const uint64_t row_bytes = (uint64_t)args->ne00*sizeof(uint16_t);
const NSUInteger tile_n = 32u;
const uint64_t header = ((uint64_t)args->ne02 + 31u) / 32u * 32u * sizeof(uint32_t);
if (rows > (NSUIntegerMax - header)/row_bytes) return 0;
const NSUInteger bytes = (NSUInteger)(header + rows*row_bytes);
if (bytes > limit ||
!ds4_gpu_ensure_scratch_buffer(&g_moe_packed_rhs_buffer,
&g_moe_packed_rhs_bytes, bytes, "ds4_moe_packed_rhs")) return 0;
const NSUInteger counts_bytes = (NSUInteger)args->ne02*sizeof(uint32_t);
const NSUInteger ids_bytes = (NSUInteger)args->ne02*args->ne21*sizeof(int32_t);
const NSUInteger work_offset = (counts_bytes + ids_bytes + 7u) & ~7u;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
f16 ? "kernel_moe_pack_rhs_f16" : "kernel_moe_pack_rhs_f32");
id<MTLComputePipelineState> offsets = ds4_gpu_get_pipeline("kernel_moe_packed_offsets");
if (!pipeline || !offsets) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
const uint32_t experts = (uint32_t)args->ne02;
[enc setComputePipelineState:offsets];
[enc setBytes:&experts length:sizeof(experts) atIndex:0];
[enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:1];
[enc setBuffer:g_moe_packed_rhs_buffer offset:0 atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(*args) atIndex:0];
[enc setBuffer:src offset:offset atIndex:1];
[enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:2];
[enc setBuffer:g_moe_id_map_buffer offset:counts_bytes atIndex:3];
[enc setBuffer:g_moe_id_map_buffer offset:work_offset atIndex:4];
[enc setBuffer:g_moe_packed_rhs_buffer offset:0 atIndex:5];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)tiles*tile_n, 1, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mm_id_mapped_tile_resources(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> mm_pipeline,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
NSUInteger threadgroup_bytes,
ds4_gpu_stream_expert_cache_entry * const *resources,
uint32_t resource_count,
uint32_t resource_kind,
id<MTLBuffer> overflow_resource,
uint32_t tile_rows) {
if (!cb || !mm_pipeline || !mm_args || !src0 || !src1 || !dst ||
!g_moe_id_map_buffer ||
mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0 ||
(tile_rows != 32u && tile_rows != 64u)) {
return 0;
}
/* The simdgroup and TensorOps kernels share the 32-token work map. */
const bool use_resource_hints =
getenv("DS4_METAL_MOE_MM_ID_USE_RESOURCES") != NULL &&
getenv("DS4_METAL_DISABLE_MOE_MM_ID_USE_RESOURCES") == NULL;
const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t);
const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t);
if (tpe_bytes > NSUIntegerMax - hids_bytes) {
return 0;
}
const NSUInteger work_offset = (tpe_bytes + hids_bytes + 7u) & ~7u;
const uint64_t pair_rows =
(uint64_t)(uint32_t)mm_args->ne20 * (uint32_t)mm_args->ne21;
const uint64_t work_cap =
(pair_rows + 31u * (uint32_t)mm_args->ne02 + 31u) / 32u;
const NSUInteger work_item_bytes = 2u * sizeof(uint32_t);
if (work_cap > NSUIntegerMax ||
work_offset > NSUIntegerMax - 8u ||
(NSUInteger)work_cap >
(NSUIntegerMax - work_offset - 8u) / work_item_bytes ||
g_moe_id_map_bytes <
work_offset + 8u + (NSUInteger)work_cap * work_item_bytes) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:mm_pipeline];
[enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0];
[enc setBuffer:src0 offset:src0_off atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3];
[enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4];
[enc setBuffer:dst offset:dst_off atIndex:5];
[enc setBuffer:g_moe_id_map_buffer offset:work_offset atIndex:6];
for (uint32_t i = 0; resources && i < resource_count; i++) {
ds4_gpu_stream_expert_cache_entry *entry = resources[i];
id<MTLBuffer> b = resource_kind == 0 ? entry->gate_buffer :
resource_kind == 1 ? entry->up_buffer : entry->down_buffer;
[enc useResource:b usage:MTLResourceUsageRead];
}
if (overflow_resource) [enc useResource:overflow_resource usage:MTLResourceUsageRead];
if (use_resource_hints) {
[enc useResource:src0 usage:MTLResourceUsageRead];
}
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)work_cap,
((NSUInteger)mm_args->ne0 + tile_rows - 1u) / tile_rows,
1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mm_id_mapped_tile(
id<MTLCommandBuffer> cb, id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mm_id_args *args,
id<MTLBuffer> src0, NSUInteger src0_off,
id<MTLBuffer> src1, NSUInteger src1_off,
id<MTLBuffer> dst, NSUInteger dst_off, NSUInteger threadgroup_bytes) {
return ds4_gpu_encode_mul_mm_id_mapped_tile_resources(
cb, pipeline, args, src0, src0_off, src1, src1_off, dst, dst_off,
threadgroup_bytes, NULL, 0, 0, nil, 64);
}
static int ds4_gpu_encode_mul_mm_id_addr_mapped_tile(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> mm_pipeline,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0_addrs,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
NSUInteger threadgroup_bytes,
ds4_gpu_stream_expert_cache_entry * const *resources,
uint32_t resource_count,
uint32_t resource_kind,
id<MTLBuffer> overflow_resource) {
if (!cb || !mm_pipeline || !mm_args || !src0_addrs || !src1 || !dst ||
!g_moe_id_map_buffer ||
mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
return 0;
}
const NSUInteger tile_n = 32u;
const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t);
const NSUInteger hids_bytes =
(NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t);
if (tpe_bytes > NSUIntegerMax - hids_bytes ||
g_moe_id_map_bytes < tpe_bytes + hids_bytes) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:mm_pipeline];
[enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0];
[enc setBuffer:src0_addrs offset:0 atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:3];
[enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:4];
[enc setBuffer:dst offset:dst_off atIndex:5];
[enc useResource:src0_addrs usage:MTLResourceUsageRead];
for (uint32_t i = 0; resources && i < resource_count; i++) {
ds4_gpu_stream_expert_cache_entry *entry = resources[i];
if (!entry) continue;
id<MTLBuffer> b =
resource_kind == 0 ? entry->gate_buffer :
resource_kind == 1 ? entry->up_buffer :
entry->down_buffer;
if (b) [enc useResource:b usage:MTLResourceUsageRead];
}
if (overflow_resource) {
[enc useResource:overflow_resource usage:MTLResourceUsageRead];
}
[enc setThreadgroupMemoryLength:threadgroup_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + tile_n - 1u) / tile_n,
((NSUInteger)mm_args->ne0 + 63u) / 64u,
(NSUInteger)mm_args->ne02)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
bool compact_tile,
const ds4_gpu_mul_mm_id_args *mm_args,
const ds4_gpu_dsv4_moe_swiglu_weight_args *act_args,
id<MTLBuffer> gate_src0,
NSUInteger gate_src0_off,
id<MTLBuffer> up_src0,
NSUInteger up_src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> mid,
NSUInteger mid_off,
id<MTLBuffer> weights,
NSUInteger weights_off) {
if (!cb || !pipeline || !mm_args || !act_args ||
!gate_src0 || !up_src0 || !src1 || !mid || !weights ||
!g_moe_id_map_buffer ||
mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
mm_args->ne20 <= 0 || mm_args->ne21 <= 0 || mm_args->ne02 <= 0) {
return 0;
}
const NSUInteger tpe_bytes = (NSUInteger)mm_args->ne02 * sizeof(int32_t);
const NSUInteger hids_bytes = (NSUInteger)mm_args->ne02 * (NSUInteger)mm_args->ne21 * sizeof(int32_t);
if (tpe_bytes > NSUIntegerMax - hids_bytes) {
return 0;
}
const NSUInteger work_offset = (tpe_bytes + hids_bytes + 7u) & ~7u;
const uint64_t pair_rows =
(uint64_t)(uint32_t)mm_args->ne20 * (uint32_t)mm_args->ne21;
const uint64_t work_cap =
(pair_rows + 31u * (uint32_t)mm_args->ne02 + 31u) / 32u;
const NSUInteger work_item_bytes = 2u * sizeof(uint32_t);
if (work_cap > NSUIntegerMax ||
work_offset > NSUIntegerMax - 8u ||
(NSUInteger)work_cap >
(NSUIntegerMax - work_offset - 8u) / work_item_bytes ||
g_moe_id_map_bytes <
work_offset + 8u + (NSUInteger)work_cap * work_item_bytes) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0];
[enc setBytes:act_args length:sizeof(*act_args) atIndex:1];
[enc setBuffer:gate_src0 offset:gate_src0_off atIndex:2];
[enc setBuffer:up_src0 offset:up_src0_off atIndex:3];
[enc setBuffer:src1 offset:src1_off atIndex:4];
[enc setBuffer:g_moe_id_map_buffer offset:0 atIndex:5];
[enc setBuffer:g_moe_id_map_buffer offset:tpe_bytes atIndex:6];
[enc setBuffer:mid offset:mid_off atIndex:7];
[enc setBuffer:weights offset:weights_off atIndex:8];
[enc setBuffer:g_moe_id_map_buffer offset:work_offset atIndex:9];
const NSUInteger tile_m = compact_tile ? 32u : 64u;
[enc setThreadgroupMemoryLength:compact_tile ? 8192u : 16384u atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)work_cap,
((NSUInteger)mm_args->ne0 + tile_m - 1u) / tile_m,
1)
threadsPerThreadgroup:MTLSizeMake(compact_tile ? 64u : 128u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_mul_mm_id_mapped(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> mm_pipeline,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off) {
return ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
mm_pipeline,
mm_args,
src0,
src0_off,
src1,
src1_off,
dst,
dst_off,
8192u);
}
static int ds4_gpu_encode_attn_out_low_mpp(
id<MTLCommandBuffer> cb,
id<MTLComputePipelineState> pipeline,
const ds4_gpu_mul_mm_id_args *mm_args,
id<MTLBuffer> src0,
NSUInteger src0_off,
id<MTLBuffer> src1,
NSUInteger src1_off,
id<MTLBuffer> dst,
NSUInteger dst_off) {
if (!cb || !pipeline || !mm_args || !src0 || !src1 || !dst ||
mm_args->ne00 <= 0 || mm_args->ne0 <= 0 ||
mm_args->ne02 <= 0 || mm_args->ne1 <= 0 || mm_args->ne21 <= 0) {
return 0;
}
const uint32_t tile_n = DS4_METAL_ATTN_OUT_MPP_TILE_N;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:mm_args length:sizeof(*mm_args) atIndex:0];
[enc setBuffer:src0 offset:src0_off atIndex:1];
[enc setBuffer:src1 offset:src1_off atIndex:2];
[enc setBuffer:dst offset:dst_off atIndex:3];
[enc setThreadgroupMemoryLength:8192u atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)mm_args->ne21 + (NSUInteger)tile_n - 1u) / (NSUInteger)tile_n,
((NSUInteger)mm_args->ne0 + 63u) / 64u,
(NSUInteger)mm_args->ne02)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_swiglu_flat(
id<MTLCommandBuffer> cb,
id<MTLBuffer> gate,
NSUInteger gate_off,
id<MTLBuffer> up,
NSUInteger up_off,
id<MTLBuffer> out,
NSUInteger out_off,
uint32_t n) {
if (!cb || !gate || !up || !out || n == 0) return 0;
ds4_gpu_glu_args args = {
.ne00 = (int32_t)n,
.nb01 = (uint64_t)n * sizeof(float),
.ne10 = (int32_t)n,
.nb11 = (uint64_t)n * sizeof(float),
.ne0 = (int32_t)n,
.nb1 = (uint64_t)n * sizeof(float),
.i00 = 0,
.i10 = 0,
.alpha = 1.0f,
.limit = 0.0f,
};
NSUInteger nth = g_swiglu_flat_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > (NSUInteger)n) nth = (NSUInteger)n;
if (nth == 0u) nth = 1u;
const NSUInteger groups = ((NSUInteger)n + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_swiglu_flat_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:gate offset:gate_off atIndex:1];
[enc setBuffer:up offset:up_off atIndex:2];
[enc setBuffer:out offset:out_off atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_moe_swiglu_weight(
id<MTLCommandBuffer> cb,
id<MTLBuffer> gate,
NSUInteger gate_off,
id<MTLBuffer> up,
NSUInteger up_off,
id<MTLBuffer> mid,
NSUInteger mid_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
uint32_t width,
uint32_t rows,
float clamp_value,
bool mid_f16) {
if (!cb || !gate || !up || !mid || !weights || width == 0 || rows == 0) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline(mid_f16 ? "kernel_dsv4_moe_swiglu_weight_f16" :
"kernel_dsv4_moe_swiglu_weight");
if (!pipeline) return 0;
ds4_gpu_dsv4_moe_swiglu_weight_args args = {
.width = width,
.rows = rows,
.gate_row_stride = (uint64_t)width * sizeof(float),
.up_row_stride = (uint64_t)width * sizeof(float),
.mid_row_stride = (uint64_t)width * (mid_f16 ? sizeof(uint16_t) : sizeof(float)),
.weight_stride = sizeof(float),
.write_clamped = getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL ? 1u : 0u,
.clamp_value = clamp_value,
};
NSUInteger nth = pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > width) nth = width;
if (nth == 0) nth = 1u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:gate offset:gate_off atIndex:1];
[enc setBuffer:up offset:up_off atIndex:2];
[enc setBuffer:mid offset:mid_off atIndex:3];
[enc setBuffer:weights offset:weights_off atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_moe_sum6(
id<MTLCommandBuffer> cb,
id<MTLBuffer> experts,
NSUInteger experts_off,
id<MTLBuffer> out,
NSUInteger out_off,
uint32_t out_dim,
uint32_t n_tokens) {
if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0;
if (!g_moe_sum6_pipeline) return 0;
const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float);
ds4_gpu_dsv4_moe_sum6_args args = {
.width = out_dim,
.tokens = n_tokens,
.src_token_stride = 6u * out_row_bytes,
.dst_token_stride = out_row_bytes,
};
NSUInteger nth = g_moe_sum6_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > out_dim) nth = out_dim;
if (nth == 0) nth = 1u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_moe_sum6_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:experts offset:experts_off atIndex:1];
[enc setBuffer:out offset:out_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_moe_sum8(
id<MTLCommandBuffer> cb,
id<MTLBuffer> experts,
NSUInteger experts_off,
id<MTLBuffer> out,
NSUInteger out_off,
uint32_t out_dim,
uint32_t n_tokens) {
if (!cb || !experts || !out || out_dim == 0 || n_tokens == 0) return 0;
if (!g_moe_sum8_pipeline) return 0;
const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float);
ds4_gpu_dsv4_moe_sum6_args args = {
.width = out_dim,
.tokens = n_tokens,
.src_token_stride = 8u * out_row_bytes,
.dst_token_stride = out_row_bytes,
};
NSUInteger nth = g_moe_sum8_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > out_dim) nth = out_dim;
if (nth == 0) nth = 1u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_moe_sum8_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:experts offset:experts_off atIndex:1];
[enc setBuffer:out offset:out_off atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static ds4_gpu_bin_args ds4_gpu_make_moe_add_args(
uint32_t out_dim,
uint32_t n_tokens,
uint64_t src0_token_stride,
uint64_t src1_token_stride,
uint64_t dst_token_stride) {
return (ds4_gpu_bin_args) {
.ne00 = (int32_t)out_dim,
.ne01 = (int32_t)n_tokens,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = src0_token_stride,
.nb02 = (uint64_t)n_tokens * src0_token_stride,
.nb03 = (uint64_t)n_tokens * src0_token_stride,
.ne10 = (int32_t)out_dim,
.ne11 = (int32_t)n_tokens,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = src1_token_stride,
.nb12 = (uint64_t)n_tokens * src1_token_stride,
.nb13 = (uint64_t)n_tokens * src1_token_stride,
.ne0 = (int32_t)out_dim,
.ne1 = (int32_t)n_tokens,
.ne2 = 1,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = dst_token_stride,
.nb2 = (uint64_t)n_tokens * dst_token_stride,
.nb3 = (uint64_t)n_tokens * dst_token_stride,
.offs = 0,
.o1 = { 0 },
};
}
static int ds4_gpu_encode_moe_sum_experts(
id<MTLCommandBuffer> cb,
id<MTLBuffer> experts,
NSUInteger experts_off,
id<MTLBuffer> out,
NSUInteger out_off,
uint32_t out_dim,
uint32_t n_expert,
uint32_t n_tokens) {
if (!cb || !experts || !out || out_dim == 0 || n_expert < 2 || n_tokens == 0) return 0;
const uint64_t out_row_bytes = (uint64_t)out_dim * sizeof(float);
const uint64_t expert_token_stride = (uint64_t)n_expert * out_row_bytes;
if (n_expert == 6 &&
ds4_gpu_encode_moe_sum6(cb,
experts,
experts_off,
out,
out_off,
out_dim,
n_tokens)) {
return 1;
}
if (n_expert == 8 &&
ds4_gpu_encode_moe_sum8(cb,
experts,
experts_off,
out,
out_off,
out_dim,
n_tokens)) {
return 1;
}
ds4_gpu_bin_args first =
ds4_gpu_make_moe_add_args(out_dim, n_tokens, expert_token_stride, expert_token_stride, out_row_bytes);
if (!ds4_gpu_encode_bin_f32_rows(cb,
g_add_pipeline,
&first,
experts,
experts_off,
experts,
experts_off + (NSUInteger)out_row_bytes,
out,
out_off)) {
return 0;
}
ds4_gpu_bin_args accum =
ds4_gpu_make_moe_add_args(out_dim, n_tokens, out_row_bytes, expert_token_stride, out_row_bytes);
for (uint32_t slot = 2; slot < n_expert; slot++) {
if (!ds4_gpu_encode_bin_f32_rows(cb,
g_add_pipeline,
&accum,
out,
out_off,
experts,
experts_off + (NSUInteger)((uint64_t)slot * out_row_bytes),
out,
out_off)) {
return 0;
}
}
return 1;
}
static int ds4_gpu_encode_get_rows_i32_token_rows(
id<MTLCommandBuffer> cb,
id<MTLBuffer> table,
NSUInteger table_off,
id<MTLBuffer> tokens,
NSUInteger tokens_off,
const int32_t *token_inline,
id<MTLBuffer> selected,
NSUInteger selected_off,
uint32_t hash_rows,
uint32_t n_cols,
uint32_t n_tokens) {
if (!cb || !table || !selected || hash_rows == 0 || n_cols == 0 || n_tokens == 0) return 0;
if (!tokens && !token_inline) return 0;
const uint64_t table_row_bytes = (uint64_t)n_cols * sizeof(int32_t);
const uint64_t token_bytes = (uint64_t)n_tokens * sizeof(int32_t);
ds4_gpu_get_rows_args args = {
.ne00t = (int64_t)n_cols,
.ne00 = (int64_t)n_cols,
.nb01 = table_row_bytes,
.nb02 = (uint64_t)hash_rows * table_row_bytes,
.nb03 = (uint64_t)hash_rows * table_row_bytes,
.ne10 = (int32_t)n_tokens,
.nb10 = sizeof(int32_t),
.nb11 = token_bytes,
.nb12 = token_bytes,
.nb1 = table_row_bytes,
.nb2 = (uint64_t)n_tokens * table_row_bytes,
.nb3 = (uint64_t)n_tokens * table_row_bytes,
};
NSUInteger nth = (NSUInteger)n_cols;
const NSUInteger max_threads = g_get_rows_i32_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > max_threads) nth = max_threads;
if (nth == 0) nth = 1u;
const NSUInteger nw0 = ((NSUInteger)n_cols + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_get_rows_i32_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:table offset:table_off atIndex:1];
if (tokens) {
[enc setBuffer:tokens offset:tokens_off atIndex:2];
} else {
[enc setBytes:token_inline length:sizeof(*token_inline) atIndex:2];
}
[enc setBuffer:selected offset:selected_off atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(nw0 * n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_get_rows_f32_router_weights(
id<MTLCommandBuffer> cb,
id<MTLBuffer> probs,
NSUInteger probs_off,
id<MTLBuffer> selected,
NSUInteger selected_off,
id<MTLBuffer> weights,
NSUInteger weights_off,
uint32_t n_expert,
uint32_t n_expert_used,
uint32_t n_tokens) {
if (!cb || !probs || !selected || !weights || n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0;
const uint64_t probs_token_bytes = (uint64_t)n_expert * sizeof(float);
const uint64_t selected_row_bytes = (uint64_t)n_expert_used * sizeof(int32_t);
const uint64_t weights_row_bytes = (uint64_t)n_expert_used * sizeof(float);
ds4_gpu_get_rows_args args = {
.ne00t = 1,
.ne00 = 1,
.nb01 = sizeof(float),
.nb02 = probs_token_bytes,
.nb03 = (uint64_t)n_tokens * probs_token_bytes,
.ne10 = (int64_t)n_expert_used,
.nb10 = sizeof(int32_t),
.nb11 = selected_row_bytes,
.nb12 = (uint64_t)n_tokens * selected_row_bytes,
.nb1 = sizeof(float),
.nb2 = weights_row_bytes,
.nb3 = (uint64_t)n_tokens * weights_row_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_get_rows_f32_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:probs offset:probs_off atIndex:1];
[enc setBuffer:selected offset:selected_off atIndex:2];
[enc setBuffer:weights offset:weights_off atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_expert_used, n_tokens, 1)
threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_sum_rows_f32(
id<MTLCommandBuffer> cb,
id<MTLBuffer> src,
NSUInteger src_off,
id<MTLBuffer> dst,
NSUInteger dst_off,
uint32_t width,
uint32_t rows) {
if (!cb || !src || !dst || width == 0 || rows == 0) return 0;
const uint64_t src_row_bytes = (uint64_t)width * sizeof(float);
ds4_gpu_kargs_sum_rows args = {
.ne00 = (int64_t)width,
.ne01 = (int64_t)rows,
.ne02 = 1,
.ne03 = 1,
.nb00 = sizeof(float),
.nb01 = src_row_bytes,
.nb02 = (uint64_t)rows * src_row_bytes,
.nb03 = (uint64_t)rows * src_row_bytes,
.ne0 = 1,
.ne1 = (int64_t)rows,
.ne2 = 1,
.ne3 = 1,
.nb0 = sizeof(float),
.nb1 = sizeof(float),
.nb2 = (uint64_t)rows * sizeof(float),
.nb3 = (uint64_t)rows * sizeof(float),
};
NSUInteger nth = 32u;
const NSUInteger max_threads = g_sum_rows_f32_f32_pipeline.maxTotalThreadsPerThreadgroup;
while (nth < (NSUInteger)args.ne00 && nth < max_threads) nth *= 2u;
if (nth > max_threads) nth = max_threads;
if (nth > (NSUInteger)args.ne00) nth = (NSUInteger)args.ne00;
if (nth == 0) nth = 1u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_sum_rows_f32_f32_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:src offset:src_off atIndex:1];
[enc setBuffer:dst offset:dst_off atIndex:2];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
static int ds4_gpu_encode_router_select(
id<MTLCommandBuffer> cb,
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
id<MTLBuffer> logitsbuf,
NSUInteger logits_off,
id<MTLBuffer> biasbuf,
NSUInteger bias_off,
id<MTLBuffer> hashbuf,
NSUInteger hash_off,
id<MTLBuffer> tokensbuf,
NSUInteger tokens_off,
id<MTLBuffer> visual_biasbuf,
NSUInteger visual_bias_off,
const int32_t *single_token,
uint32_t hash_rows,
uint32_t vocab_size,
uint32_t n_tokens,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
bool has_bias,
bool hash_mode,
bool mixed_visual) {
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
const NSUInteger selected_off = ds4_gpu_tensor_offset(selected);
const NSUInteger weights_off = ds4_gpu_tensor_offset(weights);
const NSUInteger probs_off = ds4_gpu_tensor_offset(probs);
if (!cb || !selectedbuf || !weightsbuf || !probsbuf || !logitsbuf ||
(mixed_visual && (!visual_biasbuf || !tokensbuf || single_token)) ||
n_tokens == 0 || n_expert == 0 || n_expert_used == 0) return 0;
const NSUInteger probs_bytes = (NSUInteger)n_tokens * (NSUInteger)n_expert * sizeof(float);
const bool flash_router_fast_path =
n_expert == 256u &&
n_expert_used == 6u &&
fabsf(expert_weight_scale - 1.5f) <= 1.0e-6f;
int ok = 0;
if (flash_router_fast_path && !mixed_visual &&
!g_quality_mode && n_tokens == 1 &&
getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL) {
const bool pre_m5_device =
ds4_gpu_device_is_pre_m5_apple_silicon();
const bool use_simd_finalize =
!hash_mode &&
g_dsv4_router_finalize_one_simd_pipeline != nil &&
g_dsv4_router_finalize_one_simd_pipeline.threadExecutionWidth == 32u &&
g_dsv4_router_finalize_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u &&
(pre_m5_device ||
ds4_gpu_device_name_contains("M5")) &&
(!pre_m5_device ||
getenv("DS4_METAL_DISABLE_PRE_M5_ROUTER_SIMD_FINALIZE") == NULL);
const bool use_simd_weights_fusion =
use_simd_finalize &&
g_dsv4_router_finalize_weights_one_simd_pipeline != nil &&
g_dsv4_router_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u &&
g_dsv4_router_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u &&
(pre_m5_device ||
ds4_gpu_device_name_contains("M5")) &&
(!pre_m5_device ||
getenv("DS4_METAL_DISABLE_PRE_M5_ROUTER_SIMD_WEIGHTS_FUSION") == NULL);
const bool use_pre_m5_transform_finalize_fusion_default =
!g_ssd_streaming_mode &&
(pre_m5_device || ds4_gpu_device_name_contains("M5")) &&
getenv("DS4_METAL_DISABLE_PRE_M5_ROUTER_TRANSFORM_FINALIZE_FUSION") == NULL;
const bool use_transform_finalize_fusion =
use_simd_weights_fusion &&
!hash_mode &&
use_pre_m5_transform_finalize_fusion_default &&
(pre_m5_device ||
ds4_gpu_device_name_contains("M5")) &&
g_dsv4_router_transform_finalize_weights_one_simd_pipeline != nil &&
g_dsv4_router_transform_finalize_weights_one_simd_pipeline.threadExecutionWidth == 32u &&
g_dsv4_router_transform_finalize_weights_one_simd_pipeline.maxTotalThreadsPerThreadgroup >= 256u;
id<MTLComputePipelineState> softplus_sqrt_pipeline =
use_transform_finalize_fusion
? nil
: ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline,
"kernel_dsv4_softplus_sqrt_f32_4");
id<MTLComputePipelineState> router_finalize_pipeline =
ds4_gpu_hot_pipeline(
use_transform_finalize_fusion
? g_dsv4_router_transform_finalize_weights_one_simd_pipeline
: use_simd_weights_fusion
? g_dsv4_router_finalize_weights_one_simd_pipeline
: use_simd_finalize
? g_dsv4_router_finalize_one_simd_pipeline
: g_dsv4_router_finalize_one_pipeline,
use_transform_finalize_fusion
? "kernel_dsv4_router_transform_finalize_weights_one_simd"
: use_simd_weights_fusion
? "kernel_dsv4_router_finalize_weights_one_simd"
: use_simd_finalize
? "kernel_dsv4_router_finalize_one_simd"
: "kernel_dsv4_router_finalize_one");
id<MTLComputePipelineState> router_weights_pipeline = use_simd_weights_fusion
? nil
: ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline,
"kernel_dsv4_router_weights_one");
if ((!use_transform_finalize_fusion && !softplus_sqrt_pipeline) ||
!router_finalize_pipeline ||
(!use_simd_weights_fusion && !router_weights_pipeline)) return 0;
if (!use_transform_finalize_fusion) {
ok = ds4_gpu_encode_unary_f32_rows(cb,
softplus_sqrt_pipeline,
logitsbuf,
logits_off,
probsbuf,
probs_off,
n_expert,
1,
1,
0.0f,
0.0f);
if (!ok) return 0;
}
const bool use_token_buffer = single_token == NULL;
ds4_gpu_dsv4_router_select_one_args args = {
.has_bias = has_bias ? 1u : 0u,
.hash_mode = hash_mode ? 1u : 0u,
.use_token_buffer = use_token_buffer ? 1u : 0u,
.token = single_token ? (uint32_t)*single_token : 0u,
.hash_rows = hash_rows,
};
const float zero_f32 = 0.0f;
const int32_t zero_i32 = 0;
if ((has_bias && !biasbuf) ||
(hash_mode && !hashbuf) ||
(use_token_buffer && !tokensbuf)) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:router_finalize_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
if (use_transform_finalize_fusion) {
[enc setBuffer:logitsbuf offset:logits_off atIndex:1];
[enc setBuffer:probsbuf offset:probs_off atIndex:2];
if (has_bias) {
[enc setBuffer:biasbuf offset:bias_off atIndex:3];
} else {
[enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:3];
}
[enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4];
if (use_token_buffer) {
[enc setBuffer:tokensbuf offset:tokens_off atIndex:5];
} else {
[enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:5];
}
[enc setBuffer:selectedbuf offset:selected_off atIndex:6];
[enc setBuffer:weightsbuf offset:weights_off atIndex:7];
} else {
[enc setBuffer:probsbuf offset:probs_off atIndex:1];
if (has_bias) {
[enc setBuffer:biasbuf offset:bias_off atIndex:2];
} else {
[enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:2];
}
if (hash_mode) {
[enc setBuffer:hashbuf offset:hash_off atIndex:3];
} else {
[enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:3];
}
if (use_token_buffer) {
[enc setBuffer:tokensbuf offset:tokens_off atIndex:4];
} else {
[enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4];
}
[enc setBuffer:selectedbuf offset:selected_off atIndex:5];
if (use_simd_weights_fusion) {
[enc setBuffer:weightsbuf offset:weights_off atIndex:6];
}
}
const NSUInteger router_finalize_scratch_bytes = use_simd_finalize
? 2u * (256u * sizeof(float) + 256u * sizeof(int32_t))
: 256u * sizeof(float) + 256u * sizeof(int32_t);
[enc setThreadgroupMemoryLength:router_finalize_scratch_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (use_simd_weights_fusion) return 1;
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:router_weights_pipeline];
[enc setBuffer:probsbuf offset:probs_off atIndex:0];
[enc setBuffer:selectedbuf offset:selected_off atIndex:1];
[enc setBuffer:weightsbuf offset:weights_off atIndex:2];
[enc dispatchThreads:MTLSizeMake(6, 1, 1)
threadsPerThreadgroup:MTLSizeMake(6, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
if (flash_router_fast_path && !mixed_visual && !g_quality_mode && n_tokens == 1) {
id<MTLComputePipelineState> softplus_sqrt_pipeline =
ds4_gpu_hot_pipeline(g_dsv4_softplus_sqrt_pipeline,
"kernel_dsv4_softplus_sqrt_f32_4");
ok = softplus_sqrt_pipeline &&
ds4_gpu_encode_unary_f32_rows(cb,
softplus_sqrt_pipeline,
logitsbuf,
logits_off,
probsbuf,
probs_off,
n_expert,
1,
1,
0.0f,
0.0f);
} else {
ok = ds4_gpu_encode_unary_f32_rows(cb,
g_unary_softplus_pipeline,
logitsbuf,
logits_off,
probsbuf,
probs_off,
n_expert,
n_tokens,
1,
0.0f,
0.0f) &&
ds4_gpu_encode_unary_f32_rows(cb,
g_unary_sqrt_pipeline,
probsbuf,
probs_off,
probsbuf,
probs_off,
n_expert,
n_tokens,
1,
0.0f,
0.0f);
}
if (!ok) return 0;
if (mixed_visual) {
if (!g_dsv4_router_select_visual_batch_pipeline ||
n_expert != 256u || n_expert_used != 6u) return 0;
ds4_gpu_dsv4_router_select_visual_args args = {
.hash_rows = hash_rows,
.vocab_size = vocab_size,
.n_tokens = n_tokens,
.has_bias = has_bias ? 1u : 0u,
.hash_mode = hash_mode ? 1u : 0u,
};
const float zero_f32 = 0.0f;
const int32_t zero_i32 = 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_router_select_visual_batch_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:probsbuf offset:probs_off atIndex:1];
if (has_bias) [enc setBuffer:biasbuf offset:bias_off atIndex:2];
else [enc setBytes:&zero_f32 length:sizeof(zero_f32) atIndex:2];
[enc setBuffer:visual_biasbuf offset:visual_bias_off atIndex:3];
if (hash_mode) [enc setBuffer:hashbuf offset:hash_off atIndex:4];
else [enc setBytes:&zero_i32 length:sizeof(zero_i32) atIndex:4];
[enc setBuffer:tokensbuf offset:tokens_off atIndex:5];
[enc setBuffer:selectedbuf offset:selected_off atIndex:6];
[enc setThreadgroupMemoryLength:
256u * (sizeof(float) + sizeof(int32_t)) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
ok = 1;
} else if (hash_mode) {
ok = ds4_gpu_encode_get_rows_i32_token_rows(cb,
hashbuf,
hash_off,
tokensbuf,
tokens_off,
single_token,
selectedbuf,
selected_off,
hash_rows,
n_expert_used,
n_tokens);
} else {
ds4_gpu_tensor *score_tensor = probs;
DS4MetalTensor *selection_view = nil;
if (has_bias) {
if (!biasbuf ||
!ds4_gpu_ensure_scratch_buffer(&g_router_selection_buffer,
&g_router_selection_bytes,
probs_bytes,
"ds4_router_selection")) {
return 0;
}
ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_expert, n_tokens, n_expert);
ok = ds4_gpu_encode_bin_f32_rows(cb,
g_add_pipeline,
&add_args,
probsbuf,
probs_off,
biasbuf,
bias_off,
g_router_selection_buffer,
0);
if (!ok) return 0;
selection_view = [DS4MetalTensor new];
selection_view.buffer = g_router_selection_buffer;
selection_view.offset = 0;
selection_view.bytes = probs_bytes;
selection_view.owner = 0;
score_tensor = (__bridge ds4_gpu_tensor *)selection_view;
}
ok = ds4_gpu_indexer_topk_tensor(selected, score_tensor, n_expert, n_tokens, n_expert_used) != 0;
}
if (!ok) return 0;
const bool use_batch_weights_fusion =
flash_router_fast_path && !g_quality_mode && n_tokens > 1u &&
g_dsv4_router_weights_batch_pipeline != nil &&
ds4_gpu_device_name_contains("M3") &&
getenv("DS4_METAL_DISABLE_ROUTER_WEIGHTS_BATCH_FUSION") == NULL &&
getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL;
if (use_batch_weights_fusion) {
const float scale = expert_weight_scale;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_dsv4_router_weights_batch_pipeline];
[enc setBytes:&scale length:sizeof(scale) atIndex:0];
[enc setBuffer:probsbuf offset:probs_off atIndex:1];
[enc setBuffer:selectedbuf offset:selected_off atIndex:2];
[enc setBuffer:weightsbuf offset:weights_off atIndex:3];
[enc setThreadgroupMemoryLength:40u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(6, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) {
id<MTLComputePipelineState> router_weights_pipeline =
ds4_gpu_hot_pipeline(g_dsv4_router_weights_one_pipeline,
"kernel_dsv4_router_weights_one");
if (!router_weights_pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:router_weights_pipeline];
[enc setBuffer:probsbuf offset:probs_off atIndex:0];
[enc setBuffer:selectedbuf offset:selected_off atIndex:1];
[enc setBuffer:weightsbuf offset:weights_off atIndex:2];
[enc dispatchThreads:MTLSizeMake(6, 1, 1)
threadsPerThreadgroup:MTLSizeMake(6, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
const NSUInteger sum_bytes = (NSUInteger)n_tokens * sizeof(float);
if (!ds4_gpu_ensure_scratch_buffer(&g_router_weight_sum_buffer,
&g_router_weight_sum_bytes,
sum_bytes,
"ds4_router_weight_sum")) {
return 0;
}
ok = ds4_gpu_encode_get_rows_f32_router_weights(cb,
probsbuf,
probs_off,
selectedbuf,
selected_off,
weightsbuf,
weights_off,
n_expert,
n_expert_used,
n_tokens) &&
ds4_gpu_encode_sum_rows_f32(cb,
weightsbuf,
weights_off,
g_router_weight_sum_buffer,
0,
n_expert_used,
n_tokens) &&
ds4_gpu_encode_unary_f32_rows(cb,
g_unary_clamp_pipeline,
g_router_weight_sum_buffer,
0,
g_router_weight_sum_buffer,
0,
1,
n_tokens,
0,
6.103515625e-5f,
ds4_gpu_positive_infinity());
if (!ok) return 0;
ds4_gpu_bin_args div_args = ds4_gpu_make_bin_rowwise_scalar_args(n_expert_used, n_tokens);
const float scale = expert_weight_scale;
ds4_gpu_bin_args scale_args = ds4_gpu_make_bin_rows_args(n_expert_used, n_tokens, 1);
ok = ds4_gpu_encode_bin_f32_rows(cb,
g_bin_div_row_pipeline,
&div_args,
weightsbuf,
weights_off,
g_router_weight_sum_buffer,
0,
weightsbuf,
weights_off);
if (!ok) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_bin_mul_scalar_pipeline];
[enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0];
[enc setBuffer:weightsbuf offset:weights_off atIndex:1];
[enc setBytes:&scale length:sizeof(scale) atIndex:2];
[enc setBuffer:weightsbuf offset:weights_off atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)scale_args.ne1,
(NSUInteger)scale_args.ne2,
(NSUInteger)scale_args.ne3)
threadsPerThreadgroup:MTLSizeMake(ds4_gpu_bin_threads(n_expert_used, g_bin_mul_scalar_pipeline), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return 1;
}
int ds4_gpu_glm_kv_lora_rms_norm_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv_raw,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !kv_raw || !model_map ||
n_tokens == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 ||
kv_lora_dim > kv_raw_dim || (kv_lora_dim & 3u) != 0 ||
!isfinite(eps) || eps < 0.0f) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(kv_raw);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float);
const uint64_t weight_bytes = (uint64_t)kv_lora_dim * sizeof(float);
if (!rawbuf || !outbuf ||
ds4_gpu_tensor_bytes(kv_raw) < raw_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal GLM KV RMS norm received undersized buffers\n");
return 0;
}
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal GLM KV RMS norm weight range is outside the mapped model\n");
return 0;
}
const bool exact_decode_weight_view =
n_tokens == 1u &&
weight_bytes <= (1ull << 20) &&
getenv("DS4_METAL_DISABLE_DECODE_NORM_EXACT_VIEWS") == NULL;
uint64_t weight_inner = 0;
id<MTLBuffer> weightbuf = exact_decode_weight_view ?
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
weight_offset,
weight_bytes,
&weight_inner) :
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&weight_inner);
if (!weightbuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_kv_lora_rms_norm_pipeline,
"kernel_glm_kv_lora_rms_norm");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_kv_lora_rms_norm_args args = {
.n_tokens = n_tokens,
.kv_raw_dim = kv_raw_dim,
.kv_lora_dim = kv_lora_dim,
.eps = eps,
};
const NSUInteger nth = ds4_gpu_rms_norm_threads(kv_lora_dim);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1];
[enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV RMS norm")) return 0;
}
return 1;
}
int ds4_gpu_glm_k_b_project_typed_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv_norm,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t n_head) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !kv_norm || !model_map ||
n_tokens == 0 || kv_lora_dim == 0 ||
qk_nope == 0 || n_head == 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> kvbuf = ds4_gpu_tensor_buffer(kv_norm);
uint64_t row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) {
fprintf(stderr, "ds4: Metal GLM k_b projection received unsupported weight type\n");
return 0;
}
const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim;
const uint64_t weight_bytes = weight_rows * row_bytes;
const uint64_t kv_bytes = (uint64_t)n_tokens * kv_lora_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float);
if (!outbuf || !kvbuf ||
ds4_gpu_tensor_bytes(kv_norm) < kv_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes) {
fprintf(stderr, "ds4: Metal GLM k_b projection received undersized buffers\n");
return 0;
}
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal GLM k_b projection range is outside the mapped model\n");
return 0;
}
const NSUInteger q_blocks = ((NSUInteger)qk_nope + 31u) / 32u;
if (q_blocks > 8u) {
fprintf(stderr, "ds4: Metal GLM k_b projection q width is too large for the tiled kernel\n");
return 0;
}
uint64_t weight_inner = 0;
id<MTLBuffer> weightbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&weight_inner);
if (!weightbuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_k_b_project_pipeline,
"kernel_glm_k_b_project_q8_0");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_k_b_project_args args = {
.n_tokens = n_tokens,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.n_head = n_head,
.row_bytes = (uint32_t)row_bytes,
.weight_type = weight_type,
.pad1 = 0,
.pad2 = 0,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
[enc setBuffer:kvbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1)
threadsPerThreadgroup:MTLSizeMake(32, q_blocks, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM k_b projection")) return 0;
}
return 1;
}
int ds4_gpu_glm_k_b_project_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv_norm,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t n_head) {
return ds4_gpu_glm_k_b_project_typed_tensor(out,
kv_norm,
model_map,
model_size,
weight_offset,
DS4_METAL_TENSOR_Q8_0,
n_tokens,
kv_lora_dim,
qk_nope,
n_head);
}
int ds4_gpu_glm_store_compact_kv_tensor(
ds4_gpu_tensor *kv_lora_cache,
ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *kv_norm,
const ds4_gpu_tensor *kv_raw,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_rope,
bool cache_f16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!kv_lora_cache || !k_rope_cache || !kv_norm || !kv_raw ||
n_tokens == 0 || cache_cap == 0 ||
kv_raw_dim == 0 || kv_lora_dim == 0 ||
kv_lora_dim > kv_raw_dim ||
qk_rope > kv_raw_dim - kv_lora_dim ||
pos0 > cache_cap || n_tokens > cache_cap - pos0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
id<MTLBuffer> kvnormbuf = ds4_gpu_tensor_buffer(kv_norm);
id<MTLBuffer> kvrawbuf = ds4_gpu_tensor_buffer(kv_raw);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t kv_cache_bytes =
(uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
const uint64_t rope_cache_bytes =
(uint64_t)cache_cap * qk_rope * cache_elem_bytes;
const uint64_t kv_norm_bytes =
(uint64_t)n_tokens * kv_lora_dim * sizeof(float);
const uint64_t kv_raw_bytes =
(uint64_t)n_tokens * kv_raw_dim * sizeof(float);
if (!kvcachebuf || !ropecachebuf || !kvnormbuf || !kvrawbuf ||
ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
ds4_gpu_tensor_bytes(kv_norm) < kv_norm_bytes ||
ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) {
fprintf(stderr, "ds4: Metal GLM compact KV store received undersized buffers\n");
return 0;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_store_compact_kv_pipeline,
"kernel_glm_store_compact_kv");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_store_compact_kv_args args = {
.pos0 = pos0,
.n_tokens = n_tokens,
.cache_cap = cache_cap,
.kv_raw_dim = kv_raw_dim,
.kv_lora_dim = kv_lora_dim,
.qk_rope = qk_rope,
.cache_f16 = cache_f16 ? 1u : 0u,
.pad1 = 0,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:kvnormbuf offset:ds4_gpu_tensor_offset(kv_norm) atIndex:1];
[enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:2];
[enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
[enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 2, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM compact KV store")) return 0;
}
return 1;
}
int ds4_gpu_glm_qkv_norm_store_compact_kv_tensor(
ds4_gpu_tensor *q_out,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t q_weight_offset,
uint32_t q_n,
ds4_gpu_tensor *kv_lora_cache,
ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *kv_raw,
uint64_t kv_weight_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_rope,
bool cache_f16,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!q_out || !q || !kv_lora_cache || !k_rope_cache || !kv_raw ||
!model_map || n_tokens == 0 || cache_cap == 0 ||
q_n == 0 || kv_raw_dim == 0 || kv_lora_dim == 0 ||
(q_n & 3u) != 0 || (kv_lora_dim & 3u) != 0 ||
kv_lora_dim > kv_raw_dim ||
qk_rope > kv_raw_dim - kv_lora_dim ||
pos0 > cache_cap || n_tokens > cache_cap - pos0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> qoutbuf = ds4_gpu_tensor_buffer(q_out);
id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
id<MTLBuffer> kvrawbuf = ds4_gpu_tensor_buffer(kv_raw);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t q_row_bytes = (uint64_t)q_n * sizeof(float);
const uint64_t kv_weight_bytes = (uint64_t)kv_lora_dim * sizeof(float);
const uint64_t kv_cache_bytes =
(uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
const uint64_t rope_cache_bytes =
(uint64_t)cache_cap * qk_rope * cache_elem_bytes;
const uint64_t kv_raw_bytes =
(uint64_t)n_tokens * kv_raw_dim * sizeof(float);
if (!qbuf || !qoutbuf || !kvcachebuf || !ropecachebuf || !kvrawbuf ||
ds4_gpu_tensor_bytes(q) < q_row_bytes * n_tokens ||
ds4_gpu_tensor_bytes(q_out) < q_row_bytes * n_tokens ||
ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
ds4_gpu_tensor_bytes(kv_raw) < kv_raw_bytes) {
fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store received undersized buffers\n");
return 0;
}
if (q_weight_offset > model_size || q_row_bytes > model_size - q_weight_offset ||
kv_weight_offset > model_size || kv_weight_bytes > model_size - kv_weight_offset) {
fprintf(stderr, "ds4: Metal GLM fused q/kv norm compact store weight range is outside the mapped model\n");
return 0;
}
uint64_t q_weight_inner = 0;
id<MTLBuffer> q_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size,
q_weight_offset, q_row_bytes,
&q_weight_inner);
if (!q_weightbuf) return 0;
uint64_t kv_weight_inner = 0;
id<MTLBuffer> kv_weightbuf = ds4_gpu_wrap_model_range(model_map, model_size,
kv_weight_offset, kv_weight_bytes,
&kv_weight_inner);
if (!kv_weightbuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_qkv_norm_store_compact_kv_pipeline,
"kernel_glm_qkv_norm_store_compact_kv");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_qkv_norm_store_compact_kv_args args = {
.pos0 = pos0,
.n_tokens = n_tokens,
.cache_cap = cache_cap,
.q_n = q_n,
.q_n4 = q_n / 4u,
.kv_raw_dim = kv_raw_dim,
.kv_lora_dim = kv_lora_dim,
.kv_lora_n4 = kv_lora_dim / 4u,
.qk_rope = qk_rope,
.cache_f16 = cache_f16 ? 1u : 0u,
.eps = eps,
.pad0 = 0,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:q_weightbuf offset:(NSUInteger)q_weight_inner atIndex:2];
[enc setBuffer:qoutbuf offset:ds4_gpu_tensor_offset(q_out) atIndex:3];
[enc setBuffer:kvrawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:4];
[enc setBuffer:kv_weightbuf offset:(NSUInteger)kv_weight_inner atIndex:5];
[enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:6];
[enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:7];
[enc setThreadgroupMemoryLength:32u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 3, 1)
threadsPerThreadgroup:MTLSizeMake(ds4_gpu_rms_norm_threads(q_n), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM fused q/kv norm compact store")) return 0;
}
return 1;
}
int ds4_gpu_glm_store_indexer_k_tensor(
ds4_gpu_tensor *indexer_key_cache,
const ds4_gpu_tensor *raw_k,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t bias_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t n_ctx_orig,
float eps,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!indexer_key_cache || !raw_k || !model_map ||
n_tokens == 0 || cache_cap == 0 ||
head_dim == 0 || rot_dim == 0 ||
rot_dim > head_dim || (rot_dim & 1u) != 0 ||
pos0 > cache_cap || n_tokens > cache_cap - pos0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_k);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t cache_bytes =
(uint64_t)cache_cap * head_dim * cache_elem_bytes;
const uint64_t raw_bytes =
(uint64_t)n_tokens * head_dim * sizeof(float);
const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
if (!cachebuf || !rawbuf ||
ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes ||
ds4_gpu_tensor_bytes(raw_k) < raw_bytes) {
fprintf(stderr, "ds4: Metal GLM indexer K store received undersized buffers\n");
return 0;
}
if (weight_offset > model_size || norm_bytes > model_size - weight_offset ||
bias_offset > model_size || norm_bytes > model_size - bias_offset) {
fprintf(stderr, "ds4: Metal GLM indexer K norm range is outside the mapped model\n");
return 0;
}
uint64_t weight_inner = 0;
uint64_t bias_inner = 0;
id<MTLBuffer> weightbuf =
ds4_gpu_wrap_model_range(model_map, model_size,
weight_offset, norm_bytes,
&weight_inner);
id<MTLBuffer> biasbuf =
ds4_gpu_wrap_model_range(model_map, model_size,
bias_offset, norm_bytes,
&bias_inner);
if (!weightbuf || !biasbuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_store_indexer_k_pipeline,
"kernel_glm_store_indexer_k");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_store_indexer_k_args args = {
.pos0 = pos0,
.n_tokens = n_tokens,
.cache_cap = cache_cap,
.head_dim = head_dim,
.rot_dim = rot_dim,
.n_ctx_orig = n_ctx_orig,
.cache_f16 = cache_f16 ? 1u : 0u,
.pad0 = 0,
.eps = eps,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
.pad1 = 0.0f,
};
const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_k) atIndex:1];
[enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:2];
[enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:3];
[enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:4];
[enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer K store")) return 0;
}
return 1;
}
int ds4_gpu_glm53_indexer_pool_update_tensor(
ds4_gpu_tensor *pool_cache,
ds4_gpu_tensor *tail_k,
ds4_gpu_tensor *tail_gate,
const ds4_gpu_tensor *raw_k,
const ds4_gpu_tensor *gate,
const void *model_map,
uint64_t model_size,
uint64_t norm_weight_offset,
uint64_t norm_bias_offset,
uint64_t ape_offset,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t head_dim,
uint32_t pool_size,
float eps,
bool cache_f16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!pool_cache || !tail_k || !tail_gate || !raw_k || !gate || !model_map ||
n_tokens == 0 || cache_cap == 0 || head_dim == 0 || pool_size != 4u ||
head_dim > 1024u || !isfinite(eps) || eps <= 0.0f ||
pos0 > cache_cap || n_tokens > cache_cap - pos0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> poolbuf = ds4_gpu_tensor_buffer(pool_cache);
id<MTLBuffer> tailkbuf = ds4_gpu_tensor_buffer(tail_k);
id<MTLBuffer> tailgatebuf = ds4_gpu_tensor_buffer(tail_gate);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_k);
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
const uint64_t pool_cap =
((uint64_t)cache_cap + pool_size - 1u) / pool_size;
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t pool_bytes = pool_cap * head_dim * cache_elem_bytes;
const uint64_t tail_bytes = (uint64_t)pool_size * head_dim * sizeof(float);
const uint64_t input_bytes = (uint64_t)n_tokens * head_dim * sizeof(float);
const uint64_t norm_bytes = (uint64_t)head_dim * sizeof(float);
const uint64_t ape_bytes =
(uint64_t)pool_size * head_dim * sizeof(uint16_t);
if (!poolbuf || !tailkbuf || !tailgatebuf || !rawbuf || !gatebuf ||
ds4_gpu_tensor_bytes(pool_cache) < pool_bytes ||
ds4_gpu_tensor_bytes(tail_k) < tail_bytes ||
ds4_gpu_tensor_bytes(tail_gate) < tail_bytes ||
ds4_gpu_tensor_bytes(raw_k) < input_bytes ||
ds4_gpu_tensor_bytes(gate) < input_bytes) {
fprintf(stderr, "ds4: Metal GLM-5.3 indexer pool received undersized buffers\n");
return 0;
}
if (norm_weight_offset > model_size ||
norm_bytes > model_size - norm_weight_offset ||
norm_bias_offset > model_size ||
norm_bytes > model_size - norm_bias_offset ||
ape_offset > model_size || ape_bytes > model_size - ape_offset) {
fprintf(stderr, "ds4: Metal GLM-5.3 indexer pool model range is invalid\n");
return 0;
}
uint64_t norm_weight_inner = 0;
uint64_t norm_bias_inner = 0;
uint64_t ape_inner = 0;
id<MTLBuffer> norm_weightbuf = ds4_gpu_wrap_model_range(
model_map, model_size, norm_weight_offset, norm_bytes, &norm_weight_inner);
id<MTLBuffer> norm_biasbuf = ds4_gpu_wrap_model_range(
model_map, model_size, norm_bias_offset, norm_bytes, &norm_bias_inner);
id<MTLBuffer> apebuf = ds4_gpu_wrap_model_range(
model_map, model_size, ape_offset, ape_bytes, &ape_inner);
if (!norm_weightbuf || !norm_biasbuf || !apebuf) return 0;
id<MTLComputePipelineState> pipeline = ds4_gpu_hot_pipeline(
g_glm53_indexer_pool_update_pipeline,
"kernel_glm53_indexer_pool_update");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBuffer:norm_weightbuf offset:(NSUInteger)norm_weight_inner atIndex:3];
[enc setBuffer:norm_biasbuf offset:(NSUInteger)norm_bias_inner atIndex:4];
[enc setBuffer:apebuf offset:(NSUInteger)ape_inner atIndex:5];
[enc setBuffer:poolbuf offset:ds4_gpu_tensor_offset(pool_cache) atIndex:6];
[enc setBuffer:tailkbuf offset:ds4_gpu_tensor_offset(tail_k) atIndex:7];
[enc setBuffer:tailgatebuf offset:ds4_gpu_tensor_offset(tail_gate) atIndex:8];
[enc setThreadgroupMemoryLength:
((NSUInteger)pool_size * head_dim + 2u * pool_size) * sizeof(float)
atIndex:0];
uint32_t done = 0;
const uint32_t leading = pos0 % pool_size;
if (leading != 0u) {
uint32_t chunk = pool_size - leading;
if (chunk > n_tokens) chunk = n_tokens;
ds4_gpu_glm53_indexer_pool_update_args args = {
.pos0 = pos0,
.n_tokens = chunk,
.cache_cap = cache_cap,
.head_dim = head_dim,
.pool_size = pool_size,
.cache_f16 = cache_f16 ? 1u : 0u,
.eps = eps,
.pad0 = 0,
};
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_k) atIndex:1];
[enc setBuffer:gatebuf offset:ds4_gpu_tensor_offset(gate) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(head_dim, 1, 1)];
done = chunk;
}
if (done < n_tokens) {
const uint32_t chunk = n_tokens - done;
ds4_gpu_glm53_indexer_pool_update_args args = {
.pos0 = pos0 + done,
.n_tokens = chunk,
.cache_cap = cache_cap,
.head_dim = head_dim,
.pool_size = pool_size,
.cache_f16 = cache_f16 ? 1u : 0u,
.eps = eps,
.pad0 = 0,
};
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:rawbuf
offset:ds4_gpu_tensor_offset(raw_k) +
(NSUInteger)((uint64_t)done * head_dim * sizeof(float))
atIndex:1];
[enc setBuffer:gatebuf
offset:ds4_gpu_tensor_offset(gate) +
(NSUInteger)((uint64_t)done * head_dim * sizeof(float))
atIndex:2];
const uint32_t groups = (chunk + pool_size - 1u) / pool_size;
[enc dispatchThreadgroups:MTLSizeMake(groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(head_dim, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 indexer pool update")) {
return 0;
}
}
return 1;
}
int ds4_gpu_glm53_expand_pool_selection_tensor(
ds4_gpu_tensor *raw_selected,
const ds4_gpu_tensor *pool_selected,
uint32_t n_tokens,
uint32_t pos0,
uint32_t selected_pools,
uint32_t index_topk,
uint32_t pool_size,
uint32_t output_width) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!raw_selected || !pool_selected || n_tokens == 0 ||
selected_pools == 0 || index_topk == 0 || pool_size == 0 ||
selected_pools > index_topk / pool_size ||
output_width < index_topk || output_width < index_topk + pool_size - 1u) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(raw_selected);
id<MTLBuffer> poolbuf = ds4_gpu_tensor_buffer(pool_selected);
const uint64_t raw_bytes =
(uint64_t)n_tokens * output_width * sizeof(uint32_t);
const uint64_t pool_bytes =
(uint64_t)n_tokens * selected_pools * sizeof(uint32_t);
if (!rawbuf || !poolbuf ||
ds4_gpu_tensor_bytes(raw_selected) < raw_bytes ||
ds4_gpu_tensor_bytes(pool_selected) < pool_bytes) {
fprintf(stderr, "ds4: Metal GLM-5.3 pool expansion received undersized buffers\n");
return 0;
}
id<MTLComputePipelineState> pipeline = ds4_gpu_hot_pipeline(
g_glm53_expand_pool_selection_pipeline,
"kernel_glm53_expand_pool_selection");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm53_expand_pool_selection_args args = {
.n_tokens = n_tokens,
.pos0 = pos0,
.selected_pools = selected_pools,
.index_topk = index_topk,
.pool_size = pool_size,
.output_width = output_width,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:poolbuf offset:ds4_gpu_tensor_offset(pool_selected) atIndex:1];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(raw_selected) atIndex:2];
const NSUInteger total = (NSUInteger)n_tokens * output_width;
const NSUInteger nth = MIN((NSUInteger)256, pipeline.maxTotalThreadsPerThreadgroup);
[enc dispatchThreads:MTLSizeMake(total, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 pool selection expansion")) {
return 0;
}
}
return 1;
}
int ds4_gpu_glm_build_kv_cache_tensor(
ds4_gpu_tensor *key_cache,
ds4_gpu_tensor *value_cache,
const ds4_gpu_tensor *kv_raw,
const ds4_gpu_tensor *k_nope,
const ds4_gpu_tensor *value,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t qk_dim = qk_nope + qk_rope;
if (!key_cache || !value_cache || !kv_raw || !k_nope || !value ||
n_tokens == 0 || cache_cap == 0 || n_head == 0 ||
kv_raw_dim == 0 || kv_lora_dim == 0 ||
qk_nope == 0 || value_dim == 0 ||
kv_lora_dim + qk_rope > kv_raw_dim ||
qk_dim < qk_nope || (qk_rope & 1u) != 0 ||
pos0 > cache_cap || n_tokens > cache_cap - pos0 ||
(qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) ||
(qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> keybuf = ds4_gpu_tensor_buffer(key_cache);
id<MTLBuffer> valbuf = ds4_gpu_tensor_buffer(value_cache);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(kv_raw);
id<MTLBuffer> knbuf = ds4_gpu_tensor_buffer(k_nope);
id<MTLBuffer> vbuf = ds4_gpu_tensor_buffer(value);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes;
const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes;
const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float);
const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float);
const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf ||
ds4_gpu_tensor_bytes(key_cache) < key_bytes ||
ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes ||
ds4_gpu_tensor_bytes(kv_raw) < raw_bytes ||
ds4_gpu_tensor_bytes(k_nope) < kn_bytes ||
ds4_gpu_tensor_bytes(value) < value_bytes) {
fprintf(stderr, "ds4: Metal GLM KV cache builder received undersized buffers\n");
return 0;
}
const bool decode_group4 =
n_tokens == 1u &&
n_head >= 4u &&
ds4_gpu_env_bool("DS4_METAL_DISABLE_GLM_DECODE_KV_GROUP4") <= 0;
id<MTLComputePipelineState> pipeline =
decode_group4 ?
ds4_gpu_hot_pipeline(g_glm_build_kv_cache_decode_group4_pipeline,
"kernel_glm_build_kv_cache_decode_group4") :
ds4_gpu_hot_pipeline(g_glm_build_kv_cache_pipeline,
"kernel_glm_build_kv_cache");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_build_kv_cache_args args = {
.pos0 = pos0,
.n_tokens = n_tokens,
.cache_cap = cache_cap,
.n_head = n_head,
.kv_raw_dim = kv_raw_dim,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_rope = qk_rope,
.value_dim = value_dim,
.n_ctx_orig = n_ctx_orig,
.cache_f16 = cache_f16 ? 1u : 0u,
.pad0 = 0u,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1];
[enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2];
[enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3];
[enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4];
[enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5];
const NSUInteger group_y = decode_group4 ?
((NSUInteger)n_head + 3u) / 4u :
(NSUInteger)n_head;
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, group_y, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM KV cache build")) return 0;
}
return 1;
}
int ds4_gpu_glm_build_kv_cache_flash_tensor(
ds4_gpu_tensor *key_cache,
ds4_gpu_tensor *value_cache,
const ds4_gpu_tensor *kv_raw,
const ds4_gpu_tensor *k_nope,
const ds4_gpu_tensor *value,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_cap,
uint32_t n_head,
uint32_t kv_raw_dim,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool cache_f16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t qk_dim = qk_nope + qk_rope;
if (!key_cache || !value_cache || !kv_raw || !k_nope || !value ||
pos0 != 0 || n_tokens == 0 || cache_cap == 0 || n_head == 0 ||
kv_raw_dim == 0 || kv_lora_dim == 0 ||
qk_nope == 0 || value_dim == 0 ||
kv_lora_dim + qk_rope > kv_raw_dim ||
qk_dim < qk_nope || (qk_rope & 1u) != 0 ||
n_tokens > cache_cap ||
(qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) ||
(qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> keybuf = ds4_gpu_tensor_buffer(key_cache);
id<MTLBuffer> valbuf = ds4_gpu_tensor_buffer(value_cache);
id<MTLBuffer> rawbuf = ds4_gpu_tensor_buffer(kv_raw);
id<MTLBuffer> knbuf = ds4_gpu_tensor_buffer(k_nope);
id<MTLBuffer> vbuf = ds4_gpu_tensor_buffer(value);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes;
const uint64_t cache_value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes;
const uint64_t raw_bytes = (uint64_t)n_tokens * kv_raw_dim * sizeof(float);
const uint64_t kn_bytes = (uint64_t)n_tokens * n_head * qk_nope * sizeof(float);
const uint64_t value_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
if (!keybuf || !valbuf || !rawbuf || !knbuf || !vbuf ||
ds4_gpu_tensor_bytes(key_cache) < key_bytes ||
ds4_gpu_tensor_bytes(value_cache) < cache_value_bytes ||
ds4_gpu_tensor_bytes(kv_raw) < raw_bytes ||
ds4_gpu_tensor_bytes(k_nope) < kn_bytes ||
ds4_gpu_tensor_bytes(value) < value_bytes) {
fprintf(stderr, "ds4: Metal GLM staged KV cache builder received undersized buffers\n");
return 0;
}
const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t);
const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t);
const NSUInteger key_f16_offset = 0;
const NSUInteger key_f16_bytes =
(NSUInteger)n_tokens * (NSUInteger)n_head * q_row_bytes_f16;
const NSUInteger value_f16_offset = key_f16_bytes;
const NSUInteger value_f16_bytes =
(NSUInteger)n_tokens * (NSUInteger)n_head * v_row_bytes_f16;
const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes;
if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_f16_bytes,
"ds4_glm_flash_attn_kv_f16")) {
return 0;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_build_kv_cache_flash_pipeline,
"kernel_glm_build_kv_cache_flash");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_build_kv_cache_args args = {
.pos0 = pos0,
.n_tokens = n_tokens,
.cache_cap = cache_cap,
.n_head = n_head,
.kv_raw_dim = kv_raw_dim,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_rope = qk_rope,
.value_dim = value_dim,
.n_ctx_orig = n_ctx_orig,
.cache_f16 = cache_f16 ? 1u : 0u,
.pad0 = 0u,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:rawbuf offset:ds4_gpu_tensor_offset(kv_raw) atIndex:1];
[enc setBuffer:knbuf offset:ds4_gpu_tensor_offset(k_nope) atIndex:2];
[enc setBuffer:vbuf offset:ds4_gpu_tensor_offset(value) atIndex:3];
[enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:4];
[enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:5];
[enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:6];
[enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:7];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM staged KV cache build")) return 0;
}
return 1;
}
static int ds4_gpu_glm_attention_flash_tensor_impl(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16,
int kv_pre_staged) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !key_cache || !value_cache ||
n_tokens == 0 || cache_len == 0 || cache_cap == 0 ||
n_head == 0 || qk_dim != 256u || value_dim != 256u ||
cache_len > cache_cap ||
pos0 > cache_len || n_tokens > cache_len - pos0 ||
cache_len > ds4_gpu_glm_flash_attention_max_cache_len()) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> keybuf = ds4_gpu_tensor_buffer(key_cache);
id<MTLBuffer> valbuf = ds4_gpu_tensor_buffer(value_cache);
const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes;
const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes;
if (!headsbuf || !qbuf || !keybuf || !valbuf ||
ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(key_cache) < key_bytes ||
ds4_gpu_tensor_bytes(value_cache) < value_bytes) {
fprintf(stderr, "ds4: Metal GLM FlashAttention received undersized buffers\n");
return 0;
}
const uint64_t key_elems = (uint64_t)cache_len * n_head * qk_dim;
const uint64_t value_elems = (uint64_t)cache_len * n_head * value_dim;
if (key_elems > UINT32_MAX || value_elems > UINT32_MAX) {
return 0;
}
const uint32_t nqptg = 8;
const uint32_t ncpsg = 64;
const uint32_t nsg = 4;
const bool has_kvpad = (cache_len % ncpsg) != 0;
const bool bc_mask = (n_tokens % nqptg) != 0;
const NSUInteger q_row_bytes = (NSUInteger)qk_dim * sizeof(float);
const NSUInteger q_row_bytes_f16 = (NSUInteger)qk_dim * sizeof(uint16_t);
const NSUInteger v_row_bytes = (NSUInteger)value_dim * sizeof(float);
const NSUInteger v_row_bytes_f16 = (NSUInteger)value_dim * sizeof(uint16_t);
const NSUInteger mask_bytes = (NSUInteger)n_tokens * (NSUInteger)cache_len * sizeof(uint16_t);
const NSUInteger key_f16_offset = 0;
const NSUInteger key_f16_bytes =
(NSUInteger)cache_len * (NSUInteger)n_head * q_row_bytes_f16;
const NSUInteger value_f16_offset = key_f16_bytes;
const NSUInteger value_f16_bytes =
(NSUInteger)cache_len * (NSUInteger)n_head * v_row_bytes_f16;
const NSUInteger kv_f16_bytes = key_f16_bytes + value_f16_bytes;
const NSUInteger pad_bytes = has_kvpad
? (NSUInteger)ncpsg * ((NSUInteger)n_head * (q_row_bytes_f16 + v_row_bytes_f16) +
(NSUInteger)n_tokens * sizeof(uint16_t))
: 1u;
const NSUInteger nblk0 = ((NSUInteger)cache_len + ncpsg - 1u) / ncpsg;
const NSUInteger nblk1 = ((NSUInteger)n_tokens + nqptg - 1u) / nqptg;
const NSUInteger blk_bytes = ds4_gpu_align_up_ns(nblk0 * nblk1, 32u);
id<MTLBuffer> mask_buffer =
ds4_gpu_glm_prefill_mask_buffer(pos0, n_tokens, cache_len, mask_bytes);
if (!mask_buffer) return 0;
if (kv_pre_staged) {
if (!g_flash_attn_kv_buffer || g_flash_attn_kv_bytes < kv_f16_bytes) {
fprintf(stderr, "ds4: GLM staged FlashAttention KV scratch is missing\n");
return 0;
}
} else if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_kv_buffer,
&g_flash_attn_kv_bytes,
kv_f16_bytes,
"ds4_glm_flash_attn_kv_f16")) {
return 0;
}
if (!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_pad_buffer,
&g_flash_attn_pad_bytes,
pad_bytes,
"ds4_glm_flash_attn_pad") ||
!ds4_gpu_ensure_scratch_buffer(&g_flash_attn_blk_buffer,
&g_flash_attn_blk_bytes,
blk_bytes,
"ds4_glm_flash_attn_blk")) {
return 0;
}
id<MTLComputePipelineState> pad_pipeline = nil;
if (has_kvpad) {
pad_pipeline = ds4_gpu_get_flash_attn_pad_pipeline(true, (int32_t)ncpsg);
if (!pad_pipeline) return 0;
}
id<MTLComputePipelineState> blk_pipeline =
ds4_gpu_get_flash_attn_blk_pipeline((int32_t)nqptg, (int32_t)ncpsg);
id<MTLComputePipelineState> attn_pipeline =
ds4_gpu_get_flash_attn_pipeline("kernel_flash_attn_ext_f16_dk256_dv256",
true, false, false, false, has_kvpad, bc_mask,
(int32_t)qk_dim,
(int32_t)value_dim,
(int32_t)nsg);
if (!blk_pipeline || !attn_pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (!kv_pre_staged) {
const bool copied = cache_f16 ?
(ds4_gpu_encode_cpy_f16_f16_3d(cb,
keybuf,
ds4_gpu_tensor_offset(key_cache),
g_flash_attn_kv_buffer,
key_f16_offset,
qk_dim,
cache_len,
n_head,
(uint64_t)n_head * q_row_bytes_f16,
q_row_bytes_f16,
q_row_bytes_f16,
(uint64_t)cache_len * q_row_bytes_f16) &&
ds4_gpu_encode_cpy_f16_f16_3d(cb,
valbuf,
ds4_gpu_tensor_offset(value_cache),
g_flash_attn_kv_buffer,
value_f16_offset,
value_dim,
cache_len,
n_head,
(uint64_t)n_head * v_row_bytes_f16,
v_row_bytes_f16,
v_row_bytes_f16,
(uint64_t)cache_len * v_row_bytes_f16)) :
(ds4_gpu_encode_cpy_f32_f16_3d(cb,
keybuf,
ds4_gpu_tensor_offset(key_cache),
g_flash_attn_kv_buffer,
key_f16_offset,
qk_dim,
cache_len,
n_head,
(uint64_t)n_head * q_row_bytes,
q_row_bytes,
q_row_bytes_f16,
(uint64_t)cache_len * q_row_bytes_f16) &&
ds4_gpu_encode_cpy_f32_f16_3d(cb,
valbuf,
ds4_gpu_tensor_offset(value_cache),
g_flash_attn_kv_buffer,
value_f16_offset,
value_dim,
cache_len,
n_head,
(uint64_t)n_head * v_row_bytes,
v_row_bytes,
v_row_bytes_f16,
(uint64_t)cache_len * v_row_bytes_f16));
if (!copied) {
return 0;
}
}
if (has_kvpad) {
ds4_gpu_flash_attn_pad_args pad_args = {
.ne11 = (int32_t)cache_len,
.ne_12_2 = (int32_t)n_head,
.ne_12_3 = 1,
.nb11 = q_row_bytes_f16,
.nb12 = (uint64_t)cache_len * q_row_bytes_f16,
.nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16,
.nb21 = v_row_bytes_f16,
.nb22 = (uint64_t)cache_len * v_row_bytes_f16,
.nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)cache_len * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pad_pipeline];
[enc setBytes:&pad_args length:sizeof(pad_args) atIndex:0];
[enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:2];
[enc setBuffer:mask_buffer offset:0 atIndex:3];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(ncpsg, n_head, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
ds4_gpu_flash_attn_blk_args blk_args = {
.ne01 = (int32_t)n_tokens,
.ne30 = (int32_t)cache_len,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)cache_len * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:blk_pipeline];
[enc setBytes:&blk_args length:sizeof(blk_args) atIndex:0];
[enc setBuffer:mask_buffer offset:0 atIndex:1];
[enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(nblk0, nblk1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
ds4_gpu_flash_attn_vec_args args = {
.ne01 = (int32_t)n_tokens,
.ne02 = (int32_t)n_head,
.ne03 = 1,
.nb01 = (uint64_t)n_head * q_row_bytes,
.nb02 = q_row_bytes,
.nb03 = (uint64_t)n_tokens * n_head * q_row_bytes,
.ne11 = (int32_t)cache_len,
.ne_12_2 = (int32_t)n_head,
.ne_12_3 = 1,
.ns10 = (int32_t)qk_dim,
.nb11 = q_row_bytes_f16,
.nb12 = (uint64_t)cache_len * q_row_bytes_f16,
.nb13 = (uint64_t)cache_len * (uint64_t)n_head * q_row_bytes_f16,
.ns20 = (int32_t)value_dim,
.nb21 = v_row_bytes_f16,
.nb22 = (uint64_t)cache_len * v_row_bytes_f16,
.nb23 = (uint64_t)cache_len * (uint64_t)n_head * v_row_bytes_f16,
.ne31 = (int32_t)n_tokens,
.ne32 = 1,
.ne33 = 1,
.nb31 = (uint64_t)cache_len * sizeof(uint16_t),
.nb32 = mask_bytes,
.nb33 = mask_bytes,
.ne1 = (int32_t)n_head,
.ne2 = (int32_t)n_tokens,
.ne3 = 1,
.scale = 1.0f / sqrtf((float)qk_dim),
.max_bias = 0.0f,
.m0 = 0.0f,
.m1 = 0.0f,
.n_head_log2 = 0,
.logit_softcap = 0.0f,
};
const NSUInteger padded_v = ds4_gpu_align_up_ns(value_dim, 64u);
const NSUInteger shared_elems = (NSUInteger)nqptg *
((NSUInteger)qk_dim + 2u * padded_v + 2u * (2u * (NSUInteger)ncpsg));
const NSUInteger shared_bytes = ds4_gpu_align_up_ns(shared_elems * (sizeof(float) / 2u), 16u);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:attn_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:g_flash_attn_kv_buffer offset:key_f16_offset atIndex:2];
[enc setBuffer:g_flash_attn_kv_buffer offset:value_f16_offset atIndex:3];
[enc setBuffer:mask_buffer offset:0 atIndex:4];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:5];
[enc setBuffer:g_flash_attn_pad_buffer offset:0 atIndex:6];
[enc setBuffer:g_flash_attn_blk_buffer offset:0 atIndex:7];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:8];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(nblk1, n_head, 1)
threadsPerThreadgroup:MTLSizeMake(32, nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM FlashAttention")) return 0;
}
return 1;
}
int ds4_gpu_glm_attention_flash_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16) {
return ds4_gpu_glm_attention_flash_tensor_impl(heads,
q,
key_cache,
value_cache,
pos0,
n_tokens,
cache_len,
cache_cap,
n_head,
qk_dim,
value_dim,
cache_f16,
0);
}
int ds4_gpu_glm_attention_flash_staged_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16) {
if (pos0 != 0 || n_tokens != cache_len) return 0;
return ds4_gpu_glm_attention_flash_tensor_impl(heads,
q,
key_cache,
value_cache,
pos0,
n_tokens,
cache_len,
cache_cap,
n_head,
qk_dim,
value_dim,
cache_f16,
1);
}
int ds4_gpu_glm_attention_full_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *key_cache,
const ds4_gpu_tensor *value_cache,
uint32_t pos0,
uint32_t n_tokens,
uint32_t cache_len,
uint32_t cache_cap,
uint32_t n_head,
uint32_t qk_dim,
uint32_t value_dim,
bool cache_f16) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !q || !key_cache || !value_cache ||
n_tokens == 0 || cache_len == 0 || cache_cap == 0 ||
n_head == 0 || qk_dim == 0 || value_dim == 0 ||
(qk_dim & 3u) != 0 ||
cache_len > cache_cap ||
pos0 > cache_len || n_tokens > cache_len - pos0 ||
cache_len > ds4_gpu_glm_full_attention_max_cache_len()) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> keybuf = ds4_gpu_tensor_buffer(key_cache);
id<MTLBuffer> valbuf = ds4_gpu_tensor_buffer(value_cache);
const uint64_t heads_bytes = (uint64_t)n_tokens * n_head * value_dim * sizeof(float);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t key_bytes = (uint64_t)cache_cap * n_head * qk_dim * cache_elem_bytes;
const uint64_t value_bytes = (uint64_t)cache_cap * n_head * value_dim * cache_elem_bytes;
if (!headsbuf || !qbuf || !keybuf || !valbuf ||
ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(key_cache) < key_bytes ||
ds4_gpu_tensor_bytes(value_cache) < value_bytes) {
fprintf(stderr, "ds4: Metal GLM attention received undersized buffers\n");
return 0;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_attention_full_pipeline,
"kernel_glm_attention_full");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const uint32_t full_attn_mode = 2u;
ds4_gpu_glm_attention_full_args args = {
.pos0 = pos0,
.n_tokens = n_tokens,
.cache_len = cache_len,
.cache_cap = cache_cap,
.n_head = n_head,
.qk_dim = qk_dim,
.value_dim = value_dim,
.pad0 = full_attn_mode,
.cache_f16 = cache_f16 ? 1u : 0u,
.pad1 = 0u,
.pad2 = 0u,
.scale = 1.0f / sqrtf((float)qk_dim),
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:keybuf offset:ds4_gpu_tensor_offset(key_cache) atIndex:2];
[enc setBuffer:valbuf offset:ds4_gpu_tensor_offset(value_cache) atIndex:3];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4];
[enc setThreadgroupMemoryLength:ds4_gpu_align_up_ns(
(256u + (NSUInteger)cache_len) * sizeof(float), 16u) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, (NSUInteger)n_head, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM full attention")) return 0;
}
return 1;
}
int ds4_gpu_glm_fill_selected_range_tensor(
ds4_gpu_tensor *selected,
uint32_t n_selected) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!selected || n_selected == 0) return 0;
@autoreleasepool {
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t);
if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) {
fprintf(stderr, "ds4: Metal GLM selected range received undersized buffer\n");
return 0;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_fill_selected_range_pipeline,
"kernel_glm_fill_selected_range");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_fill_selected_range_args args = {
.n_selected = n_selected,
};
const NSUInteger nth = 256u;
const NSUInteger n_groups = ((NSUInteger)n_selected + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range")) return 0;
}
return 1;
}
int ds4_gpu_glm_fill_selected_range_batch_tensor(
ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_selected,
uint32_t pad_row) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!selected || n_tokens == 0 || n_selected == 0) return 0;
@autoreleasepool {
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
const uint64_t total = (uint64_t)n_tokens * n_selected;
if (n_tokens != 0 && total / n_tokens != n_selected) return 0;
if (total > UINT64_MAX / sizeof(uint32_t)) return 0;
const uint64_t selected_bytes = total * sizeof(uint32_t);
if (!selectedbuf || ds4_gpu_tensor_bytes(selected) < selected_bytes) {
fprintf(stderr, "ds4: Metal GLM selected range batch received undersized buffer\n");
return 0;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_fill_selected_range_batch_pipeline,
"kernel_glm_fill_selected_range_batch");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_fill_selected_range_batch_args args = {
.n_tokens = n_tokens,
.pos0 = pos0,
.n_selected = n_selected,
.pad_row = pad_row,
};
const NSUInteger nth = 256u;
const NSUInteger n_groups = ((NSUInteger)total + nth - 1u) / nth;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake(n_groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM selected range batch")) return 0;
}
return 1;
}
static int ds4_gpu_glm_rope_tail_offset_tensor(
ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t rot_offset,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
const char *label) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!x || n_tokens == 0 || n_head == 0 || head_dim == 0 ||
rot_dim == 0 || rot_offset > head_dim || rot_dim > head_dim - rot_offset ||
(rot_dim & 1u) != 0 ||
pos0 > UINT32_MAX - n_tokens ||
!isfinite(freq_base) || freq_base <= 0.0f ||
!isfinite(freq_scale) || freq_scale <= 0.0f ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
const uint64_t bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
if (!xbuf || ds4_gpu_tensor_bytes(x) < bytes) {
fprintf(stderr, "ds4: Metal %s received undersized buffer\n",
label ? label : "GLM RoPE");
return 0;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_indexer_rope_tail_pipeline,
"kernel_glm_indexer_rope_tail_f32");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_rope_tail_args args = {
.n_tokens = n_tokens,
.n_head = n_head,
.head_dim = head_dim,
.rot_dim = rot_dim,
.rot_offset = rot_offset,
.pos0 = pos0,
.n_ctx_orig = n_ctx_orig,
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
};
const NSUInteger nth = ds4_gpu_rms_norm_threads(rot_dim);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, label ? label : "GLM RoPE")) return 0;
}
return 1;
}
int ds4_gpu_glm_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (rot_dim > head_dim) return 0;
return ds4_gpu_glm_rope_tail_offset_tensor(x,
n_tokens,
n_head,
head_dim,
rot_dim,
head_dim - rot_dim,
pos0,
n_ctx_orig,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow,
"GLM RoPE");
}
int ds4_gpu_glm_indexer_rope_tail_tensor(
ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t n_head,
uint32_t head_dim,
uint32_t rot_dim,
uint32_t pos0,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
return ds4_gpu_glm_rope_tail_offset_tensor(x,
n_tokens,
n_head,
head_dim,
rot_dim,
0,
pos0,
n_ctx_orig,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow,
"GLM indexer RoPE");
}
static int ds4_gpu_v41_index_score_wide_off(void) {
static int off = -1;
if (off < 0) off = getenv("DS4_METAL_DISABLE_V41_INDEX_SCORE_WIDE") != NULL;
return off;
}
static int ds4_gpu_glm_indexer_score_one_impl(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16,
bool fp4_inputs) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!scores || !q || !weights || !indexer_key_cache ||
n_rows == 0 || n_head == 0 || head_dim == 0 ||
!isfinite(scale) || scale <= 0.0f) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> scoresbuf = ds4_gpu_tensor_buffer(scores);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t score_bytes = (uint64_t)n_rows * sizeof(float);
const uint64_t q_bytes = (uint64_t)n_head * head_dim * sizeof(float);
const uint64_t weights_bytes = (uint64_t)n_head * sizeof(float);
const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes;
if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf ||
ds4_gpu_tensor_bytes(scores) < score_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(weights) < weights_bytes ||
ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) {
fprintf(stderr, "ds4: Metal GLM indexer score received undersized buffers\n");
return 0;
}
ds4_gpu_glm_indexer_score_one_args args = {
.n_rows = n_rows,
.n_head = n_head,
.head_dim = head_dim,
.cache_f16 = cache_f16 ? 1u : 0u,
.scale = scale,
};
if (n_head == 32u && head_dim == 128u) {
/* long rows stream one key per thread when FP4 query and keys make
* its summation order exact; the rest keep the per-key kernel */
const bool wide = (g_test_flags & DS4_GPU_TEST_INDEX_SCORE_WIDE) != 0u ||
((g_test_flags & DS4_GPU_TEST_INDEX_SCORE_DIRECT) == 0u && fp4_inputs &&
n_rows >= 8192u && !ds4_gpu_v41_index_score_wide_off());
id<MTLComputePipelineState> direct_pipeline = wide ?
ds4_gpu_hot_pipeline(g_glm_indexer_score_one_wide_pipeline,
"kernel_glm_indexer_score_one_wide") :
ds4_gpu_hot_pipeline(g_glm_indexer_score_one_direct_pipeline,
"kernel_glm_indexer_score_one_direct");
if (!direct_pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:direct_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
[enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3];
[enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
if (wide) {
[enc setThreadgroupMemoryLength:(4096u + 32u) * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)(n_rows + 255u) / 256u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
} else {
[enc setThreadgroupMemoryLength:(128u + 4u) * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer direct score")) return 0;
return 1;
}
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_indexer_score_one_pipeline,
"kernel_glm_indexer_score_one");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
[enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3];
[enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
[enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer score")) return 0;
}
return 1;
}
int ds4_gpu_glm_indexer_score_one_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16) {
return ds4_gpu_glm_indexer_score_one_impl(scores, q, weights, indexer_key_cache,
n_rows, n_head, head_dim, scale, cache_f16, false);
}
int ds4_gpu_dsv41_indexer_score_one_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16) {
return ds4_gpu_glm_indexer_score_one_impl(scores, q, weights, indexer_key_cache,
n_rows, n_head, head_dim, scale, cache_f16, true);
}
static int ds4_gpu_glm_indexer_scores_batch_grouped_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t row_group_size,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16,
bool tiled_f32) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!scores || !q || !weights || !indexer_key_cache ||
n_rows == 0 || n_tokens == 0 || row_group_size == 0 ||
n_head == 0 || head_dim != 128 ||
pos0 > UINT32_MAX - n_tokens ||
(uint64_t)(pos0 + n_tokens) / row_group_size > n_rows ||
!isfinite(scale) || scale <= 0.0f) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> scoresbuf = ds4_gpu_tensor_buffer(scores);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> cachebuf = ds4_gpu_tensor_buffer(indexer_key_cache);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t score_bytes = (uint64_t)n_rows * n_tokens * sizeof(float);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
const uint64_t weights_bytes = (uint64_t)n_tokens * n_head * sizeof(float);
const uint64_t cache_bytes = (uint64_t)n_rows * head_dim * cache_elem_bytes;
if (!scoresbuf || !qbuf || !weightsbuf || !cachebuf ||
ds4_gpu_tensor_bytes(scores) < score_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(weights) < weights_bytes ||
ds4_gpu_tensor_bytes(indexer_key_cache) < cache_bytes) {
fprintf(stderr, "ds4: Metal GLM indexer batch scores received undersized buffers\n");
return 0;
}
const bool force_scalar = g_quality_mode;
const bool use_tiled_f32 = tiled_f32;
/* a verify block's few rows take the tiled kernel too: the scalar one
* reduces every head once per key and row */
static int verify_tiled = -1;
if (verify_tiled < 0) verify_tiled = getenv("DS4_METAL_DISABLE_V41_VERIFY_TILED_INDEX") == NULL;
const bool use_tiled = !force_scalar && n_head == 32u && head_dim == 128u &&
(n_tokens >= 8u || (n_tokens >= 2u && g_v41_verify_rows && verify_tiled));
id<MTLComputePipelineState> pipeline =
use_tiled
? ds4_gpu_hot_pipeline(use_tiled_f32 ? g_glm_indexer_scores_tiled_f32_pipeline
: g_glm_indexer_scores_tiled_pipeline,
use_tiled_f32 ? "kernel_glm_indexer_scores_tiled_f32"
: "kernel_glm_indexer_scores_tiled")
: ds4_gpu_hot_pipeline(g_glm_indexer_scores_batch_pipeline,
"kernel_glm_indexer_scores_batch");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_indexer_scores_batch_args args = {
.n_rows = n_rows,
.n_tokens = n_tokens,
.n_head = n_head,
.head_dim = head_dim,
.pos0 = pos0,
.cache_f16 = cache_f16 ? 1u : 0u,
.row_group_size = row_group_size,
.pad0 = 0,
.q_token_stride = (uint64_t)n_head * head_dim * sizeof(float),
.q_head_stride = (uint64_t)head_dim * sizeof(float),
.weights_token_stride = (uint64_t)n_head * sizeof(float),
.score_token_stride = (uint64_t)n_rows * sizeof(float),
.scale = scale,
};
const NSUInteger nth = ds4_gpu_rms_norm_threads(head_dim);
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:2];
[enc setBuffer:cachebuf offset:ds4_gpu_tensor_offset(indexer_key_cache) atIndex:3];
[enc setBuffer:scoresbuf offset:ds4_gpu_tensor_offset(scores) atIndex:4];
if (use_tiled) {
const NSUInteger q_shared = 8u * 128u;
const NSUInteger k_shared = 32u * 128u;
const NSUInteger dot_shared = 8u * 32u;
if (use_tiled_f32) {
[enc setThreadgroupMemoryLength:(q_shared + k_shared + dot_shared) *
sizeof(float) atIndex:0];
} else {
[enc setThreadgroupMemoryLength:(q_shared + k_shared) * sizeof(uint16_t) +
dot_shared * sizeof(float) atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_rows + 31u) / 32u,
((NSUInteger)n_tokens + 7u) / 8u,
1)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
} else {
[enc setThreadgroupMemoryLength:nth * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows,
(NSUInteger)n_tokens,
1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexer batch scores")) return 0;
}
return 1;
}
int ds4_gpu_glm_indexer_scores_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16) {
return ds4_gpu_glm_indexer_scores_batch_grouped_tensor(
scores, q, weights, indexer_key_cache, n_rows, n_tokens, pos0, 1u,
n_head, head_dim, scale, cache_f16, false);
}
int ds4_gpu_glm53_indexer_scores_batch_tensor(
ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *indexer_key_cache,
uint32_t n_rows,
uint32_t n_tokens,
uint32_t pos0,
uint32_t pool_size,
uint32_t n_head,
uint32_t head_dim,
float scale,
bool cache_f16) {
if (pool_size != 4u) return 0;
return ds4_gpu_glm_indexer_scores_batch_grouped_tensor(
scores, q, weights, indexer_key_cache, n_rows, n_tokens, pos0,
pool_size, n_head, head_dim, scale, cache_f16, false);
}
int ds4_gpu_dsv41_tensor_ops_available(void) {
return ds4_gpu_mpp_available() && !g_quality_mode;
}
static uint64_t ds4_gpu_dsv41_indexer_flags_bytes(uint32_t source_rows, uint32_t rows) {
return (((uint64_t)rows + ((uint64_t)source_rows + 63u) / 64u) * 4u + 255u) & ~UINT64_C(255);
}
uint64_t ds4_gpu_dsv41_indexer_packed_bytes(uint32_t source_rows, uint32_t rows) {
return ds4_gpu_dsv41_indexer_flags_bytes(source_rows, rows) +
(uint64_t)rows * 32u * 128u * 2u + (((uint64_t)source_rows + 63u) / 64u) * 64u * 128u * 2u;
}
int ds4_gpu_dsv41_indexer_pack(ds4_gpu_tensor *packed,
const ds4_gpu_tensor *q, const ds4_gpu_tensor *keys,
uint32_t source_rows, uint32_t rows) {
if (!source_rows || source_rows > INT32_MAX || !rows || rows > INT32_MAX ||
!ds4_gpu_dsv41_tensor_ops_available()) return 0;
const ds4_gpu_tensor *tensors[] = {q, keys, packed};
const uint64_t bytes[] = {(uint64_t)rows * 32u * 128u * 4u,
(uint64_t)source_rows * 128u * 4u, ds4_gpu_dsv41_indexer_packed_bytes(source_rows, rows)};
for (uint32_t i = 0; i < 3; i++)
if (!tensors[i] || !ds4_gpu_tensor_buffer(tensors[i]) || ds4_gpu_tensor_bytes(tensors[i]) < bytes[i])
return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_indexer_pack");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
const uint32_t args[] = {source_rows, rows, 0, 0};
const uint64_t flags_bytes = ds4_gpu_dsv41_indexer_flags_bytes(source_rows, rows);
const uint64_t offset = ds4_gpu_tensor_offset(packed);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(q) offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(keys) offset:ds4_gpu_tensor_offset(keys) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(packed) offset:offset atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(packed) offset:offset + flags_bytes atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(packed) offset:offset + flags_bytes + (uint64_t)rows * 8192u atIndex:5];
[enc setThreadgroupMemoryLength:16 atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)rows + ((NSUInteger)source_rows + 63u) / 64u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 index input packing");
}
}
int ds4_gpu_dsv41_indexer_scores_packed(ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *keys,
const ds4_gpu_tensor *packed,
uint32_t source_rows, uint32_t rows,
uint32_t start, uint32_t ratio,
uint32_t packed_rows, uint32_t offset) {
if ((ratio != 1u && ratio != 2u) || !source_rows || source_rows > INT32_MAX || !rows ||
!packed_rows || packed_rows > INT32_MAX ||
rows > UINT32_MAX - start || (start + rows) / ratio > source_rows ||
offset > packed_rows || rows > packed_rows - offset ||
!ds4_gpu_dsv41_tensor_ops_available()) return 0;
const ds4_gpu_tensor *tensors[] = {q, weights, keys, scores, packed};
const uint64_t bytes[] = {(uint64_t)rows * 32u * 128u * 4u, (uint64_t)rows * 32u * 4u,
(uint64_t)source_rows * 128u * 4u, (uint64_t)rows * source_rows * 4u,
ds4_gpu_dsv41_indexer_packed_bytes(source_rows, packed_rows)};
for (uint32_t i = 0; i < 5; i++)
if (!tensors[i] || !ds4_gpu_tensor_buffer(tensors[i]) || ds4_gpu_tensor_bytes(tensors[i]) < bytes[i])
return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_indexer_scores_packed");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
const uint32_t args[] = {source_rows, rows, start, ratio}, range[] = {packed_rows, offset};
const uint64_t flags_bytes = ds4_gpu_dsv41_indexer_flags_bytes(source_rows, packed_rows);
const uint64_t base = ds4_gpu_tensor_offset(packed);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(args) atIndex:0];
[enc setBytes:range length:sizeof(range) atIndex:1];
for (NSUInteger i = 0; i < 5; i++)
[enc setBuffer:ds4_gpu_tensor_buffer(tensors[i]) offset:ds4_gpu_tensor_offset(tensors[i]) atIndex:i + 2];
[enc setBuffer:ds4_gpu_tensor_buffer(packed) offset:base + flags_bytes atIndex:7];
[enc setBuffer:ds4_gpu_tensor_buffer(packed) offset:base + flags_bytes + (uint64_t)packed_rows * 8192u atIndex:8];
[enc setThreadgroupMemoryLength:32u * 64u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)source_rows + 63u) / 64u, rows, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 packed index scores");
}
}
int ds4_gpu_dsv41_indexer_scores_batch(ds4_gpu_tensor *scores,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *keys,
uint32_t source_rows, uint32_t rows,
uint32_t start, uint32_t ratio) {
if (ratio != 1u && ratio != 2u) return 0;
return ds4_gpu_glm_indexer_scores_batch_grouped_tensor(scores, q, weights,
keys, source_rows, rows, start, ratio, 32, 128, 1.0f / 64.0f, false, true);
}
int ds4_gpu_glm_qk_lowrank_typed_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!qk_low || !q || !model_map ||
n_head == 0 || kv_lora_dim == 0 ||
qk_nope == 0 || qk_nope > qk_dim) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(qk_low);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
uint64_t row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) {
fprintf(stderr, "ds4: Metal GLM qk lowrank received unsupported weight type\n");
return 0;
}
const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim;
const uint64_t weight_bytes = weight_rows * row_bytes;
const uint64_t out_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float);
const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float);
if (!outbuf || !qbuf ||
ds4_gpu_tensor_bytes(qk_low) < out_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes) {
fprintf(stderr, "ds4: Metal GLM qk lowrank received undersized buffers\n");
return 0;
}
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal GLM qk lowrank range is outside the mapped model\n");
return 0;
}
uint64_t weight_inner = 0;
id<MTLBuffer> weightbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&weight_inner);
if (!weightbuf) return 0;
const int use_glm52 =
n_head == 64u &&
kv_lora_dim == 512u &&
qk_nope == 192u &&
qk_dim == 256u &&
row_bytes == 204u &&
weight_type == DS4_METAL_TENSOR_Q8_0;
/* Coalesced simdgroup variant: lanes split the 192-dot so weight
* reads coalesce, and 2048 threadgroups replace 64. The thread-
* per-row kernels measured ~7.5x off the weight-bandwidth floor
* (7.2ms of the decode token by skip-ablation). Covers the GLM 5.2
* shape for both Q8_0 k_b and the DenseQ4 GGUF's Q4_0 k_b (the Q8
* fast path above never engaged there). */
const int use_glm52_sg =
(n_head == 32u || n_head == 64u) &&
kv_lora_dim == 512u &&
qk_nope == 192u &&
qk_dim == 256u &&
((weight_type == DS4_METAL_TENSOR_Q8_0 && row_bytes == 204u) ||
(weight_type == DS4_METAL_TENSOR_Q4_0 && row_bytes == 108u)) &&
getenv("DS4_METAL_DISABLE_GLM_QKLOW_SG") == NULL &&
ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline,
"kernel_glm_qk_lowrank_q8_0_glm52_sg") != nil;
if (getenv("DS4_METAL_GLM_QKLOW_DEBUG")) {
static int printed = 0;
if (!printed) {
printed = 1;
fprintf(stderr, "ds4: qk_lowrank decode path: use_glm52=%d sg=%d n_head=%u kv=%u nope=%u dim=%u rb=%llu type=%u\n",
use_glm52, use_glm52_sg, n_head, kv_lora_dim, qk_nope, qk_dim,
(unsigned long long)row_bytes, weight_type);
}
}
id<MTLComputePipelineState> pipeline =
use_glm52_sg ?
ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_sg_pipeline,
"kernel_glm_qk_lowrank_q8_0_glm52_sg") :
use_glm52 ?
ds4_gpu_hot_pipeline(g_glm_qk_lowrank_glm52_pipeline,
"kernel_glm_qk_lowrank_q8_0_glm52") :
ds4_gpu_hot_pipeline(g_glm_qk_lowrank_pipeline,
"kernel_glm_qk_lowrank_q8_0");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_qk_lowrank_args args = {
.n_head = n_head,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_dim = qk_dim,
.row_bytes = (uint32_t)row_bytes,
.weight_type = weight_type,
.pad1 = 0,
.pad2 = 0,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3];
if (use_glm52_sg) {
/* 8 simdgroups x 2 rows per threadgroup: (64, 512/16) grid. */
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head,
(NSUInteger)(kv_lora_dim / 16u),
1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
} else {
if (use_glm52) {
[enc setThreadgroupMemoryLength:192u * sizeof(float) atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM qk lowrank")) return 0;
}
return 1;
}
int ds4_gpu_glm_qk_lowrank_q8_0_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim) {
return ds4_gpu_glm_qk_lowrank_typed_tensor(qk_low,
q,
model_map,
model_size,
weight_offset,
DS4_METAL_TENSOR_Q8_0,
n_head,
kv_lora_dim,
qk_nope,
qk_dim);
}
int ds4_gpu_glm_qk_lowrank_typed_batch_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!qk_low || !q || !model_map ||
n_tokens == 0 || n_head == 0 || kv_lora_dim == 0 ||
qk_nope == 0 || qk_nope > qk_dim) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(qk_low);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
uint64_t row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(weight_type, qk_nope, &row_bytes)) {
fprintf(stderr, "ds4: Metal GLM batch qk lowrank received unsupported weight type\n");
return 0;
}
const uint64_t weight_rows = (uint64_t)n_head * kv_lora_dim;
const uint64_t weight_bytes = weight_rows * row_bytes;
const uint64_t out_bytes =
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
const uint64_t q_bytes =
(uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
if (!outbuf || !qbuf ||
ds4_gpu_tensor_bytes(qk_low) < out_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes) {
fprintf(stderr, "ds4: Metal GLM batch qk lowrank received undersized buffers\n");
return 0;
}
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal GLM batch qk lowrank range is outside the mapped model\n");
return 0;
}
uint64_t weight_inner = 0;
id<MTLBuffer> weightbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&weight_inner);
if (!weightbuf) return 0;
const int use_glm52_t4 =
n_tokens >= 4u &&
n_head == 64u &&
kv_lora_dim == 512u &&
qk_nope == 192u &&
qk_dim == 256u &&
row_bytes == 204u &&
weight_type == DS4_METAL_TENSOR_Q8_0;
id<MTLComputePipelineState> pipeline =
use_glm52_t4 ?
ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_glm52_t4_pipeline,
"kernel_glm_qk_lowrank_q8_0_batch_glm52_t4") :
ds4_gpu_hot_pipeline(g_glm_qk_lowrank_batch_pipeline,
"kernel_glm_qk_lowrank_q8_0_batch");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
uint32_t head_base = 0;
uint32_t head_count = n_head;
ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count);
ds4_gpu_glm_qk_lowrank_batch_args args = {
.n_tokens = n_tokens,
.n_head = n_head,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_dim = qk_dim,
.row_bytes = (uint32_t)row_bytes,
.weight_type = weight_type,
.head_base = head_base,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:3];
if (use_glm52_t4) {
[enc setThreadgroupMemoryLength:4u * 192u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count,
((NSUInteger)n_tokens + 3u) / 4u,
1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
} else {
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count,
(NSUInteger)n_tokens,
1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch qk lowrank")) return 0;
}
return 1;
}
int ds4_gpu_glm_qk_lowrank_q8_0_batch_tensor(
ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *q,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_dim) {
return ds4_gpu_glm_qk_lowrank_typed_batch_tensor(qk_low,
q,
model_map,
model_size,
weight_offset,
DS4_METAL_TENSOR_Q8_0,
n_tokens,
n_head,
kv_lora_dim,
qk_nope,
qk_dim);
}
int ds4_gpu_glm_value_project_typed_batch_heads_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *lora,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t weight_type,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t value_dim) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!heads || !lora || !model_map ||
n_tokens == 0 || n_head == 0 ||
kv_lora_dim == 0 || value_dim == 0) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
id<MTLBuffer> lorabuf = ds4_gpu_tensor_buffer(lora);
uint64_t row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(weight_type, kv_lora_dim, &row_bytes)) {
fprintf(stderr, "ds4: Metal GLM batch value project received unsupported weight type\n");
return 0;
}
const uint64_t weight_rows = (uint64_t)n_head * value_dim;
const uint64_t weight_bytes = weight_rows * row_bytes;
const uint64_t heads_bytes =
(uint64_t)n_tokens * n_head * value_dim * sizeof(float);
const uint64_t lora_bytes =
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
if (!headsbuf || !lorabuf ||
ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(lora) < lora_bytes) {
fprintf(stderr, "ds4: Metal GLM batch value project received undersized buffers\n");
return 0;
}
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal GLM batch value project range is outside the mapped model\n");
return 0;
}
uint64_t weight_inner = 0;
id<MTLBuffer> weightbuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
weight_offset,
weight_bytes,
&weight_inner);
if (!weightbuf) return 0;
const int use_mma =
n_head == 64u &&
kv_lora_dim == 512u &&
value_dim == 256u &&
row_bytes == 544u &&
weight_type == DS4_METAL_TENSOR_Q8_0 &&
n_tokens >= 32u;
id<MTLComputePipelineState> pipeline =
use_mma ?
ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_mma_pipeline,
"kernel_glm_value_project_q8_0_batch_heads_mma") :
ds4_gpu_hot_pipeline(g_glm_value_project_q8_0_batch_heads_pipeline,
"kernel_glm_value_project_q8_0_batch_heads");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
uint32_t head_base = 0;
uint32_t head_count = n_head;
ds4_gpu_tp_attn_head_range(n_head, 8u, &head_base, &head_count);
ds4_gpu_glm_qk_lowrank_batch_args args = {
.n_tokens = n_tokens,
.n_head = n_head,
.kv_lora_dim = kv_lora_dim,
.qk_nope = 0,
.qk_dim = value_dim,
.row_bytes = (uint32_t)row_bytes,
.weight_type = weight_type,
.head_base = head_base,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weightbuf offset:(NSUInteger)weight_inner atIndex:1];
[enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora) atIndex:2];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:3];
if (use_mma) {
[enc setThreadgroupMemoryLength:16u * 1024u atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_tokens + 31u) / 32u,
((NSUInteger)value_dim + 63u) / 64u,
(NSUInteger)head_count)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
} else {
[enc setThreadgroupMemoryLength:(NSUInteger)kv_lora_dim * sizeof(float)
atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)head_count,
(NSUInteger)n_tokens,
1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch value project")) return 0;
}
return 1;
}
int ds4_gpu_glm_value_project_q8_0_batch_heads_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *lora,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n_tokens,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t value_dim) {
return ds4_gpu_glm_value_project_typed_batch_heads_tensor(heads,
lora,
model_map,
model_size,
weight_offset,
DS4_METAL_TENSOR_Q8_0,
n_tokens,
n_head,
kv_lora_dim,
value_dim);
}
int ds4_gpu_glm_attention_indexed_decode_typed_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t qk_dim = qk_nope + qk_rope;
if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
!model_map || !selected ||
n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
n_head == 0 || kv_lora_dim == 0 ||
qk_nope == 0 || (qk_rope & 1u) != 0 ||
value_dim == 0 || qk_dim < qk_nope ||
(qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) ||
(qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
uint64_t value_row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) {
fprintf(stderr, "ds4: Metal GLM indexed attention received unsupported value type\n");
return 0;
}
const uint64_t value_weight_rows = (uint64_t)n_head * value_dim;
const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes;
const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float);
const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float);
const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float);
const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t);
if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf ||
ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes) {
fprintf(stderr, "ds4: Metal GLM indexed attention received undersized buffers\n");
return 0;
}
if (value_weight_offset > model_size ||
value_weight_bytes > model_size - value_weight_offset) {
fprintf(stderr, "ds4: Metal GLM indexed attention value range is outside the mapped model\n");
return 0;
}
uint64_t value_inner = 0;
id<MTLBuffer> valuebuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
value_weight_offset,
value_weight_bytes,
&value_inner);
if (!valuebuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_attention_indexed_decode_pipeline,
"kernel_glm_attention_indexed_decode");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
ds4_gpu_glm_attention_indexed_decode_args args = {
.n_selected = n_selected,
.cache_cap = cache_cap,
.cache_f16 = cache_f16 ? 1u : 0u,
.n_head = n_head,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_rope = qk_rope,
.value_dim = value_dim,
.n_ctx_orig = n_ctx_orig,
.value_row_bytes = (uint32_t)value_row_bytes,
.scale = 1.0f / sqrtf((float)qk_dim),
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
.value_type = value_weight_type,
};
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
[enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
[enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
[enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7];
const NSUInteger scratch_floats =
256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim;
[enc setThreadgroupMemoryLength:ds4_gpu_align_up_ns(
scratch_floats * sizeof(float), 16u) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed attention decode")) return 0;
}
return 1;
}
int ds4_gpu_glm_attention_indexed_decode_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
return ds4_gpu_glm_attention_indexed_decode_typed_tensor(heads,
q,
qk_low,
kv_lora_cache,
k_rope_cache,
model_map,
model_size,
value_weight_offset,
DS4_METAL_TENSOR_Q8_0,
selected,
n_selected,
cache_cap,
cache_f16,
n_head,
kv_lora_dim,
qk_nope,
qk_rope,
value_dim,
n_ctx_orig,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow);
}
int ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(
ds4_gpu_tensor *heads,
ds4_gpu_tensor *partial_lora,
ds4_gpu_tensor *partial_ms,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
bool selected_rows_valid,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
uint32_t block_rows,
uint32_t n_blocks,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t qk_dim = qk_nope + qk_rope;
const uint32_t needed_blocks =
block_rows != 0u ? (n_selected + block_rows - 1u) / block_rows : 0u;
if (!heads || !partial_lora || !partial_ms || !q || !qk_low ||
!kv_lora_cache || !k_rope_cache || !model_map || !selected ||
n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
n_head == 0 || (n_head % 8u) != 0 ||
kv_lora_dim != 512u ||
qk_nope == 0 || qk_rope != 64u ||
value_dim == 0 || qk_dim < qk_nope ||
block_rows == 0u || needed_blocks == 0u ||
n_blocks < needed_blocks || n_blocks > 64u ||
!cache_f16 ||
(qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) ||
(qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
id<MTLBuffer> partial_lorabuf = ds4_gpu_tensor_buffer(partial_lora);
id<MTLBuffer> partial_msbuf = ds4_gpu_tensor_buffer(partial_ms);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
const uint64_t cache_elem_bytes = sizeof(uint16_t);
uint64_t value_row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) {
fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received unsupported value type\n");
return 0;
}
const uint64_t value_weight_rows = (uint64_t)n_head * value_dim;
const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes;
const uint64_t heads_bytes = (uint64_t)n_head * value_dim * sizeof(float);
const uint64_t partial_lora_bytes =
(uint64_t)n_blocks * n_head * kv_lora_dim * sizeof(float);
const uint64_t partial_ms_bytes =
(uint64_t)n_blocks * n_head * 2u * sizeof(float);
const uint64_t q_bytes = (uint64_t)n_head * qk_dim * sizeof(float);
const uint64_t low_bytes = (uint64_t)n_head * kv_lora_dim * sizeof(float);
const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
const uint64_t selected_bytes = (uint64_t)n_selected * sizeof(uint32_t);
if (!headsbuf || !partial_lorabuf || !partial_msbuf ||
!qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf ||
ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(partial_lora) < partial_lora_bytes ||
ds4_gpu_tensor_bytes(partial_ms) < partial_ms_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes) {
fprintf(stderr, "ds4: Metal GLM split grouped indexed attention received undersized buffers\n");
return 0;
}
if (value_weight_offset > model_size ||
value_weight_bytes > model_size - value_weight_offset) {
fprintf(stderr, "ds4: Metal GLM split grouped indexed attention value range is outside the mapped model\n");
return 0;
}
uint64_t value_inner = 0;
id<MTLBuffer> valuebuf =
ds4_gpu_wrap_model_range(model_map,
model_size,
value_weight_offset,
value_weight_bytes,
&value_inner);
if (!valuebuf) return 0;
const bool use_valid_fullheads =
selected_rows_valid && (n_head % 8u) == 0u;
id<MTLComputePipelineState> partial_pipeline =
ds4_gpu_hot_pipeline(use_valid_fullheads ?
g_glm_attention_indexed_decode_split_group8_partial_valid_fullheads_pipeline :
g_glm_attention_indexed_decode_split_group8_partial_pipeline,
use_valid_fullheads ?
"kernel_glm_attention_indexed_decode_split_group8_partial_valid_fullheads" :
"kernel_glm_attention_indexed_decode_split_group8_partial");
const bool use_reduce16 =
n_blocks == 16u && block_rows == 128u &&
n_selected == 2048u && value_dim == 256u;
id<MTLComputePipelineState> reduce_pipeline =
ds4_gpu_hot_pipeline(use_reduce16 ?
g_glm_attention_indexed_decode_split_group8_reduce16_pipeline :
g_glm_attention_indexed_decode_split_group8_reduce_pipeline,
use_reduce16 ?
"kernel_glm_attention_indexed_decode_split_group8_reduce16" :
"kernel_glm_attention_indexed_decode_split_group8_reduce");
if (!partial_pipeline || !reduce_pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_attention_indexed_decode_split_args args = {
.n_selected = n_selected,
.cache_cap = cache_cap,
.cache_f16 = 1u,
.n_head = n_head,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_rope = qk_rope,
.value_dim = value_dim,
.n_ctx_orig = n_ctx_orig,
.value_row_bytes = (uint32_t)value_row_bytes,
.block_rows = block_rows,
.n_blocks = n_blocks,
.scale = 1.0f / sqrtf((float)qk_dim),
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
.value_type = value_weight_type,
};
const NSUInteger stage_rows = 16u;
const NSUInteger kv_vecs = (NSUInteger)kv_lora_dim / 4u;
const NSUInteger rope_vecs = (NSUInteger)qk_rope / 4u;
const NSUInteger partial_scratch_bytes =
stage_rows * kv_vecs * sizeof(uint16_t) * 4u +
stage_rows * rope_vecs * sizeof(float) * 4u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:partial_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
[enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
[enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5];
[enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:6];
[enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:7];
[enc setThreadgroupMemoryLength:partial_scratch_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head / 8u, (NSUInteger)n_blocks, 1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
const NSUInteger reduce_scratch_floats = 256u + 64u + (NSUInteger)kv_lora_dim;
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:reduce_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:partial_lorabuf offset:ds4_gpu_tensor_offset(partial_lora) atIndex:1];
[enc setBuffer:partial_msbuf offset:ds4_gpu_tensor_offset(partial_ms) atIndex:2];
[enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:3];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:4];
[enc setThreadgroupMemoryLength:reduce_scratch_floats * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM split grouped indexed attention decode")) return 0;
}
return 1;
}
int ds4_gpu_glm_attention_indexed_decode_split_group8_tensor(
ds4_gpu_tensor *heads,
ds4_gpu_tensor *partial_lora,
ds4_gpu_tensor *partial_ms,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
const ds4_gpu_tensor *selected,
uint32_t n_selected,
bool selected_rows_valid,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
uint32_t block_rows,
uint32_t n_blocks,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
return ds4_gpu_glm_attention_indexed_decode_split_group8_typed_tensor(heads,
partial_lora,
partial_ms,
q,
qk_low,
kv_lora_cache,
k_rope_cache,
model_map,
model_size,
value_weight_offset,
DS4_METAL_TENSOR_Q8_0,
selected,
n_selected,
selected_rows_valid,
cache_cap,
cache_f16,
n_head,
kv_lora_dim,
qk_nope,
qk_rope,
value_dim,
n_ctx_orig,
block_rows,
n_blocks,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow);
}
int ds4_gpu_glm_attention_indexed_batch_typed_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
uint32_t value_weight_type,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t qk_dim = qk_nope + qk_rope;
if (!heads || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
!model_map || !selected ||
n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
n_head == 0 || kv_lora_dim == 0 ||
qk_nope == 0 || (qk_rope & 1u) != 0 ||
value_dim == 0 || qk_dim < qk_nope ||
(qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) ||
(qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> headsbuf = ds4_gpu_tensor_buffer(heads);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
uint64_t value_row_bytes = 0;
if (!ds4_gpu_quant_row_bytes(value_weight_type, kv_lora_dim, &value_row_bytes)) {
fprintf(stderr, "ds4: Metal GLM indexed batch attention received unsupported value type\n");
return 0;
}
const uint64_t value_weight_rows = (uint64_t)n_head * value_dim;
const uint64_t value_weight_bytes = value_weight_rows * value_row_bytes;
const uint64_t heads_bytes =
(uint64_t)n_tokens * n_head * value_dim * sizeof(float);
const uint64_t q_bytes =
(uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
const uint64_t low_bytes =
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
const uint64_t kv_cache_bytes = (uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
const uint64_t rope_cache_bytes = (uint64_t)cache_cap * qk_rope * cache_elem_bytes;
const uint64_t selected_bytes =
(uint64_t)n_tokens * n_selected * sizeof(uint32_t);
if (!headsbuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf ||
ds4_gpu_tensor_bytes(heads) < heads_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes) {
fprintf(stderr, "ds4: Metal GLM indexed batch attention received undersized buffers\n");
return 0;
}
if (value_weight_offset > model_size ||
value_weight_bytes > model_size - value_weight_offset) {
fprintf(stderr, "ds4: Metal GLM indexed batch attention value range is outside the mapped model\n");
return 0;
}
uint64_t value_inner = 0;
id<MTLBuffer> valuebuf =
ds4_gpu_wrap_model_range(model_map, model_size,
value_weight_offset,
value_weight_bytes,
&value_inner);
if (!valuebuf) return 0;
const uint64_t grouped_attn = (n_tokens >= 128u) ? 8u : 2u;
const NSUInteger q2_bit_words = ((NSUInteger)cache_cap + 31u) / 32u;
const NSUInteger q2_scratch_bytes =
q2_bit_words * sizeof(uint32_t) +
4u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) +
8u * (NSUInteger)kv_lora_dim * sizeof(float);
const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength];
const bool q2_fits = max_tg_mem == 0 || q2_scratch_bytes <= max_tg_mem;
const bool use_q2_group4 = grouped_attn == 4u && n_tokens >= 2u && q2_fits;
const bool use_group8 =
grouped_attn == 8u ||
(grouped_attn == 4u && !use_q2_group4);
const bool use_group2 = grouped_attn == 1u || grouped_attn == 2u;
id<MTLComputePipelineState> pipeline = use_q2_group4 ?
ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_q2_group4_pipeline,
"kernel_glm_attention_indexed_batch_q2_group4") :
(use_group8 ?
ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline,
"kernel_glm_attention_indexed_batch_group8") :
(use_group2 ?
ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group2_pipeline,
"kernel_glm_attention_indexed_batch_group2") :
ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_pipeline,
"kernel_glm_attention_indexed_batch")));
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_attention_indexed_batch_args args = {
.n_tokens = n_tokens,
.n_selected = n_selected,
.cache_cap = cache_cap,
.cache_f16 = cache_f16 ? 1u : 0u,
.n_head = n_head,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_rope = qk_rope,
.value_dim = value_dim,
.n_ctx_orig = n_ctx_orig,
.value_row_bytes = (uint32_t)value_row_bytes,
.value_type = value_weight_type,
.scale = 1.0f / sqrtf((float)qk_dim),
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
.head_base = 0,
};
const NSUInteger scratch_bytes = use_q2_group4 ?
q2_scratch_bytes :
(use_group8 ?
(8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) +
8u * (NSUInteger)kv_lora_dim * sizeof(float)) :
((use_group2 ?
(512u + 2u * (NSUInteger)n_selected + 2u * (NSUInteger)kv_lora_dim) :
(256u + (NSUInteger)n_selected + (NSUInteger)kv_lora_dim)) * sizeof(float)));
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
[enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
[enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
[enc setBuffer:valuebuf offset:(NSUInteger)value_inner atIndex:5];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6];
[enc setBuffer:headsbuf offset:ds4_gpu_tensor_offset(heads) atIndex:7];
[enc setThreadgroupMemoryLength:scratch_bytes atIndex:0];
if (use_q2_group4) {
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 3u) / 4u,
((NSUInteger)n_tokens + 1u) / 2u,
1)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
} else if (use_group8) {
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 7u) / 8u,
(NSUInteger)n_tokens,
1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
} else if (use_group2) {
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)n_head + 1u) / 2u,
(NSUInteger)n_tokens,
1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
} else {
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_head, (NSUInteger)n_tokens, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention")) return 0;
}
return 1;
}
int ds4_gpu_glm_attention_indexed_batch_tensor(
ds4_gpu_tensor *heads,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const void *model_map,
uint64_t model_size,
uint64_t value_weight_offset,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t value_dim,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
return ds4_gpu_glm_attention_indexed_batch_typed_tensor(heads,
q,
qk_low,
kv_lora_cache,
k_rope_cache,
model_map,
model_size,
value_weight_offset,
DS4_METAL_TENSOR_Q8_0,
selected,
n_tokens,
n_selected,
cache_cap,
cache_f16,
n_head,
kv_lora_dim,
qk_nope,
qk_rope,
value_dim,
n_ctx_orig,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow);
}
int ds4_gpu_sort_i32_rows_asc_tensor(
ds4_gpu_tensor *dst,
const ds4_gpu_tensor *src,
uint32_t row_width,
uint32_t n_rows) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!dst || !src || row_width == 0 || n_rows == 0 ||
(row_width & (row_width - 1u)) != 0) {
return 0;
}
@autoreleasepool {
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_dsv4_sort_i32_rows_asc_pipeline,
"kernel_dsv4_sort_i32_rows_asc");
if (!pipeline) return 0;
const uint64_t bytes = (uint64_t)row_width * n_rows * sizeof(int32_t);
id<MTLBuffer> srcbuf = ds4_gpu_tensor_buffer(src);
id<MTLBuffer> dstbuf = ds4_gpu_tensor_buffer(dst);
if (!srcbuf || !dstbuf ||
ds4_gpu_tensor_bytes(src) < bytes ||
ds4_gpu_tensor_bytes(dst) < bytes) {
fprintf(stderr, "ds4: Metal row sort received undersized buffers\n");
return 0;
}
NSUInteger threads = (NSUInteger)row_width;
const NSUInteger max_threads = pipeline.maxTotalThreadsPerThreadgroup;
if (max_threads != 0 && threads > max_threads) threads = max_threads;
if (threads == 0) return 0;
const NSUInteger scratch_bytes = (NSUInteger)row_width * sizeof(int32_t);
const NSUInteger max_tg_mem = [g_device maxThreadgroupMemoryLength];
if (max_tg_mem != 0 && scratch_bytes > max_tg_mem) {
fprintf(stderr, "ds4: Metal row sort scratch exceeds threadgroup memory limit\n");
return 0;
}
ds4_gpu_dsv4_topk_mask_args args = {
.ne00 = (int64_t)row_width,
.ne01 = (int64_t)n_rows,
.nb00 = sizeof(int32_t),
.nb01 = (uint64_t)row_width * sizeof(int32_t),
.ne0 = (int64_t)row_width,
.ne1 = (int64_t)n_rows,
.nb0 = sizeof(int32_t),
.nb1 = (uint64_t)row_width * sizeof(int32_t),
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:srcbuf offset:ds4_gpu_tensor_offset(src) atIndex:1];
[enc setBuffer:dstbuf offset:ds4_gpu_tensor_offset(dst) atIndex:2];
[enc setThreadgroupMemoryLength:scratch_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "sort i32 rows asc")) return 0;
}
return 1;
}
static int ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow,
bool selected_rows_valid) {
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t qk_dim = qk_nope + qk_rope;
if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache || !selected ||
n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
n_head == 0 || kv_lora_dim == 0 ||
qk_nope == 0 || (qk_rope & 1u) != 0 ||
qk_dim < qk_nope ||
(qk_rope != 0 && (!isfinite(freq_base) || freq_base <= 0.0f)) ||
(qk_rope != 0 && (!isfinite(freq_scale) || freq_scale <= 0.0f)) ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> lorabuf = ds4_gpu_tensor_buffer(lora_out);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
const uint64_t cache_elem_bytes = cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t lora_bytes =
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
const uint64_t q_bytes =
(uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
const uint64_t low_bytes =
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
const uint64_t kv_cache_bytes =
(uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
const uint64_t rope_cache_bytes =
(uint64_t)cache_cap * qk_rope * cache_elem_bytes;
const uint64_t selected_bytes =
(uint64_t)n_tokens * n_selected * sizeof(uint32_t);
if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf || !selectedbuf ||
ds4_gpu_tensor_bytes(lora_out) < lora_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes) {
fprintf(stderr, "ds4: Metal GLM indexed batch attention-lora received undersized buffers\n");
return 0;
}
const bool use_vec_lora =
cache_f16 && kv_lora_dim == 512u &&
(qk_rope == 0u || qk_rope == 64u);
const bool full_head_groups = (n_head % 8u) == 0u;
id<MTLComputePipelineState> pipeline = nil;
if (use_vec_lora && selected_rows_valid && full_head_groups) {
pipeline = ds4_gpu_hot_pipeline(
g_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads_pipeline,
"kernel_glm_attention_indexed_batch_lora_group8_vec_valid_fullheads");
} else if (use_vec_lora && selected_rows_valid) {
pipeline = ds4_gpu_hot_pipeline(
g_glm_attention_indexed_batch_lora_group8_vec_valid_pipeline,
"kernel_glm_attention_indexed_batch_lora_group8_vec_valid");
} else if (use_vec_lora) {
pipeline = ds4_gpu_hot_pipeline(
g_glm_attention_indexed_batch_lora_group8_vec_pipeline,
"kernel_glm_attention_indexed_batch_lora_group8_vec");
} else {
pipeline = ds4_gpu_hot_pipeline(g_glm_attention_indexed_batch_group8_pipeline,
"kernel_glm_attention_indexed_batch_group8");
}
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_attention_indexed_batch_args args = {
.n_tokens = n_tokens,
.n_selected = n_selected,
.cache_cap = cache_cap,
.cache_f16 = cache_f16 ? 1u : 0u,
.n_head = n_head,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_rope = qk_rope,
.value_dim = kv_lora_dim,
.n_ctx_orig = n_ctx_orig,
.value_row_bytes = 0,
.value_type = 1u,
.scale = 1.0f / sqrtf((float)qk_dim),
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
.head_base = 0,
};
uint32_t head_count = n_head;
ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count);
const NSUInteger scratch_bytes = use_vec_lora ?
(16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u +
16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u) :
(8u * ((NSUInteger)kv_lora_dim + (NSUInteger)qk_rope) * sizeof(uint16_t) +
8u * (NSUInteger)kv_lora_dim * sizeof(float));
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
[enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
[enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
if (use_vec_lora) {
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5];
[enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:6];
} else {
[enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:6];
[enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:7];
}
[enc setThreadgroupMemoryLength:scratch_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u,
(NSUInteger)n_tokens,
1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM indexed batch attention-lora")) return 0;
}
return 1;
}
int ds4_gpu_glm_attention_indexed_batch_lora_causal_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
uint32_t n_tokens,
uint32_t pos0,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t qk_dim = qk_nope + qk_rope;
if (!lora_out || !q || !qk_low || !kv_lora_cache || !k_rope_cache ||
n_tokens == 0 || n_selected == 0 || cache_cap == 0 || n_selected > cache_cap ||
pos0 > n_selected || n_tokens > n_selected - pos0 ||
n_head == 0 || kv_lora_dim != 512u ||
qk_nope == 0 || (qk_rope != 0u && qk_rope != 64u) ||
qk_dim < qk_nope ||
(qk_rope != 0u && (!isfinite(freq_base) || freq_base <= 0.0f)) ||
(qk_rope != 0u && (!isfinite(freq_scale) || freq_scale <= 0.0f)) ||
!isfinite(ext_factor) || !isfinite(attn_factor) ||
!isfinite(beta_fast) || !isfinite(beta_slow)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> lorabuf = ds4_gpu_tensor_buffer(lora_out);
id<MTLBuffer> qbuf = ds4_gpu_tensor_buffer(q);
id<MTLBuffer> lowbuf = ds4_gpu_tensor_buffer(qk_low);
id<MTLBuffer> kvcachebuf = ds4_gpu_tensor_buffer(kv_lora_cache);
id<MTLBuffer> ropecachebuf = ds4_gpu_tensor_buffer(k_rope_cache);
const uint64_t cache_elem_bytes =
cache_f16 ? sizeof(uint16_t) : sizeof(float);
const uint64_t lora_bytes =
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
const uint64_t q_bytes =
(uint64_t)n_tokens * n_head * qk_dim * sizeof(float);
const uint64_t low_bytes =
(uint64_t)n_tokens * n_head * kv_lora_dim * sizeof(float);
const uint64_t kv_cache_bytes =
(uint64_t)cache_cap * kv_lora_dim * cache_elem_bytes;
const uint64_t rope_cache_bytes =
(uint64_t)cache_cap * qk_rope * cache_elem_bytes;
if (!lorabuf || !qbuf || !lowbuf || !kvcachebuf || !ropecachebuf ||
ds4_gpu_tensor_bytes(lora_out) < lora_bytes ||
ds4_gpu_tensor_bytes(q) < q_bytes ||
ds4_gpu_tensor_bytes(qk_low) < low_bytes ||
ds4_gpu_tensor_bytes(kv_lora_cache) < kv_cache_bytes ||
ds4_gpu_tensor_bytes(k_rope_cache) < rope_cache_bytes) {
fprintf(stderr, "ds4: Metal GLM causal batch attention-lora received undersized buffers\n");
return 0;
}
const bool full_head_groups = (n_head % 8u) == 0u;
id<MTLComputePipelineState> pipeline = full_head_groups ?
ds4_gpu_hot_pipeline(
g_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads_pipeline,
"kernel_glm_attention_indexed_batch_lora_group8_vec_causal_fullheads") :
ds4_gpu_hot_pipeline(
g_glm_attention_indexed_batch_lora_group8_vec_causal_pipeline,
"kernel_glm_attention_indexed_batch_lora_group8_vec_causal");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_attention_indexed_batch_args args = {
.n_tokens = n_tokens,
.n_selected = n_selected,
.cache_cap = cache_cap,
.cache_f16 = cache_f16 ? 1u : 0u,
.n_head = n_head,
.kv_lora_dim = kv_lora_dim,
.qk_nope = qk_nope,
.qk_rope = qk_rope,
.value_dim = kv_lora_dim,
.n_ctx_orig = n_ctx_orig,
.value_row_bytes = 0,
.value_type = 1u,
.pos0 = pos0,
.scale = 1.0f / sqrtf((float)qk_dim),
.freq_base = freq_base,
.freq_scale = freq_scale,
.ext_factor = ext_factor,
.attn_factor = attn_factor,
.beta_fast = beta_fast,
.beta_slow = beta_slow,
.head_base = 0,
};
uint32_t head_count = n_head;
ds4_gpu_tp_attn_head_range(n_head, 8u, &args.head_base, &head_count);
const NSUInteger scratch_bytes =
16u * ((NSUInteger)kv_lora_dim / 4u) * sizeof(uint16_t) * 4u +
16u * ((NSUInteger)qk_rope / 4u) * sizeof(float) * 4u;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:qbuf offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:lowbuf offset:ds4_gpu_tensor_offset(qk_low) atIndex:2];
[enc setBuffer:kvcachebuf offset:ds4_gpu_tensor_offset(kv_lora_cache) atIndex:3];
[enc setBuffer:ropecachebuf offset:ds4_gpu_tensor_offset(k_rope_cache) atIndex:4];
[enc setBuffer:lorabuf offset:ds4_gpu_tensor_offset(lora_out) atIndex:5];
[enc setThreadgroupMemoryLength:scratch_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)head_count + 7u) / 8u,
(NSUInteger)n_tokens,
1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM causal indexed batch attention-lora")) return 0;
}
return 1;
}
int ds4_gpu_glm_attention_indexed_batch_lora_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out,
q,
qk_low,
kv_lora_cache,
k_rope_cache,
selected,
n_tokens,
n_selected,
cache_cap,
cache_f16,
n_head,
kv_lora_dim,
qk_nope,
qk_rope,
n_ctx_orig,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow,
false);
}
int ds4_gpu_glm_attention_indexed_batch_lora_valid_tensor(
ds4_gpu_tensor *lora_out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *qk_low,
const ds4_gpu_tensor *kv_lora_cache,
const ds4_gpu_tensor *k_rope_cache,
const ds4_gpu_tensor *selected,
uint32_t n_tokens,
uint32_t n_selected,
uint32_t cache_cap,
bool cache_f16,
uint32_t n_head,
uint32_t kv_lora_dim,
uint32_t qk_nope,
uint32_t qk_rope,
uint32_t n_ctx_orig,
float freq_base,
float freq_scale,
float ext_factor,
float attn_factor,
float beta_fast,
float beta_slow) {
return ds4_gpu_glm_attention_indexed_batch_lora_layout_tensor(lora_out,
q,
qk_low,
kv_lora_cache,
k_rope_cache,
selected,
n_tokens,
n_selected,
cache_cap,
cache_f16,
n_head,
kv_lora_dim,
qk_nope,
qk_rope,
n_ctx_orig,
freq_base,
freq_scale,
ext_factor,
attn_factor,
beta_fast,
beta_slow,
true);
}
int ds4_gpu_glm_router_select_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
const ds4_gpu_tensor *logits,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!selected || !weights || !probs || !logits || !model_map ||
n_expert == 0 || n_expert > 512u ||
n_expert_used == 0 || n_expert_used > n_expert) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf ||
ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int32_t) ||
ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) ||
ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) {
fprintf(stderr, "ds4: Metal GLM router received undersized buffers\n");
return 0;
}
const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
if (bias_offset > model_size || bias_bytes > model_size - bias_offset) {
fprintf(stderr, "ds4: Metal GLM router bias range is outside the mapped model\n");
return 0;
}
const bool exact_bias_view =
bias_bytes <= (1ull << 20) &&
getenv("DS4_METAL_DISABLE_DECODE_ROUTER_BIAS_EXACT_VIEWS") == NULL;
uint64_t bias_inner = 0;
id<MTLBuffer> biasbuf = exact_bias_view ?
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
bias_offset,
bias_bytes,
&bias_inner) :
ds4_gpu_wrap_model_range(model_map,
model_size,
bias_offset,
bias_bytes,
&bias_inner);
if (!biasbuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline,
"kernel_glm_router_select_one");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_router_select_one_args args = {
.n_expert = n_expert,
.n_expert_used = n_expert_used,
.expert_weight_scale = expert_weight_scale,
.pad0 = 0,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1];
[enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4];
[enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5];
const NSUInteger router_threads = n_expert > 256u ? 512u : 256u;
[enc setThreadgroupMemoryLength:router_threads *
(sizeof(float) + sizeof(int32_t)) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(router_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM router select")) return 0;
}
return 1;
}
int ds4_gpu_glm_router_select_batch_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
const ds4_gpu_tensor *logits,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!selected || !weights || !probs || !logits || !model_map ||
n_tokens == 0 ||
n_expert == 0 || n_expert > 512u ||
n_expert_used == 0 || n_expert_used > n_expert) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
const uint64_t logits_bytes = (uint64_t)n_tokens * n_expert * sizeof(float);
const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(int32_t);
const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert_used * sizeof(float);
const uint64_t probs_bytes = (uint64_t)n_tokens * n_expert * sizeof(float);
if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf ||
ds4_gpu_tensor_bytes(logits) < logits_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes ||
ds4_gpu_tensor_bytes(weights) < weights_bytes ||
ds4_gpu_tensor_bytes(probs) < probs_bytes) {
fprintf(stderr, "ds4: Metal GLM batch router received undersized buffers\n");
return 0;
}
const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
if (bias_offset > model_size || bias_bytes > model_size - bias_offset) {
fprintf(stderr, "ds4: Metal GLM batch router bias range is outside the mapped model\n");
return 0;
}
uint64_t bias_inner = 0;
id<MTLBuffer> biasbuf = ds4_gpu_wrap_model_range(model_map, model_size,
bias_offset, bias_bytes,
&bias_inner);
if (!biasbuf) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_hot_pipeline(g_glm_router_select_one_pipeline,
"kernel_glm_router_select_one");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_glm_router_select_one_args args = {
.n_expert = n_expert,
.n_expert_used = n_expert_used,
.expert_weight_scale = expert_weight_scale,
.pad0 = 0,
};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:logitsbuf offset:ds4_gpu_tensor_offset(logits) atIndex:1];
[enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:3];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:4];
[enc setBuffer:probsbuf offset:ds4_gpu_tensor_offset(probs) atIndex:5];
const NSUInteger router_threads = n_expert > 256u ? 512u : 256u;
[enc setThreadgroupMemoryLength:router_threads *
(sizeof(float) + sizeof(int32_t)) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_tokens, 1, 1)
threadsPerThreadgroup:MTLSizeMake(router_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM batch router select")) return 0;
}
return 1;
}
static bool ds4_gpu_glm_gate_pair_type_supported(
uint32_t gate_type,
uint32_t up_type) {
return gate_type == up_type &&
(gate_type == DS4_METAL_TENSOR_Q2_K ||
gate_type == DS4_METAL_TENSOR_Q4_K ||
gate_type == DS4_METAL_TENSOR_Q5_K);
}
static bool ds4_gpu_glm_down_type_supported(uint32_t down_type) {
return down_type == DS4_METAL_TENSOR_Q2_K ||
down_type == DS4_METAL_TENSOR_Q4_K ||
down_type == DS4_METAL_TENSOR_Q5_K ||
down_type == DS4_METAL_TENSOR_Q6_K;
}
int ds4_gpu_glm_routed_moe_one_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index,
const ds4_gpu_tensor *x,
bool force_resident) {
if (!g_initialized && !ds4_gpu_init()) return 0;
/* TP sharding: only the owned contiguous expert range is mapped,
* so bind from the owned base, validate only its bytes, and tell the
* kernels the first expert id present at that base. */
uint32_t first_expert = 0;
uint32_t n_bind_expert = 0;
ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
/* Streaming address tables point at individually cached experts by their
* global id. Keep the source offsets global; ownership is enforced by the
* kernels before they dereference an address-table slot. */
const bool streaming_tp =
g_ssd_streaming_mode && g_tp_split_world == 2;
if (streaming_tp) {
first_expert = 0;
n_bind_expert = n_total_expert;
}
const int32_t tp_expert_base_host = (int32_t)first_expert;
gate_offset += (uint64_t)first_expert * gate_expert_bytes;
up_offset += (uint64_t)first_expert * up_expert_bytes;
down_offset += (uint64_t)first_expert * down_expert_bytes;
if (!out || !mid || !model_map || !selected || !weights || !x ||
n_total_expert == 0 || n_expert == 0 || n_expert > 256u ||
n_expert > n_total_expert ||
expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 ||
gate_expert_bytes == 0 || gate_row_bytes == 0 ||
up_expert_bytes == 0 || up_row_bytes == 0 ||
down_expert_bytes == 0 || down_row_bytes == 0 ||
(expert_in_dim % 256u) != 0 ||
(expert_mid_dim % 256u) != 0 ||
!ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) ||
!ds4_gpu_glm_down_type_supported(down_type)) {
return 0;
}
if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes ||
(uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes ||
(uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes ||
(uint64_t)out_dim > UINT64_MAX / down_row_bytes) {
fprintf(stderr, "ds4: Metal GLM routed MoE tensor byte size overflow\n");
return 0;
}
const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes;
const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;
if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes ||
up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes ||
down_expert_bytes != (uint64_t)out_dim * down_row_bytes) {
fprintf(stderr, "ds4: Metal GLM routed MoE received inconsistent expert strides\n");
return 0;
}
ds4_gpu_stream_expert_cache_note_token(layer_index);
if (gate_offset > model_size || gate_tensor_bytes > model_size - gate_offset ||
up_offset > model_size || up_tensor_bytes > model_size - up_offset ||
down_offset > model_size || down_tensor_bytes > model_size - down_offset) {
fprintf(stderr, "ds4: Metal GLM routed MoE tensor range is outside the mapped model\n");
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float);
const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float);
if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(mid) < mid_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int32_t) ||
ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) {
fprintf(stderr, "ds4: Metal GLM routed MoE received undersized activation buffers\n");
return 0;
}
const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K;
const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K;
const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K;
const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K;
const BOOL down_simd_q4 = down_scalar_q4;
const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K;
const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K;
const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6;
const BOOL stream_addr_q2 =
gate_pair_q2 && down_scalar_q2 &&
g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil &&
g_glm_q2_k_addr_down_f32_pipeline != nil;
const BOOL stream_addr_q4 =
!gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 &&
g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil &&
g_glm_q4_k_addr_down_f32_pipeline != nil;
BOOL use_stream_expert_addr_table =
g_ssd_streaming_mode &&
!force_resident &&
getenv("DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE") == NULL &&
(stream_addr_q2 || stream_addr_q4) &&
layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER &&
n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
n_expert <= 8u &&
ds4_gpu_stream_expert_cache_configured_budget() >= n_expert &&
ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) &&
ds4_gpu_stream_expert_cache_effective_cap(layer_index,
n_total_expert,
n_expert) != 0;
int32_t stream_selected_ids[8] = {0, 0, 0, 0, 0, 0, 0, 0};
ds4_gpu_stream_expert_cache_entry *stream_entries[8] = {
NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL
};
uint64_t stream_gate_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint64_t stream_up_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint64_t stream_down_abs_offsets[8] = {0, 0, 0, 0, 0, 0, 0, 0};
uint32_t stream_missing_mask = 0;
uint32_t stream_entry_count = 0;
uint32_t stream_resident_mask = 0;
BOOL use_stream_split_deferred = false;
id<MTLBuffer> stream_gate_addr_buf = nil;
id<MTLBuffer> stream_up_addr_buf = nil;
id<MTLBuffer> stream_down_addr_buf = nil;
const int glm_stream_timing =
ds4_gpu_stream_expert_timing_summary_enabled();
if (use_stream_expert_addr_table) {
const int had_batch = g_batch_cb != nil;
int stream_ok = 1;
const int have_prefetched_selected =
ds4_gpu_glm_stream_selected_prefetch_take(model_map,
model_size,
layer_index,
n_total_expert,
n_expert,
gate_offset,
up_offset,
down_offset,
gate_expert_bytes,
down_expert_bytes,
stream_selected_ids);
if (have_prefetched_selected) {
if (had_batch && g_batch_has_work &&
g_stream_expert_pending_load.active &&
ds4_gpu_flush_commands() == 0) {
stream_ok = 0;
}
} else {
if (had_batch && ds4_gpu_end_commands() == 0) {
stream_ok = 0;
}
if (stream_ok &&
ds4_gpu_tensor_read(selected,
0,
stream_selected_ids,
(uint64_t)n_expert * sizeof(stream_selected_ids[0])) == 0) {
stream_ok = 0;
}
}
for (uint32_t i = 0; stream_ok && i < n_expert; i++) {
if (stream_selected_ids[i] < 0 ||
(uint32_t)stream_selected_ids[i] >= n_total_expert) {
fprintf(stderr,
"ds4: Metal GLM routed MoE selected expert id %d is outside 0..%u\n",
stream_selected_ids[i],
n_total_expert);
stream_ok = 0;
}
}
if (stream_ok) {
ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index,
stream_selected_ids,
n_expert);
if (!ds4_gpu_moe_selected_hotlist_record(layer_index,
stream_selected_ids,
n_expert,
n_total_expert)) {
stream_ok = 0;
}
}
if (stream_ok) {
g_glm_stream_expert_addr_table_building++;
for (uint32_t i = 0; stream_ok && i < n_expert; i++) {
const uint64_t expert_id = (uint64_t)(uint32_t)stream_selected_ids[i];
if (expert_id > UINT64_MAX / gate_expert_bytes ||
expert_id > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n");
stream_ok = 0;
break;
}
const uint64_t gate_rel = expert_id * gate_expert_bytes;
const uint64_t down_rel = expert_id * down_expert_bytes;
if (gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal GLM routed MoE selected expert offset overflow\n");
stream_ok = 0;
break;
}
stream_gate_abs_offsets[i] = gate_offset + gate_rel;
stream_up_abs_offsets[i] = up_offset + gate_rel;
stream_down_abs_offsets[i] = down_offset + down_rel;
stream_entries[i] =
ds4_gpu_stream_expert_cache_peek(model_map,
model_size,
layer_index,
(uint32_t)stream_selected_ids[i],
n_total_expert,
n_expert,
stream_gate_abs_offsets[i],
stream_up_abs_offsets[i],
stream_down_abs_offsets[i],
gate_expert_bytes,
down_expert_bytes);
if (!stream_entries[i]) {
stream_missing_mask |= 1u << i;
} else {
stream_resident_mask |= 1u << i;
}
}
if (stream_ok && glm_stream_timing) {
ds4_gpu_stream_expert_timing_note_cache_class(
stream_resident_mask,
stream_missing_mask);
}
use_stream_split_deferred =
stream_ok &&
getenv("DS4_METAL_DISABLE_GLM_STREAMING_EXPERT_SPLIT") == NULL &&
stream_missing_mask != 0 &&
stream_resident_mask != 0 &&
ds4_gpu_stream_expert_split_worthwhile(stream_resident_mask,
stream_missing_mask) &&
ds4_gpu_stream_expert_split_ready() &&
((gate_pair_q2 &&
g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline != nil) ||
(!gate_pair_q2 && !gate_pair_q5 &&
g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline != nil));
if (use_stream_split_deferred) {
const ds4_gpu_stream_expert_table table = {
.model_map = model_map,
.model_size = model_size,
.layer = layer_index,
.n_total_expert = n_total_expert,
.gate_offset = gate_offset,
.up_offset = up_offset,
.down_offset = down_offset,
.gate_expert_bytes = gate_expert_bytes,
.down_expert_bytes = down_expert_bytes,
};
if (!ds4_gpu_stream_expert_cache_begin_selected_load(
&table,
stream_selected_ids,
n_expert)) {
stream_ok = 0;
}
}
if (stream_ok && stream_missing_mask != 0 &&
!use_stream_split_deferred &&
!ds4_gpu_stream_expert_cache_load_selected_missing(
model_map,
model_size,
layer_index,
stream_selected_ids,
n_total_expert,
n_expert,
stream_gate_abs_offsets,
stream_up_abs_offsets,
stream_down_abs_offsets,
gate_expert_bytes,
down_expert_bytes,
stream_missing_mask,
stream_entries)) {
fprintf(stderr,
"ds4: Metal GLM streaming expert cache failed to load "
"layer=%u missing=0x%x budget=%u\n",
layer_index,
stream_missing_mask,
ds4_gpu_stream_expert_cache_configured_budget());
stream_ok = 0;
}
for (uint32_t i = 0; stream_ok && i < n_expert; i++) {
ds4_gpu_stream_expert_cache_entry *entry = stream_entries[i];
if (!entry) {
if (use_stream_split_deferred &&
(stream_missing_mask & (1u << i)) != 0) {
continue;
}
stream_ok = 0;
break;
}
if (!ds4_gpu_stream_expert_cache_set_addr_slot(
layer_index,
(uint32_t)stream_selected_ids[i],
entry->gate_buffer,
entry->gate_inner,
entry->up_buffer,
entry->up_inner,
entry->down_buffer,
entry->down_inner)) {
stream_ok = 0;
break;
}
stream_entry_count++;
}
if (stream_ok &&
!ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf)) {
stream_ok = 0;
}
g_glm_stream_expert_addr_table_building--;
}
if (!have_prefetched_selected && had_batch &&
ds4_gpu_begin_commands() == 0) {
stream_ok = 0;
}
if (!stream_ok) return 0;
ds4_gpu_stream_expert_cache_prune_layer(layer_index,
n_total_expert,
n_expert,
stream_selected_ids,
n_expert);
ds4_gpu_stream_expert_cache_prune_global(layer_index,
stream_selected_ids,
n_expert);
}
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
uint64_t down_inner = 0;
id<MTLBuffer> gatebuf = nil;
id<MTLBuffer> upbuf = nil;
id<MTLBuffer> downbuf = nil;
if (!use_stream_expert_addr_table) {
gatebuf = ds4_gpu_wrap_model_range(model_map, model_size,
gate_offset, gate_tensor_bytes,
&gate_inner);
upbuf = ds4_gpu_wrap_model_range(model_map, model_size,
up_offset, up_tensor_bytes,
&up_inner);
downbuf = ds4_gpu_wrap_model_range(model_map, model_size,
down_offset, down_tensor_bytes,
&down_inner);
if (!gatebuf || !upbuf || !downbuf) return 0;
}
id<MTLComputePipelineState> pair_pipeline =
(use_stream_split_deferred ?
(gate_pair_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_masked_f32_pipeline,
"kernel_glm_q2_K_addr_pair_swiglu2_f32_masked") :
ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_masked_f32_pipeline,
"kernel_glm_q4_K_addr_pair_swiglu_f32_masked")) :
use_stream_expert_addr_table ?
(gate_pair_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline,
"kernel_glm_q2_K_addr_pair_swiglu2_f32") :
ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline,
"kernel_glm_q4_K_addr_pair_swiglu_f32")) :
gate_pair_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline,
"kernel_glm_q2_K_pair_swiglu_f32") :
gate_pair_q5 ?
ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline,
"kernel_glm_q5_K_pair_swiglu_f32") :
ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline,
"kernel_glm_q4_K_pair_swiglu2_f32"));
id<MTLComputePipelineState> down_pipeline =
(use_stream_expert_addr_table ?
(down_scalar_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline,
"kernel_glm_q2_K_addr_down_f32") :
ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline,
"kernel_glm_q4_K_addr_down_f32")) :
down_scalar_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline,
"kernel_glm_q2_K_down_f32") :
down_scalar_q4 ?
ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline,
"kernel_glm_q4_K_down_f32") :
down_simd_q5 ?
ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline,
"kernel_glm_q5_K_down_f32") :
ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline,
"kernel_glm_q6_K_down_f32"));
if (!pair_pipeline || !down_pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const bool glm_moe_stage_profile =
g_batch_cb != nil &&
ds4_gpu_stage_profile_enabled_for_layer(
"DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE",
"DS4_METAL_GLM_MOE_ONE_STAGE_PROFILE_LAYER",
layer_index);
const char *glm_moe_stage_filter =
getenv("DS4_METAL_GLM_MOE_STAGE_PROFILE_FILTER");
const char *glm_pair_path =
use_stream_expert_addr_table ?
(gate_pair_q2 ? "q2_stream_addr_swiglu" :
"q4_stream_addr_swiglu") :
gate_pair_q2 ? "q2_scalar_swiglu" :
gate_pair_q5 ? "q5_pair_simd_swiglu" : "q4_pair2_simd_swiglu";
const char *glm_down_path =
use_stream_expert_addr_table ?
(down_scalar_q2 ? "q2_stream_addr_down" :
"q4_stream_addr_down_simd") :
down_scalar_q2 ? "q2_down_simd" :
down_scalar_q4 ? "q4_down_simd" :
down_simd_q5 ? "q5_down_simd" : "q6_down_simd";
double glm_moe_stage_t0 = 0.0;
if (glm_moe_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
glm_moe_stage_t0 = ds4_gpu_now_ms();
}
int ok = 1;
#define DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE(name) do { \
if (ok && glm_moe_stage_profile) { \
if (ds4_gpu_end_commands() == 0) { \
ok = 0; \
} else { \
const char *stage_name = (name); \
const double now_ms = ds4_gpu_now_ms(); \
const int print_stage = \
!glm_moe_stage_filter || !glm_moe_stage_filter[0] || \
strstr(stage_name, glm_moe_stage_filter) != NULL; \
if (print_stage) { \
fprintf(stderr, \
"ds4: Metal GLM routed MoE one stage layer=%u tokens=1 experts=%u " \
"gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \
layer_index, n_expert, \
ds4_gpu_metal_tensor_type_name(gate_type), \
ds4_gpu_metal_tensor_type_name(down_type), \
glm_pair_path, glm_down_path, \
stage_name, now_ms - glm_moe_stage_t0); \
} \
glm_moe_stage_t0 = now_ms; \
if (ds4_gpu_begin_commands() == 0) { \
ok = 0; \
} else { \
cb = ds4_gpu_command_buffer(&owned); \
if (!cb) ok = 0; \
} \
} \
} \
} while (0)
ds4_gpu_glm_routed_moe_args args = {
.tp_rank = g_tp_split_rank,
.tp_world = g_tp_split_world,
.tp_expert_base = tp_expert_base_host,
.in_dim = expert_in_dim,
.mid_dim = expert_mid_dim,
.out_dim = out_dim,
.n_total_expert = n_total_expert,
.n_expert_used = n_expert,
.n_tokens = 1,
.mid_token_stride = n_expert * expert_mid_dim,
.down_type = down_type,
.swiglu_clamp = swiglu_clamp,
.gate_expert_bytes = gate_expert_bytes,
.gate_row_bytes = gate_row_bytes,
.up_expert_bytes = up_expert_bytes,
.up_row_bytes = up_row_bytes,
.down_expert_bytes = down_expert_bytes,
.down_row_bytes = down_row_bytes,
};
const NSUInteger pair_x_groups =
gate_pair_q2 ? (use_stream_expert_addr_table ?
(NSUInteger)((expert_mid_dim + 1u) / 2u) :
(NSUInteger)((expert_mid_dim + 7u) / 8u)) :
gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) :
use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) :
(NSUInteger)((expert_mid_dim + 1u) / 2u);
const NSUInteger pair_threadgroup_bytes = 0u;
const NSUInteger pair_threads = 64u;
const NSUInteger down_x_groups =
down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) :
down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) :
down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) :
down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) :
(NSUInteger)out_dim;
const NSUInteger down_threadgroup_bytes =
(down_scalar_q2 || down_simd) ? 0u : 256u * sizeof(float);
const NSUInteger down_threads =
(down_scalar_q2 || down_simd) ? 64u : 256u;
if (use_stream_expert_addr_table &&
!ds4_gpu_stream_expert_cache_mark_entries_inflight(
stream_entries,
use_stream_split_deferred ? n_expert : stream_entry_count,
use_stream_split_deferred ? stream_resident_mask : 0)) {
return 0;
}
const int glm_stream_split_timing =
use_stream_split_deferred && glm_stream_timing;
double glm_stream_split_t0 =
glm_stream_split_timing ? ds4_gpu_now_ms() : 0.0;
double glm_stream_split_resident_ms = 0.0;
double glm_stream_split_missing_load_ms = 0.0;
double glm_stream_split_missing_wait_ms = 0.0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pair_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
if (use_stream_split_deferred) {
[enc setBytes:&stream_resident_mask length:sizeof(stream_resident_mask) atIndex:1];
[enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2];
[enc setBuffer:stream_up_addr_buf offset:0u atIndex:3];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7];
} else {
[enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf
offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner
atIndex:1];
[enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf
offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner
atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
}
if (use_stream_expert_addr_table) {
const uint32_t use_count =
use_stream_split_deferred ? n_expert : stream_entry_count;
const uint32_t use_mask =
use_stream_split_deferred ? stream_resident_mask : 0;
for (uint32_t i = 0; i < use_count; i++) {
if (use_mask != 0 && (use_mask & (1u << i)) == 0) continue;
[enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead];
[enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead];
}
}
if (pair_threadgroup_bytes != 0u) {
[enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(pair_x_groups,
(NSUInteger)n_expert,
1)
threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("pair");
if (ok && use_stream_split_deferred) {
id<MTLCommandBuffer> resident_cb = nil;
const int resident_owned = owned;
if (owned) {
resident_cb = cb;
[resident_cb commit];
cb = nil;
} else {
ok = ds4_gpu_flush_commands();
}
if (glm_stream_split_timing) {
const double now_ms = ds4_gpu_now_ms();
glm_stream_split_resident_ms = now_ms - glm_stream_split_t0;
glm_stream_split_t0 = now_ms;
}
if (ok) {
ok = ds4_gpu_stream_expert_pending_load_finish(stream_entries);
}
if (glm_stream_split_timing) {
const double now_ms = ds4_gpu_now_ms();
glm_stream_split_missing_load_ms = now_ms - glm_stream_split_t0;
glm_stream_split_t0 = now_ms;
}
if (ok) {
ds4_gpu_stream_expert_cache_prune_layer(layer_index,
n_total_expert,
n_expert,
stream_selected_ids,
n_expert);
ds4_gpu_stream_expert_cache_prune_global(layer_index,
stream_selected_ids,
n_expert);
ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf);
}
if (ok) {
if (resident_owned) {
ok = ds4_gpu_wait_command_buffer(
resident_cb,
"GLM streaming split resident pair");
ds4_gpu_stream_expert_cache_note_owned_completed();
} else {
ok = ds4_gpu_wait_pending_command_buffers(
"GLM streaming split resident pair");
}
}
if (glm_stream_split_timing) {
const double now_ms = ds4_gpu_now_ms();
glm_stream_split_missing_wait_ms = now_ms - glm_stream_split_t0;
glm_stream_split_t0 = now_ms;
ds4_gpu_stream_expert_timing_note_split(
stream_resident_mask,
stream_missing_mask,
glm_stream_split_resident_ms,
glm_stream_split_missing_load_ms +
glm_stream_split_missing_wait_ms);
ds4_gpu_stream_expert_timing_note_split_missing_detail(
glm_stream_split_missing_load_ms,
0.0,
0.0,
0.0,
glm_stream_split_missing_wait_ms);
}
if (ok &&
!ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_entries,
n_expert,
stream_missing_mask)) {
ok = 0;
}
if (ok) {
cb = ds4_gpu_command_buffer(&owned);
if (!cb) ok = 0;
}
if (ok) {
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pair_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBytes:&stream_missing_mask length:sizeof(stream_missing_mask) atIndex:1];
[enc setBuffer:stream_gate_addr_buf offset:0u atIndex:2];
[enc setBuffer:stream_up_addr_buf offset:0u atIndex:3];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:4];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:5];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:6];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:7];
for (uint32_t i = 0; i < n_expert; i++) {
if ((stream_missing_mask & (1u << i)) == 0) continue;
[enc useResource:stream_entries[i]->gate_buffer usage:MTLResourceUsageRead];
[enc useResource:stream_entries[i]->up_buffer usage:MTLResourceUsageRead];
}
if (pair_threadgroup_bytes != 0u) {
[enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(pair_x_groups,
(NSUInteger)n_expert,
1)
threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
}
if (!ok) return 0;
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:down_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf
offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner
atIndex:1];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
if (use_stream_expert_addr_table) {
const uint32_t use_count =
use_stream_split_deferred ? n_expert : stream_entry_count;
for (uint32_t i = 0; i < use_count; i++) {
[enc useResource:stream_entries[i]->down_buffer usage:MTLResourceUsageRead];
}
}
if (down_threadgroup_bytes != 0u) {
[enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(down_x_groups, 1, 1)
threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE("down");
if (!ok) return 0;
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed MoE")) return 0;
#undef DS4_METAL_PROFILE_GLM_MOE_ONE_STAGE
}
return 1;
}
static bool ds4_gpu_glm_grouped_moe_fast_default(void) {
return !g_quality_mode;
}
static bool ds4_gpu_glm_routed_moe_batch_grouped_available(
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint32_t n_expert,
uint32_t n_tokens) {
if (n_tokens < 32u ||
!ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) ||
!ds4_gpu_glm_down_type_supported(down_type) ||
ds4_gpu_mul_mm_id_map0_name(n_expert) == NULL) {
return false;
}
const bool fast_default = ds4_gpu_glm_grouped_moe_fast_default();
if (!fast_default) {
return false;
}
if (n_tokens < 96u) return false;
return ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert)) != nil &&
ds4_gpu_routed_mm_pipeline(gate_type) != nil &&
ds4_gpu_routed_mm_pipeline(up_type) != nil &&
ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) != nil;
}
static bool ds4_gpu_glm_grouped_moe_layer_enabled(uint32_t layer_index) {
(void)layer_index;
return true;
}
static int ds4_gpu_glm_routed_moe_batch_grouped_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens) {
if (!ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type,
up_type,
down_type,
n_expert,
n_tokens)) {
return 0;
}
if (n_expert > UINT32_MAX / n_tokens ||
(uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) {
return 0;
}
const uint32_t pair_rows = n_tokens * n_expert;
if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim ||
(uint64_t)pair_rows > UINT64_MAX / out_dim ||
(uint64_t)n_tokens > UINT64_MAX / expert_in_dim ||
(uint64_t)n_tokens > UINT64_MAX / out_dim) {
return 0;
}
const bool mid_f16 = true;
const NSUInteger mm_id_threadgroup_bytes = 8192u;
const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim;
const uint64_t down_values = (uint64_t)pair_rows * out_dim;
const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim;
const uint64_t out_values = (uint64_t)n_tokens * out_dim;
if (compact_mid_values > UINT64_MAX / sizeof(float) ||
compact_mid_values > UINT64_MAX / (mid_f16 ? sizeof(uint16_t) : sizeof(float)) ||
down_values > UINT64_MAX / sizeof(float) ||
x_values > UINT64_MAX / sizeof(float) ||
out_values > UINT64_MAX / sizeof(float)) {
return 0;
}
const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float);
const uint64_t mid_bytes = compact_mid_values * (mid_f16 ? sizeof(uint16_t) : sizeof(float));
const uint64_t down_scratch_bytes = down_values * sizeof(float);
const uint64_t x_bytes = x_values * sizeof(float);
const uint64_t out_bytes = out_values * sizeof(float);
const uint64_t selected_values = (uint64_t)n_tokens * n_expert;
const uint64_t selected_bytes = selected_values * sizeof(int32_t);
const uint64_t weights_bytes = selected_values * sizeof(float);
if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes ||
gate_scratch_bytes > NSUIntegerMax ||
gate_scratch_bytes * 2ull > NSUIntegerMax ||
down_scratch_bytes > NSUIntegerMax) {
return 0;
}
uint32_t first_expert = 0;
uint32_t n_bind_expert = 0;
ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
gate_offset += (uint64_t)first_expert * gate_expert_bytes;
up_offset += (uint64_t)first_expert * up_expert_bytes;
down_offset += (uint64_t)first_expert * down_expert_bytes;
const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes;
const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(mid) < mid_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes ||
ds4_gpu_tensor_bytes(weights) < weights_bytes) {
fprintf(stderr, "ds4: Metal GLM grouped routed MoE received undersized activation buffers\n");
return 0;
}
if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer,
&g_moe_gate_scratch_bytes,
(NSUInteger)(gate_scratch_bytes * 2ull),
"ds4_glm_moe_gate_up_scratch")) {
return 0;
}
if (n_expert > 1 &&
!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
&g_moe_down_scratch_bytes,
(NSUInteger)down_scratch_bytes,
"ds4_glm_moe_down_scratch")) {
return 0;
}
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
uint64_t down_inner = 0;
id<MTLBuffer> gatebuf = ds4_gpu_wrap_model_range(model_map, model_size,
gate_offset, gate_tensor_bytes,
&gate_inner);
id<MTLBuffer> upbuf = ds4_gpu_wrap_model_range(model_map, model_size,
up_offset, up_tensor_bytes,
&up_inner);
id<MTLBuffer> downbuf = ds4_gpu_wrap_model_range(model_map, model_size,
down_offset, down_tensor_bytes,
&down_inner);
if (!gatebuf || !upbuf || !downbuf) return 0;
id<MTLComputePipelineState> map_pipeline =
ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert));
id<MTLComputePipelineState> gate_pipeline = ds4_gpu_routed_mm_pipeline(gate_type);
id<MTLComputePipelineState> up_pipeline = ds4_gpu_routed_mm_pipeline(up_type);
id<MTLComputePipelineState> down_pipeline =
ds4_gpu_routed_mm_f16_rhs_pipeline(down_type);
if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) {
return 0;
}
ds4_gpu_mul_mm_id_map_args map_args =
ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens);
ds4_gpu_mul_mm_id_args gate_args =
ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
gate_row_bytes, gate_expert_bytes,
1, n_expert, n_tokens);
ds4_gpu_mul_mm_id_args up_args =
ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
up_row_bytes, up_expert_bytes,
1, n_expert, n_tokens);
ds4_gpu_mul_mm_id_args down_args =
ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert,
down_row_bytes, down_expert_bytes,
n_expert, n_expert, n_tokens,
mid_f16 ? sizeof(uint16_t) : sizeof(float));
gate_args.tp_rank = g_tp_split_rank;
gate_args.tp_world = g_tp_split_world;
gate_args.tp_expert_base = (int32_t)first_expert;
up_args.tp_rank = g_tp_split_rank;
up_args.tp_world = g_tp_split_world;
up_args.tp_expert_base = (int32_t)first_expert;
down_args.tp_rank = g_tp_split_rank;
down_args.tp_world = g_tp_split_world;
down_args.tp_expert_base = (int32_t)first_expert;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const bool glm_moe_stage_profile = false;
const char *glm_moe_stage_filter = NULL;
double glm_moe_stage_t0 = 0.0;
if (glm_moe_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
glm_moe_stage_t0 = ds4_gpu_now_ms();
}
int ok = 1;
#define DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE(name) do { \
if (ok && glm_moe_stage_profile) { \
if (ds4_gpu_end_commands() == 0) { \
ok = 0; \
} else { \
const char *stage_name = (name); \
const double now_ms = ds4_gpu_now_ms(); \
const int print_stage = \
!glm_moe_stage_filter || !glm_moe_stage_filter[0] || \
strstr(stage_name, glm_moe_stage_filter) != NULL; \
if (print_stage) { \
fprintf(stderr, \
"ds4: Metal GLM grouped routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \
"gate=%s down=%s mid=%s %s=%.3f ms\n", \
layer_index, n_tokens, pair_rows, n_expert, \
ds4_gpu_metal_tensor_type_name(gate_type), \
ds4_gpu_metal_tensor_type_name(down_type), \
mid_f16 ? "f16" : "f32", \
stage_name, now_ms - glm_moe_stage_t0); \
} \
glm_moe_stage_t0 = now_ms; \
if (ds4_gpu_begin_commands() == 0) { \
ok = 0; \
} else { \
cb = ds4_gpu_command_buffer(&owned); \
if (!cb) ok = 0; \
} \
} \
} \
} while (0)
ok = ds4_gpu_encode_mul_mm_id_map(cb,
map_pipeline,
&map_args,
&gate_args,
selectedbuf,
ds4_gpu_tensor_offset(selected));
DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("map");
if (ok) {
ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
gate_pipeline,
&gate_args,
gatebuf,
(NSUInteger)gate_inner,
xbuf,
ds4_gpu_tensor_offset(x),
g_moe_gate_scratch_buffer,
0,
mm_id_threadgroup_bytes);
}
DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("gate");
if (ok) {
ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
up_pipeline,
&up_args,
upbuf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
g_moe_gate_scratch_buffer,
(NSUInteger)gate_scratch_bytes,
mm_id_threadgroup_bytes);
}
DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("up");
if (ok) {
ok = ds4_gpu_encode_moe_swiglu_weight(cb,
g_moe_gate_scratch_buffer,
0,
g_moe_gate_scratch_buffer,
(NSUInteger)gate_scratch_bytes,
midbuf,
ds4_gpu_tensor_offset(mid),
weightsbuf,
ds4_gpu_tensor_offset(weights),
expert_mid_dim,
pair_rows,
swiglu_clamp,
mid_f16);
}
DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("activation_weight");
id<MTLBuffer> down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer;
NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0;
if (ok) {
ok = ds4_gpu_encode_mul_mm_id_mapped_tile(cb,
down_pipeline,
&down_args,
downbuf,
(NSUInteger)down_inner,
midbuf,
ds4_gpu_tensor_offset(mid),
down_dst,
down_dst_off,
mm_id_threadgroup_bytes);
}
DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("down");
if (ok && n_expert > 1) {
ok = ds4_gpu_encode_moe_sum_experts(cb,
down_dst,
down_dst_off,
outbuf,
ds4_gpu_tensor_offset(out),
out_dim,
n_expert,
n_tokens);
}
DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE("sum");
if (!ok) return 0;
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM grouped routed batch MoE")) {
return 0;
}
#undef DS4_METAL_PROFILE_GLM_GROUPED_MOE_STAGE
}
return 1;
}
static int ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
id<MTLBuffer> gate_addrs,
id<MTLBuffer> up_addrs,
id<MTLBuffer> down_addrs,
id<MTLBuffer> overflow_gate,
id<MTLBuffer> overflow_up,
id<MTLBuffer> overflow_down,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens,
ds4_gpu_stream_expert_cache_entry * const *resources,
uint32_t resource_count) {
if (n_expert > UINT32_MAX / n_tokens ||
!gate_addrs || !up_addrs || !down_addrs ||
!ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type,
up_type,
down_type,
n_expert,
n_tokens)) {
return 0;
}
const uint32_t pair_rows = n_tokens * n_expert;
if ((uint64_t)pair_rows > UINT64_MAX / expert_mid_dim ||
(uint64_t)pair_rows > UINT64_MAX / out_dim ||
(uint64_t)n_tokens > UINT64_MAX / expert_in_dim ||
(uint64_t)n_tokens > UINT64_MAX / out_dim) {
return 0;
}
const bool mid_f16 = true;
const NSUInteger mm_id_threadgroup_bytes = 8192u;
const uint64_t compact_mid_values = (uint64_t)pair_rows * expert_mid_dim;
const uint64_t down_values = (uint64_t)pair_rows * out_dim;
const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim;
const uint64_t out_values = (uint64_t)n_tokens * out_dim;
if (compact_mid_values > UINT64_MAX / sizeof(float) ||
compact_mid_values > UINT64_MAX / sizeof(uint16_t) ||
down_values > UINT64_MAX / sizeof(float) ||
x_values > UINT64_MAX / sizeof(float) ||
out_values > UINT64_MAX / sizeof(float)) {
return 0;
}
const uint64_t gate_scratch_bytes = compact_mid_values * sizeof(float);
const uint64_t mid_bytes = compact_mid_values * sizeof(uint16_t);
const uint64_t down_scratch_bytes = down_values * sizeof(float);
const uint64_t x_bytes = x_values * sizeof(float);
const uint64_t out_bytes = out_values * sizeof(float);
const uint64_t selected_values = (uint64_t)n_tokens * n_expert;
const uint64_t selected_bytes = selected_values * sizeof(int32_t);
const uint64_t weights_bytes = selected_values * sizeof(float);
if (gate_scratch_bytes > UINT64_MAX - gate_scratch_bytes ||
gate_scratch_bytes > NSUIntegerMax ||
gate_scratch_bytes * 2ull > NSUIntegerMax ||
down_scratch_bytes > NSUIntegerMax) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(mid) < mid_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes ||
ds4_gpu_tensor_bytes(weights) < weights_bytes) {
fprintf(stderr, "ds4: Metal GLM grouped-address routed MoE received undersized activation buffers\n");
return 0;
}
if (!ds4_gpu_ensure_scratch_buffer(&g_moe_gate_scratch_buffer,
&g_moe_gate_scratch_bytes,
(NSUInteger)(gate_scratch_bytes * 2ull),
"ds4_glm_moe_gate_up_scratch")) {
return 0;
}
if (n_expert > 1 &&
!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
&g_moe_down_scratch_bytes,
(NSUInteger)down_scratch_bytes,
"ds4_glm_moe_down_scratch")) {
return 0;
}
id<MTLComputePipelineState> map_pipeline =
ds4_gpu_get_pipeline(ds4_gpu_mul_mm_id_map0_name(n_expert));
id<MTLComputePipelineState> gate_pipeline =
ds4_gpu_routed_mm_addr_pipeline(gate_type);
id<MTLComputePipelineState> up_pipeline =
ds4_gpu_routed_mm_addr_pipeline(up_type);
id<MTLComputePipelineState> down_pipeline =
ds4_gpu_routed_mm_addr_f16_rhs_pipeline(down_type);
if (!map_pipeline || !gate_pipeline || !up_pipeline || !down_pipeline) {
return 0;
}
ds4_gpu_mul_mm_id_map_args map_args =
ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens);
ds4_gpu_mul_mm_id_args gate_args =
ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
gate_row_bytes, gate_expert_bytes,
1, n_expert, n_tokens);
ds4_gpu_mul_mm_id_args up_args =
ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
up_row_bytes, up_expert_bytes,
1, n_expert, n_tokens);
ds4_gpu_mul_mm_id_args down_args =
ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert,
down_row_bytes, down_expert_bytes,
n_expert, n_expert, n_tokens,
sizeof(uint16_t));
if (!ds4_gpu_stream_expert_cache_mark_entries_inflight(resources,
resource_count,
0)) {
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const bool glm_moe_stage_profile = false;
const char *glm_moe_stage_filter = NULL;
double glm_moe_stage_t0 = 0.0;
if (glm_moe_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
glm_moe_stage_t0 = ds4_gpu_now_ms();
}
int ok = 1;
#define DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE(name) do { \
if (ok && glm_moe_stage_profile) { \
if (ds4_gpu_end_commands() == 0) { \
ok = 0; \
} else { \
const char *stage_name = (name); \
const double now_ms = ds4_gpu_now_ms(); \
const int print_stage = \
!glm_moe_stage_filter || !glm_moe_stage_filter[0] || \
strstr(stage_name, glm_moe_stage_filter) != NULL; \
if (print_stage) { \
fprintf(stderr, \
"ds4: Metal GLM grouped-address routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \
"gate=%s down=%s mid=f16 %s=%.3f ms\n", \
layer_index, n_tokens, pair_rows, n_expert, \
ds4_gpu_metal_tensor_type_name(gate_type), \
ds4_gpu_metal_tensor_type_name(down_type), \
stage_name, now_ms - glm_moe_stage_t0); \
} \
glm_moe_stage_t0 = now_ms; \
if (ds4_gpu_begin_commands() == 0) { \
ok = 0; \
} else { \
cb = ds4_gpu_command_buffer(&owned); \
if (!cb) ok = 0; \
} \
} \
} \
} while (0)
ok = ds4_gpu_encode_mul_mm_id_map(cb,
map_pipeline,
&map_args,
&gate_args,
selectedbuf,
ds4_gpu_tensor_offset(selected));
DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("map");
if (ok) {
ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb,
gate_pipeline,
&gate_args,
gate_addrs,
xbuf,
ds4_gpu_tensor_offset(x),
g_moe_gate_scratch_buffer,
0,
mm_id_threadgroup_bytes,
resources,
resource_count,
0,
overflow_gate);
}
DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("gate");
if (ok) {
ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb,
up_pipeline,
&up_args,
up_addrs,
xbuf,
ds4_gpu_tensor_offset(x),
g_moe_gate_scratch_buffer,
(NSUInteger)gate_scratch_bytes,
mm_id_threadgroup_bytes,
resources,
resource_count,
1,
overflow_up);
}
DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("up");
if (ok) {
ok = ds4_gpu_encode_moe_swiglu_weight(cb,
g_moe_gate_scratch_buffer,
0,
g_moe_gate_scratch_buffer,
(NSUInteger)gate_scratch_bytes,
midbuf,
ds4_gpu_tensor_offset(mid),
weightsbuf,
ds4_gpu_tensor_offset(weights),
expert_mid_dim,
pair_rows,
swiglu_clamp,
mid_f16);
}
DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("activation_weight");
id<MTLBuffer> down_dst = n_expert == 1 ? outbuf : g_moe_down_scratch_buffer;
NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) : 0;
if (ok) {
ok = ds4_gpu_encode_mul_mm_id_addr_mapped_tile(cb,
down_pipeline,
&down_args,
down_addrs,
midbuf,
ds4_gpu_tensor_offset(mid),
down_dst,
down_dst_off,
mm_id_threadgroup_bytes,
resources,
resource_count,
2,
overflow_down);
}
DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("down");
if (ok && n_expert > 1) {
ok = ds4_gpu_encode_moe_sum_experts(cb,
down_dst,
down_dst_off,
outbuf,
ds4_gpu_tensor_offset(out),
out_dim,
n_expert,
n_tokens);
}
DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE("sum");
if (!ok) return 0;
if (!ds4_gpu_finish_command_buffer(cb, owned,
"GLM grouped-address routed batch MoE")) {
return 0;
}
#undef DS4_METAL_PROFILE_GLM_GROUPED_ADDR_MOE_STAGE
}
return 1;
}
static int ds4_gpu_glm_routed_moe_batch_tensor_impl(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t mid_token_stride,
bool allow_grouped,
bool force_scalar_q4_pair) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !mid || !model_map || !selected || !weights || !x ||
n_tokens == 0 ||
n_total_expert == 0 || n_expert == 0 || n_expert > 256u ||
n_expert > n_total_expert ||
expert_in_dim == 0 || expert_mid_dim == 0 || out_dim == 0 ||
gate_expert_bytes == 0 || gate_row_bytes == 0 ||
up_expert_bytes == 0 || up_row_bytes == 0 ||
down_expert_bytes == 0 || down_row_bytes == 0 ||
(expert_in_dim % 256u) != 0 ||
(expert_mid_dim % 256u) != 0 ||
!ds4_gpu_glm_gate_pair_type_supported(gate_type, up_type) ||
!ds4_gpu_glm_down_type_supported(down_type)) {
return 0;
}
const uint64_t per_token_mid = (uint64_t)n_expert * expert_mid_dim;
if ((uint64_t)mid_token_stride < per_token_mid ||
(uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / up_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes ||
(uint64_t)expert_mid_dim > UINT64_MAX / gate_row_bytes ||
(uint64_t)expert_mid_dim > UINT64_MAX / up_row_bytes ||
(uint64_t)out_dim > UINT64_MAX / down_row_bytes ||
(uint64_t)n_tokens > UINT64_MAX / expert_in_dim ||
(uint64_t)n_tokens > UINT64_MAX / out_dim) {
fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor byte size overflow\n");
return 0;
}
const uint64_t full_gate_tensor_bytes =
(uint64_t)n_total_expert * gate_expert_bytes;
const uint64_t full_up_tensor_bytes =
(uint64_t)n_total_expert * up_expert_bytes;
const uint64_t full_down_tensor_bytes =
(uint64_t)n_total_expert * down_expert_bytes;
if (gate_expert_bytes != (uint64_t)expert_mid_dim * gate_row_bytes ||
up_expert_bytes != (uint64_t)expert_mid_dim * up_row_bytes ||
down_expert_bytes != (uint64_t)out_dim * down_row_bytes) {
fprintf(stderr, "ds4: Metal GLM routed batch MoE received inconsistent expert strides\n");
return 0;
}
if (gate_offset > model_size || full_gate_tensor_bytes > model_size - gate_offset ||
up_offset > model_size || full_up_tensor_bytes > model_size - up_offset ||
down_offset > model_size || full_down_tensor_bytes > model_size - down_offset) {
fprintf(stderr, "ds4: Metal GLM routed batch MoE tensor range is outside the mapped model\n");
return 0;
}
if (allow_grouped &&
(!g_ssd_streaming_mode ||
ds4_gpu_glm_streaming_prefill_full_layer_active()) &&
ds4_gpu_glm_grouped_moe_layer_enabled(layer_index) &&
ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type,
up_type,
down_type,
n_expert,
n_tokens)) {
return ds4_gpu_glm_routed_moe_batch_grouped_tensor(out,
mid,
model_map,
model_size,
gate_offset,
up_offset,
down_offset,
gate_type,
up_type,
down_type,
gate_expert_bytes,
gate_row_bytes,
up_expert_bytes,
up_row_bytes,
down_expert_bytes,
down_row_bytes,
expert_in_dim,
expert_mid_dim,
out_dim,
selected,
weights,
n_total_expert,
n_expert,
swiglu_clamp,
layer_index,
x,
n_tokens);
}
uint32_t first_expert = 0;
uint32_t n_bind_expert = 0;
ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
/* Streaming address tables are indexed by global expert id. As in the
* one-token path, keep model offsets global and let the kernels enforce
* rank ownership before dereferencing the selected address. */
const bool streaming_tp =
g_ssd_streaming_mode && g_tp_split_world == 2;
if (streaming_tp) {
first_expert = 0;
n_bind_expert = n_total_expert;
}
gate_offset += (uint64_t)first_expert * gate_expert_bytes;
up_offset += (uint64_t)first_expert * up_expert_bytes;
down_offset += (uint64_t)first_expert * down_expert_bytes;
const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
const uint64_t up_tensor_bytes = (uint64_t)n_bind_expert * up_expert_bytes;
const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
const uint64_t x_values = (uint64_t)n_tokens * expert_in_dim;
const uint64_t out_values = (uint64_t)n_tokens * out_dim;
const uint64_t mid_values =
(uint64_t)(n_tokens - 1u) * mid_token_stride + per_token_mid;
const uint64_t selected_values = (uint64_t)n_tokens * n_expert;
if (x_values > UINT64_MAX / sizeof(float) ||
out_values > UINT64_MAX / sizeof(float) ||
mid_values > UINT64_MAX / sizeof(float) ||
selected_values > UINT64_MAX / sizeof(int32_t)) {
fprintf(stderr, "ds4: Metal GLM routed batch MoE activation byte size overflow\n");
return 0;
}
const uint64_t x_bytes = x_values * sizeof(float);
const uint64_t mid_bytes = mid_values * sizeof(float);
const uint64_t out_bytes = out_values * sizeof(float);
const uint64_t selected_bytes = selected_values * sizeof(int32_t);
const uint64_t weights_bytes = selected_values * sizeof(float);
if (!xbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(mid) < mid_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes ||
ds4_gpu_tensor_bytes(weights) < weights_bytes) {
fprintf(stderr, "ds4: Metal GLM routed batch MoE received undersized activation buffers\n");
return 0;
}
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
uint64_t down_inner = 0;
id<MTLBuffer> gatebuf = nil;
id<MTLBuffer> upbuf = nil;
id<MTLBuffer> downbuf = nil;
id<MTLBuffer> stream_gate_addr_buf = nil;
id<MTLBuffer> stream_up_addr_buf = nil;
id<MTLBuffer> stream_down_addr_buf = nil;
id<MTLBuffer> stream_overflow_gate = nil;
id<MTLBuffer> stream_overflow_up = nil;
id<MTLBuffer> stream_overflow_down = nil;
ds4_gpu_stream_expert_cache_entry
*stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL };
uint32_t stream_resource_count = 0;
uint32_t stream_unique = 0;
const BOOL gate_pair_q2 = gate_type == DS4_METAL_TENSOR_Q2_K;
const BOOL gate_pair_q5 = gate_type == DS4_METAL_TENSOR_Q5_K;
const BOOL down_scalar_q2 = down_type == DS4_METAL_TENSOR_Q2_K;
const BOOL down_scalar_q4 = down_type == DS4_METAL_TENSOR_Q4_K;
const BOOL down_simd_q4 = down_scalar_q4;
const BOOL down_simd_q5 = down_type == DS4_METAL_TENSOR_Q5_K;
const BOOL down_simd_q6 = down_type == DS4_METAL_TENSOR_Q6_K;
const BOOL down_simd = down_simd_q4 || down_simd_q5 || down_simd_q6;
const BOOL stream_addr_q2 =
gate_pair_q2 && down_scalar_q2 &&
g_glm_q2_k_addr_pair_swiglu2_f32_pipeline != nil &&
g_glm_q2_k_addr_down_f32_pipeline != nil;
const BOOL stream_addr_q4 =
!gate_pair_q2 && !gate_pair_q5 && down_scalar_q4 &&
g_glm_q4_k_addr_pair_swiglu_f32_pipeline != nil &&
g_glm_q4_k_addr_down_f32_pipeline != nil;
BOOL use_stream_expert_addr_table =
g_ssd_streaming_mode &&
!ds4_gpu_glm_streaming_prefill_full_layer_active() &&
n_tokens > 1 &&
(stream_addr_q2 || stream_addr_q4) &&
layer_index < DS4_METAL_STREAM_EXPERT_CACHE_MAX_LAYER &&
n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
n_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_SELECTED &&
ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) &&
ds4_gpu_stream_expert_cache_effective_cap(layer_index,
n_total_expert,
n_expert) != 0;
const BOOL use_stream_grouped_addr_table =
use_stream_expert_addr_table &&
allow_grouped &&
getenv("DS4_METAL_GLM_DISABLE_STREAMING_GROUPED_ADDR_PREFILL") == NULL &&
ds4_gpu_glm_routed_moe_batch_grouped_available(gate_type,
up_type,
down_type,
n_expert,
n_tokens);
const BOOL enable_q4_pair4 = true;
const BOOL q4_scalar_pair = false;
const BOOL q4_pair2 =
!use_stream_expert_addr_table &&
!gate_pair_q5 && !q4_scalar_pair &&
(force_scalar_q4_pair || !enable_q4_pair4);
id<MTLComputePipelineState> pair_pipeline =
use_stream_expert_addr_table ?
(gate_pair_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_addr_pair_swiglu2_f32_pipeline,
"kernel_glm_q2_K_addr_pair_swiglu2_f32") :
ds4_gpu_hot_pipeline(g_glm_q4_k_addr_pair_swiglu_f32_pipeline,
"kernel_glm_q4_K_addr_pair_swiglu_f32")) :
gate_pair_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_pair_swiglu_f32_pipeline,
"kernel_glm_q2_K_pair_swiglu_f32") :
gate_pair_q5 ?
ds4_gpu_hot_pipeline(g_glm_q5_k_pair_swiglu_f32_pipeline,
"kernel_glm_q5_K_pair_swiglu_f32") :
(q4_scalar_pair ?
ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu_f32_pipeline,
"kernel_glm_q4_K_pair_swiglu_f32") :
q4_pair2 ?
ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu2_f32_pipeline,
"kernel_glm_q4_K_pair_swiglu2_f32") :
ds4_gpu_hot_pipeline(g_glm_q4_k_pair_swiglu4_f32_pipeline,
"kernel_glm_q4_K_pair_swiglu4_f32"));
id<MTLComputePipelineState> down_pipeline =
use_stream_expert_addr_table ?
(down_scalar_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_addr_down_f32_pipeline,
"kernel_glm_q2_K_addr_down_f32") :
ds4_gpu_hot_pipeline(g_glm_q4_k_addr_down_f32_pipeline,
"kernel_glm_q4_K_addr_down_f32")) :
down_scalar_q2 ?
ds4_gpu_hot_pipeline(g_glm_q2_k_down_f32_pipeline,
"kernel_glm_q2_K_down_f32") :
down_scalar_q4 ?
ds4_gpu_hot_pipeline(g_glm_q4_k_down_f32_pipeline,
"kernel_glm_q4_K_down_f32") :
down_simd_q5 ?
ds4_gpu_hot_pipeline(g_glm_q5_k_down_f32_pipeline,
"kernel_glm_q5_K_down_f32") :
ds4_gpu_hot_pipeline(g_glm_q6_k_down_f32_pipeline,
"kernel_glm_q6_K_down_f32");
if (!pair_pipeline || !down_pipeline) return 0;
if (use_stream_expert_addr_table) {
const int had_batch = g_batch_cb != nil;
if (had_batch && ds4_gpu_end_commands() == 0) {
return 0;
}
if (!ds4_gpu_stream_expert_cache_prepare_selected_batch(
model_map,
model_size,
layer_index,
selected,
n_tokens,
n_total_expert,
n_expert,
gate_offset,
up_offset,
down_offset,
gate_expert_bytes,
down_expert_bytes,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf,
stream_resources,
&stream_resource_count,
&stream_unique,
&stream_overflow_gate,
&stream_overflow_up,
&stream_overflow_down)) {
return 0;
}
if (stream_unique == 0) {
ds4_gpu_stream_expert_cache_clear_layer(layer_index);
return 0;
}
if (had_batch && ds4_gpu_begin_commands() == 0) {
ds4_gpu_stream_expert_cache_clear_layer(layer_index);
return 0;
}
if (use_stream_grouped_addr_table) {
return ds4_gpu_glm_routed_moe_batch_grouped_addr_tensor(
out,
mid,
stream_gate_addr_buf,
stream_up_addr_buf,
stream_down_addr_buf,
stream_overflow_gate,
stream_overflow_up,
stream_overflow_down,
gate_type,
up_type,
down_type,
gate_expert_bytes,
gate_row_bytes,
up_expert_bytes,
up_row_bytes,
down_expert_bytes,
down_row_bytes,
expert_in_dim,
expert_mid_dim,
out_dim,
selected,
weights,
n_total_expert,
n_expert,
swiglu_clamp,
layer_index,
x,
n_tokens,
stream_resources,
stream_resource_count);
}
} else {
gatebuf = ds4_gpu_wrap_model_range(model_map, model_size,
gate_offset, gate_tensor_bytes,
&gate_inner);
upbuf = ds4_gpu_wrap_model_range(model_map, model_size,
up_offset, up_tensor_bytes,
&up_inner);
downbuf = ds4_gpu_wrap_model_range(model_map, model_size,
down_offset, down_tensor_bytes,
&down_inner);
if (!gatebuf || !upbuf || !downbuf) return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
const bool glm_moe_stage_profile = false;
const char *glm_moe_stage_filter = NULL;
const char *glm_pair_path = use_stream_expert_addr_table ?
(gate_pair_q2 ? "q2_stream_addr_swiglu" :
"q4_stream_addr_swiglu") :
gate_pair_q2 ? "q2_scalar_swiglu" :
gate_pair_q5 ? "q5_pair_simd_swiglu" :
(q4_scalar_pair ? "q4_scalar_swiglu" :
(q4_pair2 ? "q4_pair2_simd_swiglu" :
"q4_pair4_simd_swiglu"));
const char *glm_down_path =
use_stream_expert_addr_table ?
(down_scalar_q2 ? "q2_stream_addr_down" : "q4_stream_addr_down_simd") :
down_scalar_q2 ? "q2_down_scalar" :
down_scalar_q4 ? "q4_down_simd" :
down_simd_q5 ? "q5_down_simd" : "q6_down_simd";
double glm_moe_stage_t0 = 0.0;
if (glm_moe_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
glm_moe_stage_t0 = ds4_gpu_now_ms();
}
int ok = 1;
#define DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE(name) do { \
if (ok && glm_moe_stage_profile) { \
if (ds4_gpu_end_commands() == 0) { \
ok = 0; \
} else { \
const char *stage_name = (name); \
const double now_ms = ds4_gpu_now_ms(); \
const int print_stage = \
!glm_moe_stage_filter || !glm_moe_stage_filter[0] || \
strstr(stage_name, glm_moe_stage_filter) != NULL; \
if (print_stage) { \
fprintf(stderr, \
"ds4: Metal GLM routed MoE batch stage layer=%u tokens=%u experts=%u " \
"gate=%s down=%s pair=%s down_path=%s %s=%.3f ms\n", \
layer_index, n_tokens, n_expert, \
ds4_gpu_metal_tensor_type_name(gate_type), \
ds4_gpu_metal_tensor_type_name(down_type), \
glm_pair_path, glm_down_path, \
stage_name, now_ms - glm_moe_stage_t0); \
} \
glm_moe_stage_t0 = now_ms; \
if (ds4_gpu_begin_commands() == 0) { \
ok = 0; \
} else { \
cb = ds4_gpu_command_buffer(&owned); \
if (!cb) ok = 0; \
} \
} \
} \
} while (0)
ds4_gpu_glm_routed_moe_args args = {
.tp_rank = g_tp_split_rank,
.tp_world = g_tp_split_world,
.tp_expert_base = (int32_t)first_expert,
.in_dim = expert_in_dim,
.mid_dim = expert_mid_dim,
.out_dim = out_dim,
.n_total_expert = n_total_expert,
.n_expert_used = n_expert,
.n_tokens = n_tokens,
.mid_token_stride = mid_token_stride,
.down_type = down_type,
.swiglu_clamp = swiglu_clamp,
.gate_expert_bytes = gate_expert_bytes,
.gate_row_bytes = gate_row_bytes,
.up_expert_bytes = up_expert_bytes,
.up_row_bytes = up_row_bytes,
.down_expert_bytes = down_expert_bytes,
.down_row_bytes = down_row_bytes,
};
const NSUInteger pair_x_groups =
gate_pair_q2 ? (use_stream_expert_addr_table ?
(NSUInteger)((expert_mid_dim + 1u) / 2u) :
(NSUInteger)((expert_mid_dim + 7u) / 8u)) :
gate_pair_q5 ? (NSUInteger)((expert_mid_dim + 7u) / 8u) :
use_stream_expert_addr_table ? (NSUInteger)((expert_mid_dim + 3u) / 4u) :
q4_scalar_pair ? (NSUInteger)expert_mid_dim :
q4_pair2 ? (NSUInteger)((expert_mid_dim + 1u) / 2u) :
(NSUInteger)((expert_mid_dim + 7u) / 8u);
const NSUInteger pair_threadgroup_bytes =
q4_scalar_pair ? 512u * sizeof(float) : 0u;
const NSUInteger pair_threads =
q4_scalar_pair ? 256u : 64u;
const NSUInteger down_x_groups =
down_scalar_q2 ? (NSUInteger)((out_dim + 7u) / 8u) :
down_simd_q4 ? (NSUInteger)((out_dim + 3u) / 4u) :
down_simd_q5 ? (NSUInteger)((out_dim + 3u) / 4u) :
down_simd_q6 ? (NSUInteger)((out_dim + 3u) / 4u) :
(NSUInteger)out_dim;
const NSUInteger down_threadgroup_bytes =
down_simd ? 0u : 256u * sizeof(float);
const NSUInteger down_threads =
down_simd ? 64u : 256u;
if (use_stream_expert_addr_table &&
!ds4_gpu_stream_expert_cache_mark_entries_inflight(stream_resources,
stream_resource_count,
0)) {
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pair_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:use_stream_expert_addr_table ? stream_gate_addr_buf : gatebuf
offset:use_stream_expert_addr_table ? 0u : (NSUInteger)gate_inner
atIndex:1];
[enc setBuffer:use_stream_expert_addr_table ? stream_up_addr_buf : upbuf
offset:use_stream_expert_addr_table ? 0u : (NSUInteger)up_inner
atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:4];
[enc setBuffer:weightsbuf offset:ds4_gpu_tensor_offset(weights) atIndex:5];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:6];
if (use_stream_expert_addr_table) {
for (uint32_t i = 0; i < stream_resource_count; i++) {
[enc useResource:stream_resources[i]->gate_buffer usage:MTLResourceUsageRead];
[enc useResource:stream_resources[i]->up_buffer usage:MTLResourceUsageRead];
}
if (stream_overflow_gate) [enc useResource:stream_overflow_gate usage:MTLResourceUsageRead];
if (stream_overflow_up) [enc useResource:stream_overflow_up usage:MTLResourceUsageRead];
}
if (pair_threadgroup_bytes != 0u) {
[enc setThreadgroupMemoryLength:pair_threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(pair_x_groups,
(NSUInteger)n_expert,
(NSUInteger)n_tokens)
threadsPerThreadgroup:MTLSizeMake(pair_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("pair");
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:down_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:use_stream_expert_addr_table ? stream_down_addr_buf : downbuf
offset:use_stream_expert_addr_table ? 0u : (NSUInteger)down_inner
atIndex:1];
[enc setBuffer:selectedbuf offset:ds4_gpu_tensor_offset(selected) atIndex:2];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(mid) atIndex:3];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
if (use_stream_expert_addr_table) {
for (uint32_t i = 0; i < stream_resource_count; i++) {
[enc useResource:stream_resources[i]->down_buffer usage:MTLResourceUsageRead];
}
if (stream_overflow_down) [enc useResource:stream_overflow_down usage:MTLResourceUsageRead];
}
if (down_threadgroup_bytes != 0u) {
[enc setThreadgroupMemoryLength:down_threadgroup_bytes atIndex:0];
}
[enc dispatchThreadgroups:MTLSizeMake(down_x_groups,
(NSUInteger)n_tokens,
1)
threadsPerThreadgroup:MTLSizeMake(down_threads, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE("down");
if (!ok) return 0;
if (!ds4_gpu_finish_command_buffer(cb, owned, "GLM routed batch MoE")) return 0;
#undef DS4_METAL_PROFILE_GLM_MOE_BATCH_STAGE
}
return 1;
}
int ds4_gpu_glm_routed_moe_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t mid_token_stride,
bool force_resident) {
(void)force_resident;
return ds4_gpu_glm_routed_moe_batch_tensor_impl(out,
mid,
model_map,
model_size,
gate_offset,
up_offset,
down_offset,
gate_type,
up_type,
down_type,
gate_expert_bytes,
gate_row_bytes,
up_expert_bytes,
up_row_bytes,
down_expert_bytes,
down_row_bytes,
expert_in_dim,
expert_mid_dim,
out_dim,
selected,
weights,
n_total_expert,
n_expert,
swiglu_clamp,
layer_index,
x,
n_tokens,
mid_token_stride,
true,
false);
}
int ds4_gpu_glm_routed_moe_batch_direct_scalar_q4_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *mid,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t up_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t up_expert_bytes,
uint64_t up_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float swiglu_clamp,
uint32_t layer_index,
const ds4_gpu_tensor *x,
uint32_t n_tokens,
uint32_t mid_token_stride) {
return ds4_gpu_glm_routed_moe_batch_tensor_impl(out,
mid,
model_map,
model_size,
gate_offset,
up_offset,
down_offset,
gate_type,
up_type,
down_type,
gate_expert_bytes,
gate_row_bytes,
up_expert_bytes,
up_row_bytes,
down_expert_bytes,
down_row_bytes,
expert_in_dim,
expert_mid_dim,
out_dim,
selected,
weights,
n_total_expert,
n_expert,
swiglu_clamp,
layer_index,
x,
n_tokens,
mid_token_stride,
false,
false);
}
int ds4_gpu_router_select_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint64_t hash_offset,
uint32_t hash_rows,
uint32_t token,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_expert_groups,
uint32_t n_group_used,
bool has_bias,
bool hash_mode,
const ds4_gpu_tensor *logits) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!selected || !weights || !probs || !logits || !model_map ||
n_expert == 0 || n_expert_used == 0) return 0;
if (hash_mode && token >= hash_rows) return 0;
if (n_expert_groups > 1u || n_group_used > 0u) {
fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n");
return 0;
}
@autoreleasepool {
id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf ||
ds4_gpu_tensor_bytes(logits) < (uint64_t)n_expert * sizeof(float) ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert_used * sizeof(int) ||
ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert_used * sizeof(float) ||
ds4_gpu_tensor_bytes(probs) < (uint64_t)n_expert * sizeof(float)) {
fprintf(stderr, "ds4: Metal router select received undersized buffers\n");
return 0;
}
uint64_t bias_inner = 0;
uint64_t hash_inner = 0;
id<MTLBuffer> biasbuf = nil;
id<MTLBuffer> hashbuf = nil;
NSUInteger bias_set_offset = 0;
NSUInteger hash_set_offset = 0;
if (has_bias && !hash_mode) {
const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner);
if (!biasbuf) return 0;
bias_set_offset = (NSUInteger)bias_inner;
}
if (hash_mode) {
const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t);
hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner);
if (!hashbuf) return 0;
hash_set_offset = (NSUInteger)hash_inner;
}
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
const int32_t token_i32 = (int32_t)token;
int ok = cb &&
ds4_gpu_encode_router_select(cb,
selected,
weights,
probs,
logitsbuf,
ds4_gpu_tensor_offset(logits),
biasbuf,
bias_set_offset,
hashbuf,
hash_set_offset,
nil,
0,
nil,
0,
&token_i32,
hash_rows,
0,
1,
n_expert,
n_expert_used,
expert_weight_scale,
has_bias && !hash_mode,
hash_mode,
false);
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
if (!ok) return 0;
}
return 1;
}
int ds4_gpu_router_select_batch_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint64_t hash_offset,
uint32_t hash_rows,
uint32_t n_expert_groups,
uint32_t n_group_used,
bool has_bias,
bool hash_mode,
const ds4_gpu_tensor *logits,
const ds4_gpu_tensor *tokens,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!selected || !weights || !probs || !logits || !tokens || !model_map ||
n_expert == 0 || n_expert_used == 0 || n_tokens == 0) return 0;
if (n_expert_groups > 1u || n_group_used > 0u) {
fprintf(stderr, "ds4: Metal router group gating is not part of this DeepSeek V4 path\n");
return 0;
}
@autoreleasepool {
id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
id<MTLBuffer> tokensbuf = ds4_gpu_tensor_buffer(tokens);
if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || !tokensbuf ||
ds4_gpu_tensor_bytes(logits) < (uint64_t)n_tokens * n_expert * sizeof(float) ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)n_tokens * n_expert_used * sizeof(int) ||
ds4_gpu_tensor_bytes(weights) < (uint64_t)n_tokens * n_expert_used * sizeof(float) ||
ds4_gpu_tensor_bytes(probs) < (uint64_t)n_tokens * n_expert * sizeof(float) ||
ds4_gpu_tensor_bytes(tokens) < (uint64_t)n_tokens * sizeof(int32_t)) {
fprintf(stderr, "ds4: Metal router batch select received undersized buffers\n");
return 0;
}
uint64_t bias_inner = 0;
uint64_t hash_inner = 0;
id<MTLBuffer> biasbuf = nil;
id<MTLBuffer> hashbuf = nil;
NSUInteger bias_set_offset = 0;
NSUInteger hash_set_offset = 0;
if (has_bias && !hash_mode) {
const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset, bias_bytes, &bias_inner);
if (!biasbuf) return 0;
bias_set_offset = (NSUInteger)bias_inner;
}
if (hash_mode) {
const uint64_t hash_bytes = (uint64_t)hash_rows * n_expert_used * sizeof(int32_t);
hashbuf = ds4_gpu_wrap_model_range(model_map, model_size, hash_offset, hash_bytes, &hash_inner);
if (!hashbuf) return 0;
hash_set_offset = (NSUInteger)hash_inner;
}
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
int ok = cb &&
ds4_gpu_encode_router_select(cb,
selected,
weights,
probs,
logitsbuf,
ds4_gpu_tensor_offset(logits),
biasbuf,
bias_set_offset,
hashbuf,
hash_set_offset,
tokensbuf,
ds4_gpu_tensor_offset(tokens),
nil,
0,
NULL,
hash_rows,
0,
n_tokens,
n_expert,
n_expert_used,
expert_weight_scale,
has_bias && !hash_mode,
hash_mode,
false);
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
if (!ok) return 0;
}
return 1;
}
int ds4_gpu_router_select_batch_visual_tensor(
ds4_gpu_tensor *selected,
ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint64_t hash_offset,
uint32_t hash_rows,
bool has_bias,
bool hash_mode,
const void *vision_map,
uint64_t vision_size,
uint64_t visual_bias_offset,
const ds4_gpu_tensor *logits,
const ds4_gpu_tensor *tokens,
uint32_t vocab_size,
uint32_t n_expert,
uint32_t n_expert_used,
float expert_weight_scale,
uint32_t n_tokens) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!selected || !weights || !probs || !logits || !tokens ||
!model_map || !vision_map || vocab_size == 0 || n_tokens == 0 ||
n_expert != 256u || n_expert_used != 6u ||
fabsf(expert_weight_scale - 1.5f) > 1.0e-6f) return 0;
@autoreleasepool {
id<MTLBuffer> logitsbuf = ds4_gpu_tensor_buffer(logits);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
id<MTLBuffer> tokensbuf = ds4_gpu_tensor_buffer(tokens);
if (!logitsbuf || !selectedbuf || !weightsbuf || !probsbuf || !tokensbuf ||
ds4_gpu_tensor_bytes(logits) < (uint64_t)n_tokens * n_expert * sizeof(float) ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)n_tokens * n_expert_used * sizeof(int) ||
ds4_gpu_tensor_bytes(weights) < (uint64_t)n_tokens * n_expert_used * sizeof(float) ||
ds4_gpu_tensor_bytes(probs) < (uint64_t)n_tokens * n_expert * sizeof(float) ||
ds4_gpu_tensor_bytes(tokens) < (uint64_t)n_tokens * sizeof(int32_t)) {
fprintf(stderr, "ds4: Metal visual router received undersized buffers\n");
return 0;
}
uint64_t bias_inner = 0, hash_inner = 0, visual_inner = 0;
id<MTLBuffer> biasbuf = nil, hashbuf = nil;
NSUInteger bias_set_offset = 0, hash_set_offset = 0;
if (has_bias && !hash_mode) {
biasbuf = ds4_gpu_wrap_model_range(model_map, model_size,
bias_offset,
(uint64_t)n_expert * sizeof(float),
&bias_inner);
if (!biasbuf) return 0;
bias_set_offset = (NSUInteger)bias_inner;
}
if (hash_mode) {
hashbuf = ds4_gpu_wrap_model_range(
model_map, model_size, hash_offset,
(uint64_t)hash_rows * n_expert_used * sizeof(int32_t),
&hash_inner);
if (!hashbuf) return 0;
hash_set_offset = (NSUInteger)hash_inner;
}
id<MTLBuffer> visual_biasbuf = ds4_gpu_wrap_model_range(
vision_map, vision_size, visual_bias_offset,
(uint64_t)n_expert * sizeof(float), &visual_inner);
if (!visual_biasbuf) return 0;
const bool had_batch = g_batch_cb != nil;
if (!had_batch && ds4_gpu_begin_commands() == 0) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
int ok = cb && ds4_gpu_encode_router_select(
cb, selected, weights, probs,
logitsbuf, ds4_gpu_tensor_offset(logits),
biasbuf, bias_set_offset,
hashbuf, hash_set_offset,
tokensbuf, ds4_gpu_tensor_offset(tokens),
visual_biasbuf, (NSUInteger)visual_inner,
NULL, hash_rows, vocab_size, n_tokens,
n_expert, n_expert_used, expert_weight_scale,
has_bias && !hash_mode, hash_mode, true);
if (!had_batch) ok = ds4_gpu_end_commands() != 0 && ok;
if (!ok) return 0;
}
return 1;
}
int ds4_gpu_routed_moe_set_selected_override(const int32_t *selected, uint32_t n_selected) {
if (n_selected > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
(!selected && n_selected != 0)) return 0;
for (uint32_t i = 0; i < n_selected; i++) {
g_routed_moe_selected_override[i] = selected[i];
}
g_routed_moe_selected_override_n = n_selected;
return 1;
}
/* Single-simdgroup MXFP4 decode MoE policy: the nsg=1 pair_swiglu/sum6
* variants are the pre-M5 Apple-Silicon default for single-token decode
* (bit-exact; only the row-to-simdgroup mapping changes). Read on the
* dispatch path so A/B harnesses can toggle it after engine creation. */
static bool ds4_gpu_mxfp4_moe_decode_nsg1_enabled(uint32_t n_tokens) {
if (n_tokens != 1) return false;
return ds4_gpu_device_is_pre_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_NSG1") == NULL;
}
int ds4_gpu_routed_moe_one_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *experts,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float clamp,
const ds4_gpu_tensor *x,
const ds4_gpu_tensor *add_in,
uint32_t layer_index,
bool force_resident) {
BOOL parallel_ffn_scope
__attribute__((cleanup(ds4_gpu_parallel_ffn_scope_cleanup))) =
g_parallel_q8_pending;
if (!g_initialized && !ds4_gpu_init()) return 0;
/* TP sharding: only the owned contiguous expert range is mapped,
* so bind from the owned base, validate only its bytes, and tell the
* kernels the first expert id present at that base. */
uint32_t first_expert = 0;
uint32_t n_bind_expert = 0;
ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
const int32_t tp_expert_base_host = (int32_t)first_expert;
gate_offset += (uint64_t)first_expert * gate_expert_bytes;
up_offset += (uint64_t)first_expert * gate_expert_bytes;
down_offset += (uint64_t)first_expert * down_expert_bytes;
if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights ||
n_total_expert == 0 || n_expert == 0 ||
n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED ||
gate_expert_bytes == 0 || down_expert_bytes == 0 ||
gate_row_bytes == 0 || down_row_bytes == 0) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32044);
return 0;
}
if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32047); return 0; }
ds4_gpu_stream_expert_cache_note_token(layer_index);
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> expertsbuf = ds4_gpu_tensor_buffer(experts);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> selected_exec_buf = selectedbuf;
NSUInteger selected_exec_off = ds4_gpu_tensor_offset(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
const uint64_t x_bytes = (uint64_t)expert_in_dim * sizeof(float);
const uint64_t mid_bytes = (uint64_t)n_expert * expert_mid_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)out_dim * sizeof(float);
if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(gate) < mid_bytes ||
ds4_gpu_tensor_bytes(up) < mid_bytes ||
ds4_gpu_tensor_bytes(mid) < mid_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)n_expert * sizeof(int) ||
ds4_gpu_tensor_bytes(weights) < (uint64_t)n_expert * sizeof(float)) {
fprintf(stderr, "ds4: Metal routed tensor MoE received undersized activation buffers\n");
return 0;
}
if (n_expert > 1 &&
(!expertsbuf ||
ds4_gpu_tensor_bytes(experts) < (uint64_t)n_expert * out_dim * sizeof(float))) {
fprintf(stderr, "ds4: Metal routed tensor MoE received undersized expert output buffer\n");
return 0;
}
if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal routed MoE tensor byte size overflow\n");
return 0;
}
const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
uint64_t down_inner = 0;
id<MTLBuffer> gate_buf = nil;
id<MTLBuffer> up_buf = nil;
id<MTLBuffer> down_buf = nil;
__unsafe_unretained id<MTLBuffer> gate_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil };
__unsafe_unretained id<MTLBuffer> up_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil };
__unsafe_unretained id<MTLBuffer> down_slot_bufs[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { nil };
ds4_gpu_stream_expert_cache_entry *stream_slot_entries[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { NULL };
NSUInteger gate_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
NSUInteger up_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
NSUInteger down_slot_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
uint64_t stream_gate_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
uint64_t stream_up_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
uint64_t stream_down_abs_offsets[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
id<MTLBuffer> stream_gate_addr_buf = nil;
id<MTLBuffer> stream_up_addr_buf = nil;
id<MTLBuffer> stream_down_addr_buf = nil;
bool use_stream_expert_addr_table = false;
bool use_stream_expert_masked_addr_table = false;
bool use_stream_compact_addr_table = false;
bool use_stream_expert_cache = false;
bool use_stream_expert_split_candidate = false;
bool use_stream_expert_split_deferred = false;
bool stream_expert_split_completed = false;
uint32_t stream_expert_resident_mask = 0;
uint32_t stream_expert_missing_mask = 0;
__unsafe_unretained id<MTLBuffer> gate_group6_bufs[6] = { nil, nil, nil, nil, nil, nil };
__unsafe_unretained id<MTLBuffer> up_group6_bufs[6] = { nil, nil, nil, nil, nil, nil };
__unsafe_unretained id<MTLBuffer> down_group6_bufs[6] = { nil, nil, nil, nil, nil, nil };
NSUInteger gate_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 };
NSUInteger up_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 };
NSUInteger down_group6_offsets[6] = { 0, 0, 0, 0, 0, 0 };
__unsafe_unretained id<MTLBuffer> gate_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil };
__unsafe_unretained id<MTLBuffer> up_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil };
__unsafe_unretained id<MTLBuffer> down_group8_bufs[8] = { nil, nil, nil, nil, nil, nil, nil, nil };
NSUInteger gate_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
NSUInteger up_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
NSUInteger down_group8_offsets[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
__unsafe_unretained id<MTLBuffer> gate_group24_bufs[24] = { nil };
__unsafe_unretained id<MTLBuffer> up_group24_bufs[24] = { nil };
__unsafe_unretained id<MTLBuffer> down_group24_bufs[24] = { nil };
NSUInteger gate_group24_offsets[24] = { 0 };
NSUInteger up_group24_offsets[24] = { 0 };
NSUInteger down_group24_offsets[24] = { 0 };
DS4MetalQ4ExpertTable *gate_table = nil;
DS4MetalQ4ExpertTable *up_table = nil;
DS4MetalQ4ExpertTable *down_table = nil;
id q4_table_layer_residency = nil;
int32_t selected_ids[DS4_METAL_MAX_ROUTED_EXPERT_USED] = { 0 };
const uint32_t n_tokens = 1;
const uint32_t pair_rows = n_tokens * n_expert;
const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float);
if ((n_expert > 1 && !expertsbuf &&
!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
&g_moe_down_scratch_bytes,
(NSUInteger)down_scratch_bytes,
"ds4_moe_down_scratch"))) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32148);
return 0;
}
const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type);
const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type);
id<MTLComputePipelineState> gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type);
id<MTLComputePipelineState> down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type);
if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) {
fprintf(stderr, "ds4: unsupported Metal routed MoE quant types gate=%u down=%u\n",
gate_type, down_type);
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32158);
return 0;
}
ds4_gpu_mul_mv_id_args gate_args =
ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
gate_row_bytes, gate_expert_bytes,
1, n_expert, n_tokens, gate_nr0);
ds4_gpu_mul_mv_id_args down_args =
ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert,
down_row_bytes, down_expert_bytes,
n_expert, n_expert, n_tokens, down_nr0);
/* Tensor-parallel expert ownership; non-TP calls keep tp_world at 1. */
gate_args.tp_rank = g_tp_split_rank;
gate_args.tp_world = g_tp_full_expert_bind ? 1u : g_tp_split_world;
gate_args.tp_expert_base = tp_expert_base_host;
down_args.tp_rank = g_tp_split_rank;
down_args.tp_world = g_tp_full_expert_bind ? 1u : g_tp_split_world;
down_args.tp_addend = add_in != NULL;
down_args.tp_expert_base = tp_expert_base_host;
const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type);
const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type);
const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type);
const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type);
const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type);
const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type);
int pair_swiglu_nsg = 2;
int down_sum6_nsg = 2;
int ok = 1;
const bool write_clamped_moe =
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") != NULL;
/* The selected nsg=1 encoders dispatch exactly 32 threads. */
const bool use_mxfp4_moe_decode_tg_multiple =
ds4_gpu_device_is_pre_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_TG_MULTIPLE") == NULL &&
gate_type == DS4_METAL_TENSOR_MXFP4 &&
down_type == DS4_METAL_TENSOR_MXFP4 &&
n_tokens == 1 &&
n_expert == 6 &&
g_tp_split_world == 1 &&
add_in == NULL &&
!g_quality_mode &&
!write_clamped_moe &&
(force_resident || !g_ssd_streaming_mode);
const bool use_mxfp4_moe_decode_fixed_route_pair =
use_mxfp4_moe_decode_tg_multiple &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_FIXED_ROUTE_PAIR") == NULL &&
((expert_in_dim == 4096 && expert_mid_dim == 2048 && n_total_expert == 256 &&
gate_row_bytes == 2176) ||
(expert_in_dim == 5120 && expert_mid_dim == 2304 && n_total_expert == 384 &&
gate_row_bytes == 2720)) &&
gate_args.nei0 == 6 &&
gate_args.nei1 == 1 &&
gate_args.ne11 == 1 &&
gate_args.tp_world == 1 &&
gate_args.tp_expert_base == 0 &&
g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 != nil;
const bool use_mxfp4_moe_decode_fixed_route_sum6 =
use_mxfp4_moe_decode_fixed_route_pair &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_FIXED_ROUTE_SUM6") == NULL &&
(out_dim == 4096 || out_dim == 5120) &&
down_row_bytes == (uint64_t)expert_mid_dim / 32u * 17u &&
down_expert_bytes == (uint64_t)out_dim * down_row_bytes &&
down_args.nei0 == 6 &&
down_args.nei1 == 1 &&
down_args.ne11 == 6 &&
down_args.tp_world == 1 &&
down_args.tp_expert_base == 0 &&
!down_args.tp_addend &&
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 != nil;
const bool use_mxfp4_moe_decode_sum6_full_rows =
use_mxfp4_moe_decode_fixed_route_sum6 &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_SUM6_FULL_ROWS") == NULL &&
down_args.ne00 == (int32_t)expert_mid_dim &&
down_args.ne01 == (int32_t)out_dim &&
down_args.ne0 == (int32_t)out_dim &&
down_args.nr0 == 2 &&
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 != nil;
/* Compile-time trip counts for the one-token fixed-route decode MoE
* kernels. Every shape the specialized kernels assume is proven above
* (gate ne00 4096 / row bytes 2176; down ne00 2048 / row bytes 1088 /
* nei0 6, or the DeepSeek V4.1 Flash 5120 / 2720; 2304 / 1224 shape),
* so the specialization only removes runtime loop bounds; the
* per-lane K walk and accumulate order stay byte-identical. */
const bool use_mxfp4_moe_decode_static_trip_pair =
use_mxfp4_moe_decode_fixed_route_pair &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_DECODE_STATIC_TRIP") == NULL &&
((gate_args.ne00 == 4096 && gate_args.nb01 == 2176) ||
(gate_args.ne00 == 5120 && gate_args.nb01 == 2720)) &&
g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 != nil;
const bool use_mxfp4_moe_decode_static_trip_down =
use_mxfp4_moe_decode_sum6_full_rows &&
use_mxfp4_moe_decode_static_trip_pair &&
(down_args.nb01 == 1088 || down_args.nb01 == 1224) &&
down_args.nei0 == 6 &&
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 != nil;
const bool use_m5_tp_mxfp4_static =
ds4_gpu_device_is_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_M5_TP_MXFP4_STATIC") == NULL &&
gate_type == DS4_METAL_TENSOR_MXFP4 &&
down_type == DS4_METAL_TENSOR_MXFP4 &&
n_expert == 6 && n_total_expert == 256 &&
expert_in_dim == 4096 && expert_mid_dim == 2048 &&
out_dim == 4096 && gate_row_bytes == 2176 &&
gate_expert_bytes == 4456448 && down_row_bytes == 1088 &&
down_expert_bytes == 4456448 &&
gate_args.nei0 == 6 && gate_args.nei1 == 1 &&
gate_args.ne11 == 1 && down_args.nei0 == 6 &&
down_args.nei1 == 1 && down_args.ne11 == 6 &&
gate_args.tp_world == 2 && down_args.tp_world == 2 &&
((g_tp_split_rank == 0 && first_expert == 0u) ||
(g_tp_split_rank == 1 &&
first_expert + n_bind_expert == n_total_expert)) &&
!g_quality_mode && !write_clamped_moe &&
(force_resident || !g_ssd_streaming_mode) &&
g_moe_mul_mv_id_mxfp4_pair_swiglu_tp_static_pipeline_nsg1 != nil &&
g_moe_mul_mv_id_mxfp4_sum6_tp_full_rows_static_pipeline_nsg1 != nil;
id<MTLComputePipelineState> pair_swiglu_pipeline = nil;
if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) {
pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline;
} else if (gate_type == DS4_METAL_TENSOR_Q4_K) {
pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline;
} else if (gate_type == DS4_METAL_TENSOR_MXFP4) {
pair_swiglu_pipeline = g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline;
if (use_m5_tp_mxfp4_static) {
pair_swiglu_pipeline =
g_moe_mul_mv_id_mxfp4_pair_swiglu_tp_static_pipeline_nsg1;
pair_swiglu_nsg = 1;
} else if (ds4_gpu_mxfp4_moe_decode_nsg1_enabled(n_tokens)) {
pair_swiglu_pipeline =
use_mxfp4_moe_decode_static_trip_pair ?
g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_static_pipeline_nsg1 :
use_mxfp4_moe_decode_fixed_route_pair ?
g_moe_mul_mv_id_mxfp4_pair_swiglu_fixed_route_pipeline_nsg1 :
(use_mxfp4_moe_decode_tg_multiple ?
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1_tg_multiple :
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline_nsg1);
pair_swiglu_nsg = 1;
}
}
const bool use_iq2_pair_pack2 =
gate_type == DS4_METAL_TENSOR_IQ2_XXS && n_tokens == 1 &&
n_expert == 6 && n_total_expert == 256 &&
expert_in_dim == 4096 && expert_mid_dim == 2048 &&
getenv("DS4_METAL_DISABLE_M5_IQ2_PAIR_PACK2") == NULL &&
ds4_gpu_device_is_m5_apple_silicon();
if (use_iq2_pair_pack2 &&
g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline) {
pair_swiglu_pipeline =
g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pack2_pipeline;
pair_swiglu_nsg = 4;
}
const bool fuse_pair_swiglu =
!g_quality_mode &&
!write_clamped_moe &&
getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
pair_swiglu_pipeline != nil;
id<MTLComputePipelineState> down_sum6_pipeline = nil;
if (down_type == DS4_METAL_TENSOR_Q2_K) {
down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline;
} else if (down_type == DS4_METAL_TENSOR_Q4_K) {
down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline;
} else if (down_type == DS4_METAL_TENSOR_MXFP4) {
down_sum6_pipeline = g_moe_mul_mv_id_mxfp4_sum6_pipeline;
if (use_m5_tp_mxfp4_static) {
down_sum6_pipeline =
g_moe_mul_mv_id_mxfp4_sum6_tp_full_rows_static_pipeline_nsg1;
down_sum6_nsg = 1;
} else if (ds4_gpu_mxfp4_moe_decode_nsg1_enabled(n_tokens)) {
down_sum6_pipeline =
use_mxfp4_moe_decode_static_trip_down ?
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_static_pipeline_nsg1 :
use_mxfp4_moe_decode_sum6_full_rows ?
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_full_rows_pipeline_nsg1 :
(use_mxfp4_moe_decode_fixed_route_sum6 ?
g_moe_mul_mv_id_mxfp4_sum6_fixed_route_pipeline_nsg1 :
(use_mxfp4_moe_decode_tg_multiple ?
g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1_tg_multiple :
g_moe_mul_mv_id_mxfp4_sum6_pipeline_nsg1));
down_sum6_nsg = 1;
}
} else if (down_type == DS4_METAL_TENSOR_IQ2_XXS &&
g_tp_split_world == 2) {
/* IQ2 down-sum exists for the GLM TP resident split only; the
* streaming paths must keep their addr/masked chain (expert
* bytes are not at their model-map offsets when streamed). */
down_sum6_pipeline = g_moe_mul_mv_id_iq2_xxs_sum6_pipeline;
}
const bool direct_down_sum =
!g_quality_mode &&
(n_expert == 6 || (n_expert == 8 && g_tp_split_world == 2) ||
g_tp_full_expert_bind) &&
n_tokens == 1 &&
down_sum6_pipeline != nil;
if (g_parallel_q8_pending) {
/* A concurrent encoder invalidates every implicit dependency in
* this generic function. Admit only resident fused pair-SwiGLU
* producers followed by direct sum6 consumers. Those are the two
* levels explicitly bracketed by the concurrent encoder. */
const bool parallel_iq2_route =
g_parallel_ffn_mode == 2 &&
gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
down_type == DS4_METAL_TENSOR_Q2_K &&
n_tokens == 1 && n_expert == 6 && n_total_expert == 256 &&
expert_in_dim == 4096 && expert_mid_dim == 2048 &&
out_dim == 4096 && gate_row_bytes == 1056 &&
gate_expert_bytes == 2162688 && down_row_bytes == 672 &&
down_expert_bytes == 2752512 &&
first_expert == 0 && n_bind_expert == n_total_expert &&
g_tp_split_rank == 0 && g_tp_split_world == 1 &&
add_in == NULL && (force_resident || !g_ssd_streaming_mode) &&
!write_clamped_moe && fuse_pair_swiglu && direct_down_sum;
const bool parallel_mxfp4_tp_route =
g_parallel_ffn_mode == 2 &&
gate_type == DS4_METAL_TENSOR_MXFP4 &&
down_type == DS4_METAL_TENSOR_MXFP4 &&
n_tokens == 1 && n_expert == 6 &&
((n_total_expert == 256 &&
expert_in_dim == 4096 && expert_mid_dim == 2048 &&
out_dim == 4096 && gate_row_bytes == 2176 &&
gate_expert_bytes == 4456448 && down_row_bytes == 1088 &&
down_expert_bytes == 4456448) ||
(n_total_expert == 384 &&
expert_in_dim == 5120 && expert_mid_dim == 2304 &&
out_dim == 5120 && gate_row_bytes == 2720 &&
gate_expert_bytes == 6266880 && down_row_bytes == 1224 &&
down_expert_bytes == 6266880)) &&
((g_tp_split_world == 1 && first_expert == 0u &&
n_bind_expert == n_total_expert) ||
(g_tp_split_world == 2 &&
((g_tp_split_rank == 0 && first_expert == 0u) ||
(g_tp_split_rank == 1 &&
first_expert + n_bind_expert == n_total_expert)))) &&
add_in == NULL &&
(force_resident || !g_ssd_streaming_mode) &&
!write_clamped_moe && fuse_pair_swiglu && direct_down_sum;
if (!parallel_iq2_route && !parallel_mxfp4_tp_route) {
fprintf(stderr,
"ds4: concurrent FFN requires a supported resident "
"fused pair-SwiGLU + direct sum6 route\n");
return 0;
}
}
/* The expert-ownership split lives only in the fused id pair+sum6
* kernels; every other routed variant would silently compute full
* sums on both ranks and double the combine. Fail fast instead. */
if ((g_tp_split_world > 1 || add_in) && !(fuse_pair_swiglu && direct_down_sum)) {
fprintf(stderr, "ds4: tensor-parallel routed MoE requires the fused pair+sum6 decode path\n");
return 0;
}
const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull;
/*
* The grouped Q4 experiment keeps selected IDs on GPU, but it also walks
* every expert window in the layer. On PRO Q4 this measured far slower
* than the active selected-slot path, so keep it opt-in for profiling.
*/
const bool enable_q4_grouped_experts =
!g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL;
const bool use_q4_grouped_experts =
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert >= 128 &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
fuse_pair_swiglu &&
direct_down_sum &&
g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_group_q4_k_sum6_pipeline != nil &&
enable_q4_grouped_experts;
const uint32_t q4_expert_group_size =
use_q4_grouped_experts ? ds4_gpu_q4_expert_group_size(n_total_expert) : 0;
const bool q4_grouped_boundary =
use_q4_grouped_experts &&
g_batch_cb != nil &&
getenv("DS4_METAL_DISABLE_Q4_GROUPED_BOUNDARY") == NULL;
const bool q4_grouped_cache_views =
getenv("DS4_METAL_Q4_GROUPED_CACHE_VIEWS") != NULL;
const uint32_t q4_group6_expert_group_size = 64;
const bool use_q4_group6_experts =
!g_ssd_streaming_mode &&
!use_q4_grouped_experts &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert == q4_group6_expert_group_size * 6u &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
fuse_pair_swiglu &&
direct_down_sum &&
g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil &&
getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL;
const uint32_t q4_group8_expert_group_size = 48;
const bool use_q4_group8_experts =
!g_ssd_streaming_mode &&
!use_q4_grouped_experts &&
!use_q4_group6_experts &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert == q4_group8_expert_group_size * 8u &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
fuse_pair_swiglu &&
direct_down_sum &&
g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil &&
getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL;
const uint32_t q4_group24_expert_group_size = 16;
const bool use_q4_group24_experts =
!g_ssd_streaming_mode &&
!use_q4_grouped_experts &&
!use_q4_group6_experts &&
!use_q4_group8_experts &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert == q4_group24_expert_group_size * 24u &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
direct_down_sum &&
g_moe_mul_mv_group24_q4_k_id_pipeline != nil &&
g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil &&
getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL;
const bool q4_group24_exact_views =
use_q4_group24_experts &&
getenv("DS4_METAL_Q4_GROUP24_EXACT_VIEWS") != NULL &&
getenv("DS4_METAL_Q4_GROUP24_BASE_VIEWS") == NULL;
const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0;
const bool can_wrap_q4_exact_tensors =
max_buffer_len != 0 &&
gate_tensor_bytes <= max_buffer_len &&
down_tensor_bytes <= max_buffer_len;
/*
* The full-tensor Q4 ID path is the closest arithmetic analogue to IQ2,
* but PRO Q4 routed tensors are multi-GiB. Keep it opt-in: even with
* per-layer command boundaries it is much slower than binding only the
* six active experts on current M3 Ultra Metal.
*/
const bool enable_q4_exact_tensor_id =
!g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL;
const bool use_q4_exact_tensor_id =
!use_q4_grouped_experts &&
!use_q4_group6_experts &&
!use_q4_group8_experts &&
!use_q4_group24_experts &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert == 384 &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
can_wrap_q4_exact_tensors &&
fuse_pair_swiglu &&
direct_down_sum &&
enable_q4_exact_tensor_id;
const bool q4_exact_boundary =
use_q4_exact_tensor_id &&
g_batch_cb != nil &&
getenv("DS4_METAL_DISABLE_Q4_EXACT_BOUNDARY") == NULL;
const bool q4_expert_table_auto =
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert,
n_expert,
gate_tensor_bytes,
down_tensor_bytes);
const bool q4_expert_address_auto =
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
ds4_gpu_pro_q4_expert_address_auto_enabled(n_total_expert,
n_expert,
gate_tensor_bytes,
down_tensor_bytes);
const bool q4_table_queue_residency =
ds4_gpu_q4_table_queue_residency_enabled(q4_expert_table_auto ||
q4_expert_address_auto);
const bool use_q4_expert_address_table =
!g_ssd_streaming_mode &&
!use_q4_grouped_experts &&
!use_q4_group6_experts &&
!use_q4_group8_experts &&
!use_q4_group24_experts &&
!use_q4_exact_tensor_id &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert == 384 &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
fuse_pair_swiglu &&
direct_down_sum &&
g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil &&
(getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL ||
q4_expert_address_auto) &&
(getenv("DS4_METAL_Q4_ADDR_USE_RESOURCES") != NULL ||
getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL ||
q4_table_queue_residency ||
ds4_gpu_q4_table_model_residency_enabled() ||
getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) &&
getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL;
const bool enable_q4_expert_table =
getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL ||
q4_expert_table_auto;
const bool use_q4_expert_table =
!g_ssd_streaming_mode &&
!use_q4_grouped_experts &&
!use_q4_group6_experts &&
!use_q4_group8_experts &&
!use_q4_group24_experts &&
!use_q4_exact_tensor_id &&
!use_q4_expert_address_table &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert == 384 &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
fuse_pair_swiglu &&
direct_down_sum &&
g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_table_q4_k_sum6_pipeline != nil &&
g_moe_table_q4_pair_gate_encoder != nil &&
g_moe_table_q4_pair_up_encoder != nil &&
g_moe_table_q4_sum_down_encoder != nil &&
enable_q4_expert_table &&
(getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL ||
getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL ||
q4_table_queue_residency ||
ds4_gpu_q4_table_model_residency_enabled() ||
getenv("DS4_METAL_USE_QUEUE_RESIDENCY_SET") != NULL) &&
getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL;
const bool q4_table_boundary =
use_q4_expert_table &&
g_batch_cb != nil &&
getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL &&
!q4_table_queue_residency &&
getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL;
const bool enable_q4_gather_slots =
!g_ssd_streaming_mode &&
getenv("DS4_METAL_ENABLE_Q4_GATHER_SLOTS") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GATHER_SLOTS") == NULL;
const bool use_q4_gather_slots =
!use_q4_grouped_experts &&
!use_q4_group6_experts &&
!use_q4_group8_experts &&
!use_q4_group24_experts &&
!use_q4_exact_tensor_id &&
!use_q4_expert_address_table &&
!use_q4_expert_table &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert == q4_group6_expert_group_size * 6u &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
fuse_pair_swiglu &&
direct_down_sum &&
g_moe_q4_gather_slots6_pipeline != nil &&
g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil &&
enable_q4_gather_slots;
const bool use_q4_selected_slots =
!force_resident &&
!use_q4_grouped_experts &&
!use_q4_group6_experts &&
!use_q4_group8_experts &&
!use_q4_group24_experts &&
!use_q4_exact_tensor_id &&
!use_q4_expert_address_table &&
!use_q4_expert_table &&
!use_q4_gather_slots &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert >= 128 &&
(g_ssd_streaming_mode ||
(gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes)) &&
fuse_pair_swiglu &&
direct_down_sum &&
ds4_gpu_q4_selected_paths_allowed() &&
g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil &&
getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL;
const bool use_iq2_selected_slots =
!force_resident &&
g_ssd_streaming_mode &&
gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
down_type == DS4_METAL_TENSOR_Q2_K &&
n_expert == 6 &&
n_tokens == 1 &&
fuse_pair_swiglu &&
direct_down_sum &&
g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_slots6_q2_k_sum6_pipeline != nil &&
getenv("DS4_METAL_DISABLE_IQ2_SELECTED_EXPERT_VIEWS") == NULL;
const bool use_mxfp4_selected_slots =
!force_resident &&
g_ssd_streaming_mode &&
gate_type == DS4_METAL_TENSOR_MXFP4 &&
down_type == DS4_METAL_TENSOR_MXFP4 &&
n_expert == 6 &&
n_tokens == 1 &&
n_total_expert >= 128 &&
fuse_pair_swiglu &&
direct_down_sum &&
g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_slots6_mxfp4_sum6_pipeline != nil &&
getenv("DS4_METAL_DISABLE_MXFP4_SELECTED_EXPERT_VIEWS") == NULL;
const bool use_iq2_stream_addr_table =
!force_resident &&
g_ssd_streaming_mode &&
gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
down_type == DS4_METAL_TENSOR_IQ2_XXS &&
n_expert <= DS4_METAL_MAX_ROUTED_EXPERT_USED &&
n_tokens == 1 &&
fuse_pair_swiglu &&
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_iq2_xxs_pipeline != nil &&
getenv("DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE") == NULL;
const bool use_selected_slots =
use_q4_selected_slots || use_iq2_selected_slots ||
use_mxfp4_selected_slots || use_iq2_stream_addr_table;
id<MTLComputePipelineState> slots_pair_swiglu_pipeline =
use_iq2_selected_slots ? g_moe_mul_mv_slots6_iq2_xxs_pair_swiglu_pipeline :
(use_mxfp4_selected_slots ? g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline :
g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline);
id<MTLComputePipelineState> slots_sum6_pipeline =
use_iq2_selected_slots ? g_moe_mul_mv_slots6_q2_k_sum6_pipeline :
(use_mxfp4_selected_slots ? g_moe_mul_mv_slots6_mxfp4_sum6_pipeline :
g_moe_mul_mv_slots6_q4_k_sum6_pipeline);
const char *selected_profile_env = getenv("DS4_METAL_SELECTED_PROFILE");
if (!selected_profile_env) {
selected_profile_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE");
}
const char *selected_profile_layer_env = getenv("DS4_METAL_SELECTED_PROFILE_LAYER");
if (!selected_profile_layer_env) {
selected_profile_layer_env = getenv("DS4_METAL_Q4_SELECTED_PROFILE_LAYER");
}
bool selected_profile_layer_match = true;
if (selected_profile_layer_env && selected_profile_layer_env[0]) {
char *end = NULL;
long layer = strtol(selected_profile_layer_env, &end, 10);
selected_profile_layer_match =
end && *end == '\0' && layer >= 0 && (uint32_t)layer == layer_index;
}
const bool selected_profile =
use_selected_slots &&
selected_profile_env != NULL &&
selected_profile_layer_match;
const bool q4_selected_shared_event =
use_q4_selected_slots &&
getenv("DS4_METAL_Q4_SELECTED_SHARED_EVENT") != NULL;
const bool q4_selected_base_views =
use_q4_selected_slots &&
getenv("DS4_METAL_Q4_SELECTED_USE_BASE_VIEWS") != NULL &&
getenv("DS4_METAL_Q4_SELECTED_EXACT_VIEWS") == NULL;
const bool q4_selected_transient_views =
use_q4_selected_slots &&
!q4_selected_base_views &&
getenv("DS4_METAL_Q4_SELECTED_TRANSIENT_VIEWS") != NULL;
const char *q4_selected_view_mode =
q4_selected_base_views ? "base" :
(q4_selected_transient_views ? "transient" : "cached");
if (!use_selected_slots) {
g_routed_moe_selected_override_n = 0;
}
if (use_q4_expert_address_table) {
gate_table = ds4_gpu_q4_expert_address_table(model_map,
model_size,
gate_offset,
gate_expert_bytes,
n_total_expert);
up_table = ds4_gpu_q4_expert_address_table(model_map,
model_size,
up_offset,
gate_expert_bytes,
n_total_expert);
down_table = ds4_gpu_q4_expert_address_table(model_map,
model_size,
down_offset,
down_expert_bytes,
n_total_expert);
if (!gate_table || !up_table || !down_table) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32546);
return 0;
}
q4_table_layer_residency =
ds4_gpu_q4_expert_layer_residency_set(gate_table,
up_table,
down_table,
q4_expert_address_auto);
} else if (use_q4_expert_table) {
gate_table = ds4_gpu_q4_expert_table(model_map,
model_size,
gate_offset,
gate_expert_bytes,
n_total_expert,
g_moe_table_q4_pair_gate_encoder);
up_table = ds4_gpu_q4_expert_table(model_map,
model_size,
up_offset,
gate_expert_bytes,
n_total_expert,
g_moe_table_q4_pair_up_encoder);
down_table = ds4_gpu_q4_expert_table(model_map,
model_size,
down_offset,
down_expert_bytes,
n_total_expert,
g_moe_table_q4_sum_down_encoder);
if (!gate_table || !up_table || !down_table) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32573);
return 0;
}
q4_table_layer_residency =
ds4_gpu_q4_expert_layer_residency_set(gate_table,
up_table,
down_table,
q4_expert_table_auto);
} else if (use_q4_exact_tensor_id) {
gate_buf = ds4_gpu_wrap_model_exact_range(model_map,
model_size,
gate_offset,
gate_tensor_bytes,
&gate_inner);
up_buf = ds4_gpu_wrap_model_exact_range(model_map,
model_size,
up_offset,
gate_tensor_bytes,
&up_inner);
down_buf = ds4_gpu_wrap_model_exact_range(model_map,
model_size,
down_offset,
down_tensor_bytes,
&down_inner);
if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32601); return 0; }
} else if (use_q4_group6_experts || use_q4_gather_slots) {
if ((uint64_t)q4_group6_expert_group_size > UINT64_MAX / gate_expert_bytes ||
(uint64_t)q4_group6_expert_group_size > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal routed MoE Q4 group6 byte size overflow\n");
return 0;
}
const uint64_t gate_group_bytes =
(uint64_t)q4_group6_expert_group_size * gate_expert_bytes;
const uint64_t down_group_bytes =
(uint64_t)q4_group6_expert_group_size * down_expert_bytes;
for (uint32_t i = 0; i < 6; i++) {
const uint64_t gate_rel = (uint64_t)i * gate_group_bytes;
const uint64_t down_rel = (uint64_t)i * down_group_bytes;
if (gate_rel > gate_tensor_bytes ||
gate_group_bytes > gate_tensor_bytes - gate_rel ||
down_rel > down_tensor_bytes ||
down_group_bytes > down_tensor_bytes - down_rel ||
gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal routed MoE Q4 group6 offset overflow\n");
return 0;
}
uint64_t group_inner = 0;
gate_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map,
model_size,
gate_offset + gate_rel,
gate_group_bytes,
&group_inner);
gate_group6_offsets[i] = (NSUInteger)group_inner;
group_inner = 0;
up_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map,
model_size,
up_offset + gate_rel,
gate_group_bytes,
&group_inner);
up_group6_offsets[i] = (NSUInteger)group_inner;
group_inner = 0;
down_group6_bufs[i] = ds4_gpu_wrap_model_range(model_map,
model_size,
down_offset + down_rel,
down_group_bytes,
&group_inner);
down_group6_offsets[i] = (NSUInteger)group_inner;
if (!gate_group6_bufs[i] || !up_group6_bufs[i] || !down_group6_bufs[i]) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32643);
return 0;
}
}
if (use_q4_gather_slots) {
if (gate_expert_bytes > NSUIntegerMax ||
down_expert_bytes > NSUIntegerMax ||
gate_expert_bytes > UINT64_MAX / 6u ||
down_expert_bytes > UINT64_MAX / 6u ||
6ull * gate_expert_bytes > NSUIntegerMax ||
6ull * down_expert_bytes > NSUIntegerMax) {
fprintf(stderr, "ds4: Metal routed MoE Q4 gather scratch byte size overflow\n");
return 0;
}
const NSUInteger gate_slots_bytes = (NSUInteger)(6ull * gate_expert_bytes);
const NSUInteger down_slots_bytes = (NSUInteger)(6ull * down_expert_bytes);
if (!ds4_gpu_ensure_scratch_buffer(&g_moe_q4_gate_slots_buffer,
&g_moe_q4_gate_slots_bytes,
gate_slots_bytes,
"ds4_moe_q4_gate_slots") ||
!ds4_gpu_ensure_scratch_buffer(&g_moe_q4_up_slots_buffer,
&g_moe_q4_up_slots_bytes,
gate_slots_bytes,
"ds4_moe_q4_up_slots") ||
!ds4_gpu_ensure_scratch_buffer(&g_moe_q4_down_slots_buffer,
&g_moe_q4_down_slots_bytes,
down_slots_bytes,
"ds4_moe_q4_down_slots")) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32670);
return 0;
}
for (uint32_t i = 0; i < 6; i++) {
gate_slot_bufs[i] = g_moe_q4_gate_slots_buffer;
up_slot_bufs[i] = g_moe_q4_up_slots_buffer;
down_slot_bufs[i] = g_moe_q4_down_slots_buffer;
gate_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes);
up_slot_offsets[i] = (NSUInteger)((uint64_t)i * gate_expert_bytes);
down_slot_offsets[i] = (NSUInteger)((uint64_t)i * down_expert_bytes);
}
}
} else if (use_q4_group8_experts) {
if ((uint64_t)q4_group8_expert_group_size > UINT64_MAX / gate_expert_bytes ||
(uint64_t)q4_group8_expert_group_size > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal routed MoE Q4 group8 byte size overflow\n");
return 0;
}
const uint64_t gate_group_bytes =
(uint64_t)q4_group8_expert_group_size * gate_expert_bytes;
const uint64_t down_group_bytes =
(uint64_t)q4_group8_expert_group_size * down_expert_bytes;
for (uint32_t i = 0; i < 8; i++) {
const uint64_t gate_rel = (uint64_t)i * gate_group_bytes;
const uint64_t down_rel = (uint64_t)i * down_group_bytes;
if (gate_rel > gate_tensor_bytes ||
gate_group_bytes > gate_tensor_bytes - gate_rel ||
down_rel > down_tensor_bytes ||
down_group_bytes > down_tensor_bytes - down_rel ||
gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal routed MoE Q4 group8 offset overflow\n");
return 0;
}
uint64_t group_inner = 0;
gate_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map,
model_size,
gate_offset + gate_rel,
gate_group_bytes,
&group_inner);
gate_group8_offsets[i] = (NSUInteger)group_inner;
group_inner = 0;
up_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map,
model_size,
up_offset + gate_rel,
gate_group_bytes,
&group_inner);
up_group8_offsets[i] = (NSUInteger)group_inner;
group_inner = 0;
down_group8_bufs[i] = ds4_gpu_wrap_model_range(model_map,
model_size,
down_offset + down_rel,
down_group_bytes,
&group_inner);
down_group8_offsets[i] = (NSUInteger)group_inner;
if (!gate_group8_bufs[i] || !up_group8_bufs[i] || !down_group8_bufs[i]) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32727);
return 0;
}
}
} else if (use_q4_group24_experts) {
if ((uint64_t)q4_group24_expert_group_size > UINT64_MAX / gate_expert_bytes ||
(uint64_t)q4_group24_expert_group_size > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal routed MoE Q4 group24 byte size overflow\n");
return 0;
}
const uint64_t gate_group_bytes =
(uint64_t)q4_group24_expert_group_size * gate_expert_bytes;
const uint64_t down_group_bytes =
(uint64_t)q4_group24_expert_group_size * down_expert_bytes;
for (uint32_t i = 0; i < 24; i++) {
const uint64_t gate_rel = (uint64_t)i * gate_group_bytes;
const uint64_t down_rel = (uint64_t)i * down_group_bytes;
if (gate_rel > gate_tensor_bytes ||
gate_group_bytes > gate_tensor_bytes - gate_rel ||
down_rel > down_tensor_bytes ||
down_group_bytes > down_tensor_bytes - down_rel ||
gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal routed MoE Q4 group24 offset overflow\n");
return 0;
}
uint64_t group_inner = 0;
gate_group24_bufs[i] = q4_group24_exact_views ?
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
gate_offset + gate_rel,
gate_group_bytes,
&group_inner) :
ds4_gpu_wrap_model_range(model_map,
model_size,
gate_offset + gate_rel,
gate_group_bytes,
&group_inner);
gate_group24_offsets[i] = (NSUInteger)group_inner;
group_inner = 0;
up_group24_bufs[i] = q4_group24_exact_views ?
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
up_offset + gate_rel,
gate_group_bytes,
&group_inner) :
ds4_gpu_wrap_model_range(model_map,
model_size,
up_offset + gate_rel,
gate_group_bytes,
&group_inner);
up_group24_offsets[i] = (NSUInteger)group_inner;
group_inner = 0;
down_group24_bufs[i] = q4_group24_exact_views ?
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
down_offset + down_rel,
down_group_bytes,
&group_inner) :
ds4_gpu_wrap_model_range(model_map,
model_size,
down_offset + down_rel,
down_group_bytes,
&group_inner);
down_group24_offsets[i] = (NSUInteger)group_inner;
if (!gate_group24_bufs[i] || !up_group24_bufs[i] || !down_group24_bufs[i]) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32794);
return 0;
}
}
} else if (use_selected_slots) {
const bool selected_timing =
selected_profile ||
ds4_gpu_stream_expert_timing_summary_enabled();
double selected_t0 = selected_timing ? ds4_gpu_now_ms() : 0.0;
double selected_read_ms = 0.0;
double selected_sync_ms = 0.0;
double selected_copy_ms = 0.0;
double selected_wrap_ms = 0.0;
uint64_t selected_cache_hits0 = g_stream_expert_cache_hits;
uint64_t selected_cache_misses0 = g_stream_expert_cache_misses;
uint64_t selected_cache_wraps0 = g_stream_expert_cache_wraps;
uint64_t selected_cache_evictions0 = g_stream_expert_cache_evictions;
const char *selected_id_source = "readback";
bool selected_ids_available = true;
bool selected_exec_ids_from_host = false;
const int stream_expert_cache_size_known =
ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes);
const bool use_iq2_full_expert_addr_table =
use_iq2_selected_slots &&
ds4_gpu_stream_full_expert_addr_table_requested() &&
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil;
use_stream_expert_cache =
!use_iq2_full_expert_addr_table &&
(use_iq2_selected_slots || use_iq2_stream_addr_table ||
use_q4_selected_slots || use_mxfp4_selected_slots) &&
stream_expert_cache_size_known &&
ds4_gpu_stream_expert_cache_effective_cap(layer_index,
n_total_expert,
n_expert) != 0;
if (use_iq2_stream_addr_table && !use_stream_expert_cache) {
fprintf(stderr,
"ds4: Metal IQ2/IQ2 streaming decode requires a non-empty expert cache\n");
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32831);
return 0;
}
const bool stream_split_ready =
use_stream_expert_cache &&
ds4_gpu_stream_expert_split_ready();
const bool use_stream_compact_addr =
use_stream_expert_cache &&
use_iq2_selected_slots &&
ds4_gpu_stream_compact_addr_requested() &&
!stream_split_ready &&
!ds4_gpu_stream_expert_masked_addr_requested() &&
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil;
use_stream_expert_split_candidate =
use_stream_expert_cache &&
use_iq2_selected_slots &&
!use_stream_compact_addr &&
stream_split_ready &&
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil &&
g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil;
const bool use_stream_hit_validator =
use_stream_expert_cache &&
use_iq2_selected_slots &&
ds4_gpu_stream_expert_hit_validator_requested() &&
g_moe_stream_expert_cache_validate_pipeline != nil &&
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil &&
ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf);
if (use_iq2_full_expert_addr_table) {
selected_id_source = "gpu-full-addr";
selected_ids_available = false;
g_routed_moe_selected_override_n = 0;
ds4_gpu_stream_expert_cache_entry *full_entry = NULL;
if (!ds4_gpu_stream_full_expert_addr_table_prepare(model_map,
model_size,
layer_index,
n_total_expert,
gate_offset,
up_offset,
down_offset,
gate_expert_bytes,
down_expert_bytes,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf,
&full_entry)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32880);
return 0;
}
for (uint32_t i = 0; i < n_expert; i++) {
stream_slot_entries[i] = full_entry;
}
use_stream_expert_addr_table = true;
} else {
const int replayed_selected_ids =
ds4_gpu_moe_selected_trace_replay(selected_ids, n_expert);
if (replayed_selected_ids < 0) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32890);
return 0;
} else if (replayed_selected_ids > 0) {
selected_id_source = "replay";
selected_exec_ids_from_host = true;
g_routed_moe_selected_override_n = 0;
} else if (g_routed_moe_selected_override_n == n_expert) {
memcpy(selected_ids,
g_routed_moe_selected_override,
(size_t)n_expert * sizeof(selected_ids[0]));
selected_id_source = "override";
selected_exec_ids_from_host = true;
g_routed_moe_selected_override_n = 0;
} else if (use_stream_expert_cache &&
ds4_gpu_glm_stream_selected_prefetch_take(
model_map, model_size, layer_index,
n_total_expert, n_expert,
gate_offset, up_offset, down_offset,
gate_expert_bytes, down_expert_bytes,
selected_ids)) {
/* GLM already read these IDs to start loading experts.
* Submit the shared expert work while the reads finish. */
selected_id_source = "prefetched";
if (g_batch_cb && g_batch_has_work &&
g_stream_expert_pending_load.active &&
!ds4_gpu_flush_commands()) return 0;
} else if (use_stream_hit_validator) {
g_routed_moe_selected_override_n = 0;
uint32_t validator_all_cached = 0;
uint32_t validator_miss_mask = 0;
uint32_t validator_invalid_mask = 0;
if (!ds4_gpu_stream_expert_cache_validate_selected(selected,
stream_gate_addr_buf,
stream_up_addr_buf,
stream_down_addr_buf,
n_total_expert,
n_expert,
selected_ids,
&validator_all_cached,
&validator_miss_mask,
&validator_invalid_mask)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32917);
return 0;
}
selected_id_source =
validator_invalid_mask != 0 ? "validator-invalid" :
(validator_miss_mask != 0 ? "validator-miss" :
(validator_all_cached != 0 ? "validator-hit" : "validator-miss"));
} else {
g_routed_moe_selected_override_n = 0;
if (g_batch_cb != nil) {
double selected_boundary_t0 =
selected_timing ? ds4_gpu_now_ms() : 0.0;
if (q4_selected_shared_event) {
if (ds4_gpu_signal_batch_and_wait_event("selected-id readback") == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32942); return 0; }
} else if (ds4_gpu_end_commands() == 0) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32931);
return 0;
}
if (selected_timing) {
selected_sync_ms +=
ds4_gpu_now_ms() - selected_boundary_t0;
}
double selected_copy_t0 =
selected_timing ? ds4_gpu_now_ms() : 0.0;
if (ds4_gpu_tensor_read(selected,
0,
selected_ids,
(uint64_t)n_expert * sizeof(selected_ids[0])) == 0) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32943);
return 0;
}
if (selected_timing) {
selected_copy_ms +=
ds4_gpu_now_ms() - selected_copy_t0;
}
if (!q4_selected_shared_event) {
selected_boundary_t0 =
selected_timing ? ds4_gpu_now_ms() : 0.0;
if (ds4_gpu_begin_commands() == 0) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32967); return 0; }
if (selected_timing) {
selected_sync_ms +=
ds4_gpu_now_ms() - selected_boundary_t0;
}
}
} else {
double selected_copy_t0 =
selected_timing ? ds4_gpu_now_ms() : 0.0;
if (ds4_gpu_tensor_read(selected,
0,
selected_ids,
(uint64_t)n_expert * sizeof(selected_ids[0])) == 0) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32965);
return 0;
}
if (selected_timing) {
selected_copy_ms +=
ds4_gpu_now_ms() - selected_copy_t0;
}
}
}
}
if (selected_timing) {
selected_read_ms = selected_sync_ms + selected_copy_ms;
selected_t0 = ds4_gpu_now_ms();
}
if (selected_ids_available) {
for (uint32_t i = 0; i < n_expert; i++) {
if (selected_ids[i] < 0 || (uint32_t)selected_ids[i] >= n_total_expert) {
fprintf(stderr,
"ds4: Metal routed MoE selected expert id %d is outside 0..%u\n",
selected_ids[i],
n_total_expert);
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32986);
return 0;
}
}
ds4_gpu_stream_expert_cache_note_selected_hotness(layer_index,
selected_ids,
n_expert);
if (!ds4_gpu_moe_selected_trace_record(selected_ids, n_expert) ||
!ds4_gpu_moe_selected_hotlist_record(layer_index,
selected_ids,
n_expert,
n_total_expert)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 32997);
return 0;
}
for (uint32_t i = 0; i < n_expert; i++) {
const uint64_t expert_id = (uint64_t)(uint32_t)selected_ids[i];
if (expert_id > UINT64_MAX / gate_expert_bytes ||
expert_id > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n");
return 0;
}
const uint64_t gate_rel = expert_id * gate_expert_bytes;
const uint64_t down_rel = expert_id * down_expert_bytes;
if (gate_rel > UINT64_MAX - gate_offset ||
gate_rel > UINT64_MAX - up_offset ||
down_rel > UINT64_MAX - down_offset) {
fprintf(stderr, "ds4: Metal routed MoE selected expert offset overflow\n");
return 0;
}
stream_gate_abs_offsets[i] = gate_offset + gate_rel;
stream_up_abs_offsets[i] = up_offset + gate_rel;
stream_down_abs_offsets[i] = down_offset + down_rel;
if (use_stream_expert_cache) {
ds4_gpu_stream_expert_cache_entry *entry = NULL;
entry = ds4_gpu_stream_expert_cache_peek(model_map,
model_size,
layer_index,
(uint32_t)selected_ids[i],
n_total_expert,
n_expert,
stream_gate_abs_offsets[i],
stream_up_abs_offsets[i],
stream_down_abs_offsets[i],
gate_expert_bytes,
down_expert_bytes);
if (!entry) {
stream_expert_missing_mask |= 1u << i;
continue;
}
stream_expert_resident_mask |= 1u << i;
stream_slot_entries[i] = entry;
gate_slot_bufs[i] = entry->gate_buffer;
gate_slot_offsets[i] = entry->gate_inner;
up_slot_bufs[i] = entry->up_buffer;
up_slot_offsets[i] = entry->up_inner;
down_slot_bufs[i] = entry->down_buffer;
down_slot_offsets[i] = entry->down_inner;
continue;
}
uint64_t slot_inner = 0;
gate_slot_bufs[i] = q4_selected_base_views ?
ds4_gpu_wrap_model_range(model_map,
model_size,
gate_offset + gate_rel,
gate_expert_bytes,
&slot_inner) :
(q4_selected_transient_views ?
ds4_gpu_wrap_model_exact_range_transient(model_map,
model_size,
gate_offset + gate_rel,
gate_expert_bytes,
&slot_inner) :
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
gate_offset + gate_rel,
gate_expert_bytes,
&slot_inner));
gate_slot_offsets[i] = (NSUInteger)slot_inner;
slot_inner = 0;
up_slot_bufs[i] = q4_selected_base_views ?
ds4_gpu_wrap_model_range(model_map,
model_size,
up_offset + gate_rel,
gate_expert_bytes,
&slot_inner) :
(q4_selected_transient_views ?
ds4_gpu_wrap_model_exact_range_transient(model_map,
model_size,
up_offset + gate_rel,
gate_expert_bytes,
&slot_inner) :
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
up_offset + gate_rel,
gate_expert_bytes,
&slot_inner));
up_slot_offsets[i] = (NSUInteger)slot_inner;
slot_inner = 0;
down_slot_bufs[i] = q4_selected_base_views ?
ds4_gpu_wrap_model_range(model_map,
model_size,
down_offset + down_rel,
down_expert_bytes,
&slot_inner) :
(q4_selected_transient_views ?
ds4_gpu_wrap_model_exact_range_transient(model_map,
model_size,
down_offset + down_rel,
down_expert_bytes,
&slot_inner) :
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
down_offset + down_rel,
down_expert_bytes,
&slot_inner));
down_slot_offsets[i] = (NSUInteger)slot_inner;
if (!gate_slot_bufs[i] || !up_slot_bufs[i] || !down_slot_bufs[i]) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33105);
return 0;
}
}
}
if (use_stream_expert_cache) {
use_stream_expert_addr_table =
((use_iq2_selected_slots &&
ds4_gpu_stream_expert_addr_table_kernel_requested() &&
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil) ||
use_iq2_stream_addr_table) &&
ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf);
if (use_iq2_stream_addr_table && !use_stream_expert_addr_table) {
fprintf(stderr,
"ds4: Metal IQ2/IQ2 streaming decode could not prepare expert address buffers\n");
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33123);
return 0;
}
use_stream_expert_masked_addr_table =
use_stream_expert_addr_table &&
use_iq2_selected_slots &&
ds4_gpu_stream_expert_masked_addr_requested() &&
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline != nil &&
g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline != nil;
use_stream_expert_split_deferred =
use_stream_expert_split_candidate &&
use_stream_expert_masked_addr_table &&
stream_expert_resident_mask != 0 &&
stream_expert_missing_mask != 0 &&
ds4_gpu_stream_expert_split_worthwhile(stream_expert_resident_mask,
stream_expert_missing_mask) &&
g_batch_cb != nil &&
getenv("DS4_METAL_MOE_ONE_STAGE_PROFILE") == NULL;
if (use_stream_expert_split_deferred) {
const ds4_gpu_stream_expert_table table = {
.model_map = model_map,
.model_size = model_size,
.layer = layer_index,
.n_total_expert = n_total_expert,
.gate_offset = gate_offset,
.up_offset = up_offset,
.down_offset = down_offset,
.gate_expert_bytes = gate_expert_bytes,
.down_expert_bytes = down_expert_bytes,
};
if (!ds4_gpu_stream_expert_cache_begin_selected_load(
&table,
selected_ids,
n_expert)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33156);
return 0;
}
}
if (stream_expert_missing_mask != 0 &&
!use_stream_expert_split_deferred) {
if (!ds4_gpu_stream_expert_cache_load_selected_missing(
model_map,
model_size,
layer_index,
selected_ids,
n_total_expert,
n_expert,
stream_gate_abs_offsets,
stream_up_abs_offsets,
stream_down_abs_offsets,
gate_expert_bytes,
down_expert_bytes,
stream_expert_missing_mask,
stream_slot_entries)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33175);
return 0;
}
for (uint32_t i = 0; i < n_expert; i++) {
if ((stream_expert_missing_mask & (1u << i)) == 0) continue;
ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i];
if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202); return 0; }
gate_slot_bufs[i] = entry->gate_buffer;
gate_slot_offsets[i] = entry->gate_inner;
up_slot_bufs[i] = entry->up_buffer;
up_slot_offsets[i] = entry->up_inner;
down_slot_bufs[i] = entry->down_buffer;
down_slot_offsets[i] = entry->down_inner;
}
}
if (use_iq2_stream_addr_table && use_stream_expert_addr_table) {
for (uint32_t i = 0; i < n_expert; i++) {
ds4_gpu_stream_expert_cache_entry *entry = stream_slot_entries[i];
if (!entry) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33214); return 0; }
if (!ds4_gpu_stream_expert_cache_set_addr_slot_raw(
layer_index,
(uint32_t)selected_ids[i],
entry->gate_buffer,
entry->gate_inner,
entry->up_buffer,
entry->up_inner,
entry->down_buffer,
entry->down_inner)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33202);
return 0;
}
}
}
ds4_gpu_stream_expert_cache_prune_layer(layer_index,
n_total_expert,
n_expert,
selected_ids,
n_expert);
ds4_gpu_stream_expert_cache_prune_global(layer_index,
selected_ids,
n_expert);
if (use_stream_compact_addr) {
id<MTLBuffer> compact_selected = nil;
if (!ds4_gpu_stream_compact_addr_prepare(layer_index,
stream_slot_entries,
n_expert,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf,
&compact_selected)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33223);
return 0;
}
selected_exec_buf = compact_selected;
selected_exec_off = 0;
use_stream_expert_addr_table = true;
use_stream_expert_masked_addr_table = false;
use_stream_compact_addr_table = true;
}
if (use_stream_expert_addr_table &&
!use_stream_compact_addr_table &&
selected_exec_ids_from_host) {
if (!ds4_gpu_stream_selected_ids_prepare(layer_index,
selected_ids,
n_expert,
&selected_exec_buf,
&selected_exec_off)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33239);
return 0;
}
}
}
if (selected_timing) {
selected_wrap_ms = ds4_gpu_now_ms() - selected_t0;
if (use_stream_expert_cache) {
ds4_gpu_stream_expert_timing_note_cache_class(
stream_expert_resident_mask,
stream_expert_missing_mask);
}
ds4_gpu_stream_expert_timing_note_selected(selected_sync_ms,
selected_copy_ms,
selected_wrap_ms);
}
if (selected_profile) {
const uint64_t selected_cache_hits =
g_stream_expert_cache_hits - selected_cache_hits0;
const uint64_t selected_cache_misses =
g_stream_expert_cache_misses - selected_cache_misses0;
const uint64_t selected_cache_wraps =
g_stream_expert_cache_wraps - selected_cache_wraps0;
const uint64_t selected_cache_evictions =
g_stream_expert_cache_evictions - selected_cache_evictions0;
const char *selected_path =
use_iq2_stream_addr_table ? "iq2/iq2" :
(use_iq2_selected_slots ? "iq2/q2" :
(use_mxfp4_selected_slots ? "mxfp4/mxfp4" : "q4/q4"));
const char *selected_view_mode =
use_stream_expert_split_deferred ? "stream-split" :
use_stream_expert_masked_addr_table ? "stream-addr-mask" :
use_stream_compact_addr_table ? "stream-compact-addr" :
use_stream_expert_addr_table ? "stream-addr" :
(use_stream_expert_cache ? "stream-cache" :
(use_iq2_selected_slots ? "exact-cache" : q4_selected_view_mode));
fprintf(stderr,
"ds4: Metal selected views layer=%u path=%s mode=%s ids=%s "
"experts=",
layer_index, selected_path, selected_view_mode,
selected_id_source);
for (uint32_t i = 0; i < n_expert; i++)
fprintf(stderr, "%s%d", i ? "," : "",
selected_ids_available ? selected_ids[i] : -1);
fprintf(stderr,
" expert_gate=%.2f MiB "
"expert_down=%.2f MiB read=%.3f ms bind=%.3f ms "
"cache_hits=%llu cache_misses=%llu cache_wraps=%llu cache_evictions=%llu\n",
ds4_gpu_mib(gate_expert_bytes),
ds4_gpu_mib(down_expert_bytes),
selected_read_ms,
selected_wrap_ms,
(unsigned long long)selected_cache_hits,
(unsigned long long)selected_cache_misses,
(unsigned long long)selected_cache_wraps,
(unsigned long long)selected_cache_evictions);
}
} else if (!use_q4_grouped_experts) {
gate_buf = ds4_gpu_wrap_model_range(model_map, model_size, gate_offset, gate_tensor_bytes, &gate_inner);
up_buf = ds4_gpu_wrap_model_range(model_map, model_size, up_offset, gate_tensor_bytes, &up_inner);
down_buf = ds4_gpu_wrap_model_range(model_map, model_size, down_offset, down_tensor_bytes, &down_inner);
if (!gate_buf || !up_buf || !down_buf) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33326); return 0; }
}
if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33305);
return 0;
}
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33337); return 0; }
if ((use_q4_expert_address_table || use_q4_expert_table) &&
!ds4_gpu_use_model_residency_set(cb)) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33314);
return 0;
}
if (!q4_table_queue_residency &&
q4_table_layer_residency &&
[cb respondsToSelector:@selector(useResidencySet:)]) {
[cb useResidencySet:q4_table_layer_residency];
}
const bool moe_one_stage_profile =
g_batch_cb != nil &&
ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_ONE_STAGE_PROFILE",
"DS4_METAL_MOE_ONE_STAGE_PROFILE_LAYER",
layer_index);
const char *moe_one_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER");
const char *moe_one_path =
use_q4_grouped_experts ? "q4_grouped_pair_swiglu" :
use_q4_group6_experts ? "q4_group6_pair_swiglu" :
use_q4_group8_experts ? "q4_group8_pair_swiglu" :
use_q4_group24_experts ? "q4_group24_split_gate_up" :
use_q4_exact_tensor_id ? "q4_exact_pair_swiglu" :
use_q4_expert_address_table ? "q4_addr_pair_swiglu" :
use_q4_expert_table ? "q4_table_pair_swiglu" :
use_q4_gather_slots ? "q4_gather_slots6_pair_swiglu" :
use_stream_expert_split_deferred ? "iq2_stream_split_pair_swiglu" :
use_stream_expert_masked_addr_table ? "iq2_stream_addr_mask_pair_swiglu" :
use_stream_expert_addr_table ? "iq2_stream_addr_pair_swiglu" :
use_iq2_selected_slots ? "iq2_slots6_pair_swiglu" :
use_mxfp4_selected_slots ? "mxfp4_slots6_pair_swiglu" :
use_q4_selected_slots ? "q4_slots6_pair_swiglu" :
(fuse_pair_swiglu ? "pair_swiglu" :
((!g_quality_mode &&
((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) ||
(gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline))) ? "pair" : "single"));
double moe_one_stage_t0 = moe_one_stage_profile ? ds4_gpu_now_ms() : 0.0;
if (moe_one_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33349);
return 0;
}
cb = ds4_gpu_command_buffer(&owned);
if (!cb) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33380); return 0; }
moe_one_stage_t0 = ds4_gpu_now_ms();
}
#define DS4_METAL_PROFILE_MOE_ONE_STAGE(name) do { \
if (ok && moe_one_stage_profile) { \
if (ds4_gpu_end_commands() == 0) { \
ok = 0; \
} else { \
const char *stage_name = (name); \
const double now_ms = ds4_gpu_now_ms(); \
const int print_stage = \
!moe_one_stage_filter || !moe_one_stage_filter[0] || \
strstr(stage_name, moe_one_stage_filter) != NULL; \
if (print_stage) { \
fprintf(stderr, \
"ds4: Metal routed MoE one stage layer=%u pairs=%u experts=%u " \
"gate=%s down=%s path=%s %s=%.3f ms\n", \
layer_index, pair_rows, n_expert, \
ds4_gpu_metal_tensor_type_name(gate_type), \
ds4_gpu_metal_tensor_type_name(down_type), \
moe_one_path, \
stage_name, now_ms - moe_one_stage_t0); \
} \
moe_one_stage_t0 = now_ms; \
if (ds4_gpu_begin_commands() == 0) { \
ok = 0; \
} else { \
cb = ds4_gpu_command_buffer(&owned); \
if (!cb) ok = 0; \
} \
} \
} \
} while (0)
if (use_q4_gather_slots) {
ds4_gpu_q4_gather_slots6_args gate_gather_args = {
.expert_bytes = gate_expert_bytes,
.group_size = q4_group6_expert_group_size,
.n_slots = n_expert,
};
ds4_gpu_q4_gather_slots6_args down_gather_args = {
.expert_bytes = down_expert_bytes,
.group_size = q4_group6_expert_group_size,
.n_slots = n_expert,
};
ok = ds4_gpu_encode_q4_gather_slots6(cb,
g_moe_q4_gather_slots6_pipeline,
&gate_gather_args,
gate_group6_bufs,
gate_group6_offsets,
selectedbuf,
ds4_gpu_tensor_offset(selected),
g_moe_q4_gate_slots_buffer,
0) &&
ds4_gpu_encode_q4_gather_slots6(cb,
g_moe_q4_gather_slots6_pipeline,
&gate_gather_args,
up_group6_bufs,
up_group6_offsets,
selectedbuf,
ds4_gpu_tensor_offset(selected),
g_moe_q4_up_slots_buffer,
0) &&
ds4_gpu_encode_q4_gather_slots6(cb,
g_moe_q4_gather_slots6_pipeline,
&down_gather_args,
down_group6_bufs,
down_group6_offsets,
selectedbuf,
ds4_gpu_tensor_offset(selected),
g_moe_q4_down_slots_buffer,
0);
}
if (use_q4_gather_slots) {
DS4_METAL_PROFILE_MOE_ONE_STAGE("q4_gather");
}
if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 33455); return 0; }
if (use_q4_grouped_experts) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) {
const uint32_t expert_count =
q4_expert_group_size < n_total_expert - expert_base ?
q4_expert_group_size : n_total_expert - expert_base;
if ((uint64_t)expert_base > UINT64_MAX / gate_expert_bytes ||
(uint64_t)expert_count > UINT64_MAX / gate_expert_bytes) {
ok = 0;
break;
}
const uint64_t group_rel = (uint64_t)expert_base * gate_expert_bytes;
const uint64_t group_bytes = (uint64_t)expert_count * gate_expert_bytes;
if (group_rel > UINT64_MAX - gate_offset ||
group_rel > UINT64_MAX - up_offset) {
ok = 0;
break;
}
uint64_t gate_group_inner = 0;
uint64_t up_group_inner = 0;
id<MTLBuffer> gate_group_buf =
q4_grouped_cache_views ?
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
gate_offset + group_rel,
group_bytes,
&gate_group_inner) :
ds4_gpu_wrap_model_exact_range_transient(model_map,
model_size,
gate_offset + group_rel,
group_bytes,
&gate_group_inner);
id<MTLBuffer> up_group_buf =
q4_grouped_cache_views ?
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
up_offset + group_rel,
group_bytes,
&up_group_inner) :
ds4_gpu_wrap_model_exact_range_transient(model_map,
model_size,
up_offset + group_rel,
group_bytes,
&up_group_inner);
if (!gate_group_buf || !up_group_buf) {
ok = 0;
break;
}
ds4_gpu_moe_expert_group_args group_args = {
.expert_base = expert_base,
.expert_count = expert_count,
.accumulate = 0,
.pad0 = 0,
};
ok = ds4_gpu_encode_mul_mv_group_q4_pair_swiglu(cb,
g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline,
&gate_args,
&act_args,
&group_args,
gate_group_buf,
(NSUInteger)gate_group_inner,
up_group_buf,
(NSUInteger)up_group_inner,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false);
}
} else if (use_q4_expert_address_table) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mv_addr_q4_pair_swiglu(cb,
g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline,
&gate_args,
&act_args,
gate_table,
up_table,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false);
} else if (use_q4_expert_table) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb,
g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline,
&gate_args,
&act_args,
gate_table,
up_table,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false,
q4_table_queue_residency);
} else if (use_q4_group6_experts) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mv_group6_pair_swiglu(cb,
g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline,
&gate_args,
&act_args,
gate_group6_bufs,
gate_group6_offsets,
up_group6_bufs,
up_group6_offsets,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false);
} else if (use_q4_group8_experts) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mv_group8_pair_swiglu(cb,
g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline,
&gate_args,
&act_args,
gate_group8_bufs,
gate_group8_offsets,
up_group8_bufs,
up_group8_offsets,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false);
} else if (use_q4_group24_experts) {
ok = ds4_gpu_encode_mul_mv_group24_id(cb,
g_moe_mul_mv_group24_q4_k_id_pipeline,
&gate_args,
gate_group24_bufs,
gate_group24_offsets,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
2,
false) &&
ds4_gpu_encode_mul_mv_group24_id(cb,
g_moe_mul_mv_group24_q4_k_id_pipeline,
&gate_args,
up_group24_bufs,
up_group24_offsets,
xbuf,
ds4_gpu_tensor_offset(x),
upbuf,
ds4_gpu_tensor_offset(up),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
2,
false);
} else if (use_q4_gather_slots || use_selected_slots) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
if (use_stream_expert_addr_table) {
if (use_stream_expert_masked_addr_table) {
if (use_stream_expert_split_deferred) {
const bool stream_split_profile =
getenv("DS4_METAL_STREAMING_EXPERT_SPLIT_PROFILE") != NULL;
const bool stream_split_timing =
stream_split_profile ||
ds4_gpu_stream_expert_timing_summary_enabled();
double stream_split_t0 =
stream_split_timing ? ds4_gpu_now_ms() : 0.0;
ds4_gpu_stream_expert_split_args resident_pair_args = {
.active_mask = stream_expert_resident_mask,
.accumulate = 0u,
};
ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb,
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline,
&gate_args,
&act_args,
&resident_pair_args,
stream_slot_entries,
stream_gate_addr_buf,
stream_up_addr_buf,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selected_exec_buf,
selected_exec_off,
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false);
if (ok) {
ok = ds4_gpu_flush_commands();
if (ok) {
cb = ds4_gpu_command_buffer(&owned);
if (!cb) ok = 0;
}
}
const double stream_split_resident_ms =
stream_split_timing ? ds4_gpu_now_ms() - stream_split_t0 : 0.0;
if (stream_split_timing) stream_split_t0 = ds4_gpu_now_ms();
const double stream_split_missing_start_ms = stream_split_t0;
double stream_split_missing_load_ms = 0.0;
double stream_split_missing_slot_ms = 0.0;
double stream_split_missing_prune_ms = 0.0;
double stream_split_missing_addr_ms = 0.0;
double stream_split_missing_wait_ms = 0.0;
if (ok) {
ok = ds4_gpu_stream_expert_cache_load_selected_missing(
model_map,
model_size,
layer_index,
selected_ids,
n_total_expert,
n_expert,
stream_gate_abs_offsets,
stream_up_abs_offsets,
stream_down_abs_offsets,
gate_expert_bytes,
down_expert_bytes,
stream_expert_missing_mask,
stream_slot_entries);
if (stream_split_timing) {
const double now_ms = ds4_gpu_now_ms();
stream_split_missing_load_ms =
now_ms - stream_split_t0;
stream_split_t0 = now_ms;
}
if (ok) {
for (uint32_t i = 0; i < n_expert; i++) {
if ((stream_expert_missing_mask & (1u << i)) == 0) continue;
ds4_gpu_stream_expert_cache_entry *entry =
stream_slot_entries[i];
if (!entry) {
ok = 0;
break;
}
gate_slot_bufs[i] = entry->gate_buffer;
gate_slot_offsets[i] = entry->gate_inner;
up_slot_bufs[i] = entry->up_buffer;
up_slot_offsets[i] = entry->up_inner;
down_slot_bufs[i] = entry->down_buffer;
down_slot_offsets[i] = entry->down_inner;
}
}
}
if (stream_split_timing) {
const double now_ms = ds4_gpu_now_ms();
stream_split_missing_slot_ms =
now_ms - stream_split_t0;
stream_split_t0 = now_ms;
}
if (ok) {
ds4_gpu_stream_expert_cache_prune_layer(layer_index,
n_total_expert,
n_expert,
selected_ids,
n_expert);
ds4_gpu_stream_expert_cache_prune_global(layer_index,
selected_ids,
n_expert);
if (stream_split_timing) {
const double now_ms = ds4_gpu_now_ms();
stream_split_missing_prune_ms =
now_ms - stream_split_t0;
stream_split_t0 = now_ms;
}
ok = ds4_gpu_stream_expert_cache_addr_buffers(layer_index,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf);
if (stream_split_timing) {
const double now_ms = ds4_gpu_now_ms();
stream_split_missing_addr_ms =
now_ms - stream_split_t0;
stream_split_t0 = now_ms;
}
}
if (ok) {
/*
* The resident stage was submitted before the
* CPU read of missing experts so I/O can overlap
* with GPU work. The missing stage reuses the same
* gate/up/mid scratch buffers, so it must not
* execute until the resident command buffer has
* finished. The down/sum pass is issued once after
* all six mid slots exist; this keeps the final
* accumulation order stable regardless of the
* resident/missing split.
*/
ok = ds4_gpu_wait_pending_command_buffers(
"streaming expert split resident");
if (stream_split_timing) {
const double now_ms = ds4_gpu_now_ms();
stream_split_missing_wait_ms =
now_ms - stream_split_t0;
stream_split_t0 = now_ms;
}
}
const double stream_split_missing_ms =
stream_split_timing ?
ds4_gpu_now_ms() - stream_split_missing_start_ms :
0.0;
ds4_gpu_stream_expert_split_args missing_pair_args = {
.active_mask = stream_expert_missing_mask,
.accumulate = 0u,
};
ds4_gpu_stream_expert_split_args all_down_args = {
.active_mask = 0x3fu,
.accumulate = 0u,
};
if (ok) {
ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb,
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline,
&gate_args,
&act_args,
&missing_pair_args,
stream_slot_entries,
stream_gate_addr_buf,
stream_up_addr_buf,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selected_exec_buf,
selected_exec_off,
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false);
}
if (ok) {
ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb,
g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline,
&down_args,
&all_down_args,
stream_slot_entries,
stream_down_addr_buf,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selected_exec_buf,
selected_exec_off,
down_smem,
2);
}
stream_expert_split_completed = ok;
if (stream_split_timing) {
ds4_gpu_stream_expert_timing_note_split(
stream_expert_resident_mask,
stream_expert_missing_mask,
stream_split_resident_ms,
stream_split_missing_ms);
ds4_gpu_stream_expert_timing_note_split_missing_detail(
stream_split_missing_load_ms,
stream_split_missing_slot_ms,
stream_split_missing_prune_ms,
stream_split_missing_addr_ms,
stream_split_missing_wait_ms);
}
if (stream_split_profile) {
fprintf(stderr,
"ds4: Metal streaming expert split layer=%u "
"resident=0x%02x missing=0x%02x resident_submit=%.3f ms "
"missing_bind=%.3f ms\n",
layer_index,
stream_expert_resident_mask,
stream_expert_missing_mask,
stream_split_resident_ms,
stream_split_missing_ms);
}
} else {
ds4_gpu_stream_expert_split_args split_args = {
.active_mask = 0x3fu,
.accumulate = 0u,
};
ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu_masked(cb,
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_masked_pipeline,
&gate_args,
&act_args,
&split_args,
stream_slot_entries,
stream_gate_addr_buf,
stream_up_addr_buf,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selected_exec_buf,
selected_exec_off,
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false);
}
} else {
ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(cb,
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline,
&gate_args,
&act_args,
stream_slot_entries,
n_expert,
stream_gate_addr_buf,
stream_up_addr_buf,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selected_exec_buf,
selected_exec_off,
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false,
nil,
nil);
}
} else {
ok = (!use_stream_expert_cache ||
ds4_gpu_stream_expert_cache_mark_entries_inflight(
stream_slot_entries,
n_expert,
0)) &&
ds4_gpu_encode_mul_mv_slots6_pair_swiglu(cb,
slots_pair_swiglu_pipeline,
&gate_args,
&act_args,
gate_slot_bufs,
gate_slot_offsets,
up_slot_bufs,
up_slot_offsets,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false);
}
} else if (fuse_pair_swiglu) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mv_id_pair_swiglu(cb,
pair_swiglu_pipeline,
&gate_args,
&act_args,
gate_buf,
(NSUInteger)gate_inner,
up_buf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
pair_swiglu_nsg,
false);
} else if (!g_quality_mode &&
gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
g_moe_mul_mv_id_iq2_xxs_pair_pipeline) {
ok = ds4_gpu_encode_mul_mv_id_pair(cb,
g_moe_mul_mv_id_iq2_xxs_pair_pipeline,
&gate_args,
gate_buf,
(NSUInteger)gate_inner,
up_buf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
2,
false);
} else if (!g_quality_mode &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
g_moe_mul_mv_id_q4_k_pair_pipeline) {
ok = ds4_gpu_encode_mul_mv_id_pair(cb,
g_moe_mul_mv_id_q4_k_pair_pipeline,
&gate_args,
gate_buf,
(NSUInteger)gate_inner,
up_buf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
2,
false);
} else {
ok = ds4_gpu_encode_mul_mv_id(cb,
gate_mv_pipeline,
&gate_args,
gate_buf,
(NSUInteger)gate_inner,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
gate_nsg,
gate_rows_per_group_is_nr0) &&
ds4_gpu_encode_mul_mv_id(cb,
gate_mv_pipeline,
&gate_args,
up_buf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
upbuf,
ds4_gpu_tensor_offset(up),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
gate_nsg,
gate_rows_per_group_is_nr0);
}
DS4_METAL_PROFILE_MOE_ONE_STAGE("gate_up");
if (ok && (!fuse_pair_swiglu || use_q4_group24_experts)) {
ok = ds4_gpu_encode_moe_swiglu_weight(cb,
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
weightsbuf,
ds4_gpu_tensor_offset(weights),
expert_mid_dim,
pair_rows,
clamp,
false);
}
DS4_METAL_PROFILE_MOE_ONE_STAGE("activation_weight");
if (ok && g_parallel_q8_pending) {
ok = ds4_gpu_parallel_q8_matvec_encode_pending(cb, midbuf) != 0;
}
id<MTLBuffer> down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer);
NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) :
(expertsbuf ? ds4_gpu_tensor_offset(experts) : 0);
if (ok && stream_expert_split_completed) {
/* The split path already wrote the resident partial output and
* accumulated the missing experts into out. */
} else if (ok && use_q4_grouped_experts) {
bool first_group = true;
for (uint32_t expert_base = 0; ok && expert_base < n_total_expert; expert_base += q4_expert_group_size) {
const uint32_t expert_count =
q4_expert_group_size < n_total_expert - expert_base ?
q4_expert_group_size : n_total_expert - expert_base;
if ((uint64_t)expert_base > UINT64_MAX / down_expert_bytes ||
(uint64_t)expert_count > UINT64_MAX / down_expert_bytes) {
ok = 0;
break;
}
const uint64_t group_rel = (uint64_t)expert_base * down_expert_bytes;
const uint64_t group_bytes = (uint64_t)expert_count * down_expert_bytes;
if (group_rel > UINT64_MAX - down_offset) {
ok = 0;
break;
}
uint64_t down_group_inner = 0;
id<MTLBuffer> down_group_buf =
q4_grouped_cache_views ?
ds4_gpu_wrap_model_exact_range(model_map,
model_size,
down_offset + group_rel,
group_bytes,
&down_group_inner) :
ds4_gpu_wrap_model_exact_range_transient(model_map,
model_size,
down_offset + group_rel,
group_bytes,
&down_group_inner);
if (!down_group_buf) {
ok = 0;
break;
}
ds4_gpu_moe_expert_group_args group_args = {
.expert_base = expert_base,
.expert_count = expert_count,
.accumulate = first_group ? 0u : 1u,
.pad0 = 0,
};
ok = ds4_gpu_encode_mul_mv_group_q4_sum6(cb,
g_moe_mul_mv_group_q4_k_sum6_pipeline,
&down_args,
&group_args,
down_group_buf,
(NSUInteger)down_group_inner,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
2);
first_group = false;
}
} else if (ok && use_q4_expert_address_table) {
ok = ds4_gpu_encode_mul_mv_addr_q4_sum6(cb,
g_moe_mul_mv_addr_q4_k_sum6_pipeline,
&down_args,
down_table,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
2);
} else if (ok && use_q4_expert_table) {
ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb,
g_moe_mul_mv_table_q4_k_sum6_pipeline,
&down_args,
down_table,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
2,
q4_table_queue_residency);
} else if (ok && use_q4_group6_experts) {
ok = ds4_gpu_encode_mul_mv_group6_sum6(cb,
g_moe_mul_mv_group6_q4_k_sum6_pipeline,
&down_args,
down_group6_bufs,
down_group6_offsets,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
2);
} else if (ok && use_q4_group8_experts) {
ok = ds4_gpu_encode_mul_mv_group8_sum6(cb,
g_moe_mul_mv_group8_q4_k_sum6_pipeline,
&down_args,
down_group8_bufs,
down_group8_offsets,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
2);
} else if (ok && use_q4_group24_experts) {
ok = ds4_gpu_encode_mul_mv_group24_sum6(cb,
g_moe_mul_mv_group24_q4_k_sum6_pipeline,
&down_args,
down_group24_bufs,
down_group24_offsets,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
2);
} else if (ok && (use_q4_gather_slots || use_selected_slots)) {
if (use_stream_expert_addr_table) {
if (down_type == DS4_METAL_TENSOR_IQ2_XXS) {
ok = ds4_gpu_encode_mul_mv_addr_iq2(cb,
g_moe_mul_mv_addr_iq2_xxs_pipeline,
&down_args,
stream_slot_entries,
n_expert,
stream_down_addr_buf,
midbuf,
ds4_gpu_tensor_offset(mid),
down_dst,
down_dst_off,
selected_exec_buf,
selected_exec_off,
down_smem,
2,
false, 2, nil);
} else if (use_stream_expert_masked_addr_table) {
ds4_gpu_stream_expert_split_args split_args = {
.active_mask = 0x3fu,
.accumulate = 0u,
};
ok = ds4_gpu_encode_mul_mv_addr_q2_sum6_masked(cb,
g_moe_mul_mv_addr_q2_k_sum6_masked_pipeline,
&down_args,
&split_args,
stream_slot_entries,
stream_down_addr_buf,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selected_exec_buf,
selected_exec_off,
down_smem,
2);
} else {
ok = ds4_gpu_encode_mul_mv_addr_q2_sum6(cb,
g_moe_mul_mv_addr_q2_k_sum6_pipeline,
&down_args,
stream_slot_entries,
n_expert,
stream_down_addr_buf,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selected_exec_buf,
selected_exec_off,
down_smem,
2,
nil);
}
} else {
ok = (!use_stream_expert_cache ||
ds4_gpu_stream_expert_cache_mark_entries_inflight(
stream_slot_entries,
n_expert,
0)) &&
ds4_gpu_encode_mul_mv_slots6_sum6(cb,
slots_sum6_pipeline,
&down_args,
down_slot_bufs,
down_slot_offsets,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
down_smem,
2);
}
} else if (ok && direct_down_sum) {
ok = ds4_gpu_encode_mul_mv_id_sum6(cb,
down_sum6_pipeline,
&down_args,
down_buf,
(NSUInteger)down_inner,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
add_in ? ds4_gpu_tensor_buffer(add_in) : nil,
add_in ? ds4_gpu_tensor_offset(add_in) : 0,
down_smem,
down_sum6_nsg);
} else if (ok) {
ok = ds4_gpu_encode_mul_mv_id(cb,
down_mv_pipeline,
&down_args,
down_buf,
(NSUInteger)down_inner,
midbuf,
ds4_gpu_tensor_offset(mid),
down_dst,
down_dst_off,
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
down_nsg,
down_rows_per_group_is_nr0);
}
DS4_METAL_PROFILE_MOE_ONE_STAGE("down");
if (ok && n_expert > 1 && !direct_down_sum && !stream_expert_split_completed) {
ok = ds4_gpu_encode_moe_sum_experts(cb,
down_dst,
down_dst_off,
outbuf,
ds4_gpu_tensor_offset(out),
out_dim,
n_expert,
n_tokens);
}
DS4_METAL_PROFILE_MOE_ONE_STAGE("sum");
if (ok && g_parallel_ffn_mode == 2 && g_parallel_ffn_stage == 1) {
ok = ds4_gpu_parallel_ffn_encode_second_stage(cb) != 0;
}
if (!ok) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34395); return 0; }
if (!ds4_gpu_finish_command_buffer(cb, owned, "routed tensor MoE")) { if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34397); return 0; }
if (q4_grouped_boundary || q4_exact_boundary || q4_table_boundary) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
if (getenv("DS4_GLM_TP_DEBUG")) fprintf(stderr, "ds4: routed_moe_one silent return at line %d\n", 34372);
return 0;
}
}
#undef DS4_METAL_PROFILE_MOE_ONE_STAGE
}
parallel_ffn_scope = NO;
return 1;
}
int ds4_gpu_routed_moe_batch_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *gate,
ds4_gpu_tensor *up,
ds4_gpu_tensor *mid,
ds4_gpu_tensor *experts,
const void *model_map,
uint64_t model_size,
uint64_t gate_offset,
uint64_t up_offset,
uint64_t down_offset,
uint32_t gate_type,
uint32_t down_type,
uint64_t gate_expert_bytes,
uint64_t gate_row_bytes,
uint64_t down_expert_bytes,
uint64_t down_row_bytes,
uint32_t expert_in_dim,
uint32_t expert_mid_dim,
uint32_t out_dim,
const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *weights,
uint32_t n_total_expert,
uint32_t n_expert,
float clamp,
const ds4_gpu_tensor *x,
uint32_t layer_index,
uint32_t n_tokens,
bool *mid_is_f16,
bool force_resident) {
if (!g_initialized && !ds4_gpu_init()) return 0;
/* TP sharding (see ds4_gpu_routed_moe_one_tensor): bind from the owned
* expert range and rebase ids in the kernels. */
uint32_t first_expert = 0;
uint32_t n_bind_expert = 0;
ds4_gpu_tp_expert_range(n_total_expert, &first_expert, &n_bind_expert);
const int32_t tp_expert_base_host = (int32_t)first_expert;
const int32_t tp_world_host = g_tp_full_expert_bind ? 1 : (int32_t)g_tp_split_world;
gate_offset += (uint64_t)first_expert * gate_expert_bytes;
up_offset += (uint64_t)first_expert * gate_expert_bytes;
down_offset += (uint64_t)first_expert * down_expert_bytes;
if (!out || !gate || !up || !mid || !x || !model_map || !selected || !weights ||
n_tokens == 0 || n_total_expert == 0 || n_expert == 0 ||
n_expert > DS4_METAL_MAX_ROUTED_EXPERT_USED) {
return 0;
}
if (gate_expert_bytes == 0 || down_expert_bytes == 0 ||
gate_row_bytes == 0 || down_row_bytes == 0) {
return 0;
}
if ((expert_in_dim % 256u) != 0 || (expert_mid_dim % 256u) != 0) return 0;
if ((uint64_t)n_total_expert > UINT64_MAX / gate_expert_bytes ||
(uint64_t)n_total_expert > UINT64_MAX / down_expert_bytes) {
fprintf(stderr, "ds4: Metal routed batch MoE tensor byte size overflow\n");
return 0;
}
const uint64_t gate_tensor_bytes = (uint64_t)n_bind_expert * gate_expert_bytes;
const uint64_t down_tensor_bytes = (uint64_t)n_bind_expert * down_expert_bytes;
/*
* PRO Q4 routed expert tensors are multi-GiB per layer. A one-token
* layer-slice prefill should use the one-token path so it can either bind
* exact tensor views with GPU-selected IDs or fall back to selected-expert
* views. Keep this guarded by tensor size so Flash and mixed-Flash Q4 keep
* their existing fast path.
*/
const uint64_t q4_selected_min_tensor_bytes = 2ull * 1024ull * 1024ull * 1024ull;
const bool can_single_token_q4_grouped =
getenv("DS4_METAL_ENABLE_Q4_GROUPED_EXPERTS") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GROUPED_EXPERTS") == NULL &&
g_moe_mul_mv_group_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_group_q4_k_sum6_pipeline != nil;
const bool can_single_token_q4_selected_slots =
ds4_gpu_q4_selected_paths_allowed() &&
getenv("DS4_METAL_DISABLE_Q4_SELECTED_EXPERT_VIEWS") == NULL &&
g_moe_mul_mv_slots6_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_slots6_q4_k_sum6_pipeline != nil;
const bool can_single_token_q4_group6 =
n_total_expert == 384 &&
getenv("DS4_METAL_ENABLE_Q4_GROUP6_EXPERT_TABLE") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GROUP6_EXPERT_TABLE") == NULL &&
g_moe_mul_mv_group6_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_group6_q4_k_sum6_pipeline != nil;
const bool can_single_token_q4_group8 =
n_total_expert == 384 &&
getenv("DS4_METAL_ENABLE_Q4_GROUP8_EXPERT_TABLE") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GROUP8_EXPERT_TABLE") == NULL &&
g_moe_mul_mv_group8_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_group8_q4_k_sum6_pipeline != nil;
const bool can_single_token_q4_group24 =
n_total_expert == 384 &&
getenv("DS4_METAL_ENABLE_Q4_GROUP24_EXPERT_TABLE") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_GROUP24_EXPERT_TABLE") == NULL &&
g_moe_mul_mv_group24_q4_k_id_pipeline != nil &&
g_moe_mul_mv_group24_q4_k_sum6_pipeline != nil;
const uint64_t max_buffer_len = g_device ? (uint64_t)[g_device maxBufferLength] : 0;
const bool can_single_token_q4_exact_tensor_id =
n_total_expert == 384 &&
max_buffer_len != 0 &&
gate_tensor_bytes <= max_buffer_len &&
down_tensor_bytes <= max_buffer_len &&
getenv("DS4_METAL_ENABLE_Q4_EXACT_TENSOR_ID") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_EXACT_TENSOR_ID") == NULL;
const bool enable_single_token_q4_expert_table =
getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL ||
ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert,
n_expert,
gate_tensor_bytes,
down_tensor_bytes);
const bool can_single_token_q4_expert_table =
n_total_expert == 384 &&
enable_single_token_q4_expert_table &&
getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL &&
g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_table_q4_k_sum6_pipeline != nil &&
g_moe_table_q4_pair_gate_encoder != nil &&
g_moe_table_q4_pair_up_encoder != nil &&
g_moe_table_q4_sum_down_encoder != nil;
const bool can_single_token_q4_expert_address_table =
n_total_expert == 384 &&
getenv("DS4_METAL_ENABLE_Q4_EXPERT_ADDRESS_TABLE") != NULL &&
getenv("DS4_METAL_DISABLE_Q4_EXPERT_ADDRESS_TABLE") == NULL &&
g_moe_mul_mv_addr_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_q4_k_sum6_pipeline != nil;
const bool use_single_token_q4_one_tensor =
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_tokens == 1 &&
n_expert == 6 &&
n_total_expert >= 128 &&
(g_ssd_streaming_mode ||
(gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes)) &&
!g_quality_mode &&
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
(can_single_token_q4_grouped ||
can_single_token_q4_group6 ||
can_single_token_q4_group8 ||
can_single_token_q4_group24 ||
can_single_token_q4_exact_tensor_id ||
can_single_token_q4_expert_address_table ||
can_single_token_q4_expert_table ||
can_single_token_q4_selected_slots);
const bool use_single_token_mxfp4_one_tensor =
gate_type == DS4_METAL_TENSOR_MXFP4 &&
down_type == DS4_METAL_TENSOR_MXFP4 &&
n_tokens == 1 &&
n_expert == 6 &&
n_total_expert >= 128 &&
!g_quality_mode &&
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_id_mxfp4_sum6_pipeline != nil &&
(!g_ssd_streaming_mode ||
(g_moe_mul_mv_slots6_mxfp4_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_slots6_mxfp4_sum6_pipeline != nil));
if (use_single_token_q4_one_tensor || use_single_token_mxfp4_one_tensor) {
if (mid_is_f16) *mid_is_f16 = false;
/* the one-token path binds the owned expert range itself */
return ds4_gpu_routed_moe_one_tensor(out,
gate,
up,
mid,
experts,
model_map,
model_size,
gate_offset - (uint64_t)first_expert * gate_expert_bytes,
up_offset - (uint64_t)first_expert * gate_expert_bytes,
down_offset - (uint64_t)first_expert * down_expert_bytes,
gate_type,
down_type,
gate_expert_bytes,
gate_row_bytes,
down_expert_bytes,
down_row_bytes,
expert_in_dim,
expert_mid_dim,
out_dim,
selected,
weights,
n_total_expert,
n_expert,
clamp,
x,
NULL,
layer_index,
force_resident);
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> gatebuf = ds4_gpu_tensor_buffer(gate);
id<MTLBuffer> upbuf = ds4_gpu_tensor_buffer(up);
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(mid);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> expertsbuf = ds4_gpu_tensor_buffer(experts);
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
const uint64_t x_bytes = (uint64_t)n_tokens * expert_in_dim * sizeof(float);
const uint64_t mid_bytes = (uint64_t)n_tokens * n_expert * expert_mid_dim * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_tokens * out_dim * sizeof(float);
const uint64_t selected_bytes = (uint64_t)n_tokens * n_expert * sizeof(int);
const uint64_t weights_bytes = (uint64_t)n_tokens * n_expert * sizeof(float);
if (!xbuf || !gatebuf || !upbuf || !midbuf || !outbuf || !selectedbuf || !weightsbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(gate) < mid_bytes ||
ds4_gpu_tensor_bytes(up) < mid_bytes ||
ds4_gpu_tensor_bytes(mid) < mid_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(selected) < selected_bytes ||
ds4_gpu_tensor_bytes(weights) < weights_bytes) {
fprintf(stderr, "ds4: Metal routed batch MoE received undersized activation buffers\n");
return 0;
}
uint64_t gate_inner = 0;
uint64_t up_inner = 0;
uint64_t down_inner = 0;
id<MTLBuffer> gate_buf = nil;
id<MTLBuffer> up_buf = nil;
id<MTLBuffer> down_buf = nil;
DS4MetalQ4ExpertTable *gate_table = nil;
DS4MetalQ4ExpertTable *up_table = nil;
DS4MetalQ4ExpertTable *down_table = nil;
id q4_table_layer_residency = nil;
id<MTLBuffer> stream_gate_addr_buf = nil;
id<MTLBuffer> stream_up_addr_buf = nil;
id<MTLBuffer> stream_down_addr_buf = nil;
id<MTLBuffer> stream_overflow_gate = nil;
id<MTLBuffer> stream_overflow_up = nil;
id<MTLBuffer> stream_overflow_down = nil;
ds4_gpu_stream_expert_cache_entry
*stream_resources[DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT] = { NULL };
uint32_t stream_resource_count = 0;
uint32_t stream_unique = 0;
const uint32_t pair_rows = n_tokens * n_expert;
const uint64_t down_scratch_bytes = (uint64_t)pair_rows * out_dim * sizeof(float);
const uint32_t gate_nr0 = ds4_gpu_routed_mv_nr0(gate_type);
const uint32_t down_nr0 = ds4_gpu_routed_mv_nr0(down_type);
id<MTLComputePipelineState> gate_mv_pipeline = ds4_gpu_routed_mv_pipeline(gate_type);
id<MTLComputePipelineState> down_mv_pipeline = ds4_gpu_routed_mv_pipeline(down_type);
id<MTLComputePipelineState> gate_mm_pipeline = nil;
id<MTLComputePipelineState> up_mm_pipeline = nil;
id<MTLComputePipelineState> down_mm_pipeline = nil;
id<MTLComputePipelineState> pair_swiglu_mm_pipeline = nil;
if (gate_nr0 == 0 || down_nr0 == 0 || !gate_mv_pipeline || !down_mv_pipeline) {
fprintf(stderr, "ds4: unsupported Metal routed batch MoE quant types gate=%u down=%u\n",
gate_type, down_type);
return 0;
}
const bool use_iq2_batch_selected_addr =
!force_resident &&
ds4_gpu_stream_prefill_batch_selected_addr_enabled(n_tokens,
n_total_expert,
n_expert,
gate_type,
down_type) &&
n_tokens > 1 &&
n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
!g_quality_mode &&
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_addr_q2_k_sum6_pipeline != nil;
/* Small full-GLM appends reuse cached bytes with the same MV/MM
* arithmetic. Large prefills retain sequential whole-layer reads. */
const bool use_iq2_cached_batch =
g_ssd_streaming_mode && !force_resident && g_tp_split_world == 1 &&
!ds4_gpu_glm_streaming_prefill_full_layer_active() &&
gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
down_type == DS4_METAL_TENSOR_IQ2_XXS && n_expert == 8 &&
n_tokens < 256 &&
n_total_expert <= DS4_METAL_STREAM_EXPERT_CACHE_MAX_EXPERT &&
ds4_gpu_stream_expert_cache_note_expert_size(gate_expert_bytes,
down_expert_bytes) &&
ds4_gpu_stream_expert_cache_configured_count() >= n_total_expert &&
getenv("DS4_METAL_GLM_DISABLE_STREAMING_EXPERT_CACHE") == NULL &&
getenv("DS4_METAL_DISABLE_IQ2_STREAM_ADDR_TABLE") == NULL &&
getenv("DS4_METAL_DISABLE_STREAMING_EXPERT_ADDR_TABLE") == NULL &&
getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL &&
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
getenv("DS4_METAL_DISABLE_STREAMING_PREFILL_BATCH_SELECTED_ADDR") == NULL;
const bool use_cached_batch = use_iq2_batch_selected_addr || use_iq2_cached_batch;
ds4_gpu_mul_mv_id_args gate_args =
ds4_gpu_make_mul_mv_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
gate_row_bytes, gate_expert_bytes,
1, n_expert, n_tokens, gate_nr0);
gate_args.tp_rank = g_tp_split_rank;
gate_args.tp_world = tp_world_host;
gate_args.tp_expert_base = tp_expert_base_host;
ds4_gpu_mul_mv_id_args down_args =
ds4_gpu_make_mul_mv_id_args(expert_mid_dim, out_dim, n_total_expert,
down_row_bytes, down_expert_bytes,
n_expert, n_expert, n_tokens, down_nr0);
down_args.tp_rank = g_tp_split_rank;
down_args.tp_world = tp_world_host;
down_args.tp_expert_base = tp_expert_base_host;
const bool q4_batch_expert_table_auto =
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
ds4_gpu_pro_q4_expert_table_auto_enabled(n_total_expert,
n_expert,
gate_tensor_bytes,
down_tensor_bytes);
const bool q4_batch_table_queue_residency =
ds4_gpu_q4_table_queue_residency_enabled(q4_batch_expert_table_auto);
const bool enable_q4_batch_expert_table =
getenv("DS4_METAL_ENABLE_Q4_BATCH_EXPERT_TABLE") != NULL ||
getenv("DS4_METAL_ENABLE_Q4_EXPERT_TABLE") != NULL ||
q4_batch_expert_table_auto;
const bool use_q4_batch_expert_table =
!g_ssd_streaming_mode &&
gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K &&
n_expert == 6 &&
n_tokens > 1 &&
n_total_expert == 384 &&
gate_tensor_bytes >= q4_selected_min_tensor_bytes &&
down_tensor_bytes >= q4_selected_min_tensor_bytes &&
!g_quality_mode &&
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
getenv("DS4_METAL_DISABLE_ROUTED_PAIR_SWIGLU_FUSION") == NULL &&
getenv("DS4_METAL_DISABLE_Q4_BATCH_EXPERT_TABLE") == NULL &&
getenv("DS4_METAL_DISABLE_Q4_EXPERT_TABLE") == NULL &&
g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline != nil &&
g_moe_mul_mv_table_q4_k_sum6_pipeline != nil &&
g_moe_table_q4_pair_gate_encoder != nil &&
g_moe_table_q4_pair_up_encoder != nil &&
g_moe_table_q4_sum_down_encoder != nil &&
enable_q4_batch_expert_table &&
(getenv("DS4_METAL_Q4_TABLE_USE_RESOURCES") != NULL ||
getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL ||
q4_batch_table_queue_residency ||
ds4_gpu_q4_table_model_residency_enabled());
const bool use_mm_id =
!use_q4_batch_expert_table &&
!use_iq2_batch_selected_addr &&
n_tokens >= 32u &&
ds4_gpu_mul_mm_id_map0_name(n_expert) != NULL;
bool mm_id_mpp_active = false;
/* Reuse half activations across gate/up and all output tiles. V4.1
* admits 512 MiB for 8k tiles; other shapes retain the 256 MiB limit. */
const uint64_t packed_rows = (uint64_t)n_tokens * n_expert;
const uint64_t packed_header = ((uint64_t)n_total_expert + 31u) / 32u * 32u * sizeof(uint32_t);
const uint64_t packed_width = expert_in_dim > expert_mid_dim ? expert_in_dim : expert_mid_dim;
const bool v41_packed = n_total_expert == 384u && expert_in_dim == 5120u &&
expert_mid_dim == 2304u;
const uint64_t packed_limit = (v41_packed ? UINT64_C(512) : UINT64_C(256)) << 20;
const bool use_packed_mpp = use_mm_id && !g_quality_mode &&
(!g_ssd_streaming_mode || force_resident) && n_tokens >= 512u &&
n_tokens <= (v41_packed ? 8192u : 4096u) &&
n_expert == 6u && n_total_expert <= 384u &&
expert_in_dim <= 5120u && expert_mid_dim <= 4096u &&
packed_header + packed_rows * packed_width * sizeof(uint16_t) <= packed_limit &&
ds4_gpu_routed_mm_mpp_mask() == 7 &&
((gate_type == DS4_METAL_TENSOR_IQ2_XXS && down_type == DS4_METAL_TENSOR_Q2_K) ||
(g_tp_split_world == 2 && gate_type == DS4_METAL_TENSOR_MXFP4 &&
down_type == DS4_METAL_TENSOR_MXFP4 &&
getenv("DS4_METAL_DISABLE_MXFP4_MM_ID_MPP") == NULL)) &&
getenv("DS4_METAL_DISABLE_ROUTED_MPP_PACKED") == NULL;
const bool packed_m32n128 = use_packed_mpp && n_tokens >= 8192u &&
n_total_expert == 384u && expert_in_dim == 5120u && expert_mid_dim == 2304u &&
gate_type == DS4_METAL_TENSOR_IQ2_XXS && down_type == DS4_METAL_TENSOR_Q2_K &&
!getenv("DS4_METAL_DISABLE_V41_PACKED_M32N128");
/* Tiny verification and session batches reuse the fused decode kernels.
* V4.1 needs scalar reduction order before its BF16/router boundaries
* through eight rows, for both routed recipes. Larger prefills retain
* the grouped matmul path. */
static int tiny_batch_off = -1;
if (tiny_batch_off < 0) tiny_batch_off = getenv("DS4_METAL_DISABLE_V41_MXFP4_TINY_BATCH") != NULL;
const bool v41_decode_batch = n_tokens <= 8u && n_total_expert == 384u &&
n_expert == 6u && expert_in_dim == 5120u && expert_mid_dim == 2304u &&
out_dim == 5120u &&
((gate_type == DS4_METAL_TENSOR_IQ2_XXS && down_type == DS4_METAL_TENSOR_Q2_K) ||
(gate_type == DS4_METAL_TENSOR_Q4_K && down_type == DS4_METAL_TENSOR_Q4_K) ||
(gate_type == DS4_METAL_TENSOR_MXFP4 && down_type == DS4_METAL_TENSOR_MXFP4 && !tiny_batch_off));
const bool use_tiny_pair_mv =
!g_quality_mode &&
(n_tokens <= 5u || v41_decode_batch ||
(n_tokens == 6u && ds4_gpu_device_is_m5_apple_silicon() &&
(gate_type == DS4_METAL_TENSOR_IQ2_XXS ||
gate_type == DS4_METAL_TENSOR_Q4_K ||
(g_tp_split_world == 2 && gate_type == DS4_METAL_TENSOR_MXFP4)))) &&
!use_q4_batch_expert_table &&
!use_mm_id &&
((gate_type == DS4_METAL_TENSOR_IQ2_XXS && g_moe_mul_mv_id_iq2_xxs_pair_pipeline) ||
(gate_type == DS4_METAL_TENSOR_Q4_K && g_moe_mul_mv_id_q4_k_pair_pipeline) ||
(gate_type == DS4_METAL_TENSOR_MXFP4 &&
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline &&
getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL &&
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL));
id<MTLComputePipelineState> tiny_pair_swiglu_pipeline = nil;
const bool use_tp_mxfp4_static_batch = use_tiny_pair_mv &&
ds4_gpu_device_is_m5_apple_silicon() && g_tp_split_world == 2 &&
gate_type == DS4_METAL_TENSOR_MXFP4 && down_type == DS4_METAL_TENSOR_MXFP4 &&
n_expert == 6 && n_total_expert == 256 &&
expert_in_dim == 4096 && expert_mid_dim == 2048 && out_dim == 4096 &&
gate_row_bytes == 2176 && gate_expert_bytes == 4456448 &&
down_row_bytes == 1088 && down_expert_bytes == 4456448 &&
!g_ssd_streaming_mode &&
getenv("DS4_METAL_DISABLE_M5_TP_MXFP4_STATIC") == NULL &&
g_moe_mul_mv_id_mxfp4_pair_swiglu_tp_static_pipeline_nsg1 != nil &&
g_moe_mul_mv_id_mxfp4_sum6_tp_full_rows_static_pipeline_nsg1 != nil;
if (gate_type == DS4_METAL_TENSOR_IQ2_XXS) {
tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_iq2_xxs_pair_swiglu_pipeline;
} else if (gate_type == DS4_METAL_TENSOR_Q4_K) {
tiny_pair_swiglu_pipeline = g_moe_mul_mv_id_q4_k_pair_swiglu_pipeline;
} else if (gate_type == DS4_METAL_TENSOR_MXFP4) {
tiny_pair_swiglu_pipeline = use_tp_mxfp4_static_batch ?
g_moe_mul_mv_id_mxfp4_pair_swiglu_tp_static_pipeline_nsg1 :
g_moe_mul_mv_id_mxfp4_pair_swiglu_pipeline;
}
const bool use_tiny_pair_swiglu =
use_tiny_pair_mv &&
tiny_pair_swiglu_pipeline != nil &&
getenv("DS4_METAL_DISABLE_TINY_PAIR_SWIGLU_FUSION") == NULL &&
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL;
ds4_gpu_mul_mm_id_map_args gate_map_args = { 0 };
ds4_gpu_mul_mm_id_args gate_mm_args = { 0 };
ds4_gpu_mul_mm_id_args down_mm_args = { 0 };
id<MTLComputePipelineState> map_pipeline = nil;
/*
* The grouped routed-MoE matmul loads activation tiles as half before
* using SIMD-group MMA. Store the SwiGLU/route-weight intermediate in
* that same precision so the down projection avoids a large F32 mid
* write/read. --quality keeps the older F32 intermediate.
*/
const bool request_mid_f16 =
!g_quality_mode &&
!use_q4_batch_expert_table &&
!use_iq2_batch_selected_addr;
/*
* Fused gate+up grouped matmul with the SwiGLU epilogue. Both IQ2 and
* Q4_K use the compact expert work list, the same MMA accumulation
* order, and the same epilogue math as separate GEMMs + SwiGLU, so the
* mid tensor is bit-identical.
*/
const bool use_mm_id_pair_swiglu =
use_mm_id &&
!(gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
(ds4_gpu_routed_mm_mpp_mask() & 3) == 3) &&
g_tp_split_world != 2 && /* pair-swiglu mm kernel lacks expert ownership */
request_mid_f16 &&
n_expert == 6 &&
((gate_type == DS4_METAL_TENSOR_IQ2_XXS &&
down_type == DS4_METAL_TENSOR_Q2_K) ||
(gate_type == DS4_METAL_TENSOR_Q4_K &&
down_type == DS4_METAL_TENSOR_Q4_K) ||
(gate_type == DS4_METAL_TENSOR_MXFP4 &&
down_type == DS4_METAL_TENSOR_MXFP4)) &&
getenv("DS4_METAL_DISABLE_MOE_MM_ID_PAIR_SWIGLU") == NULL &&
getenv("DS4_METAL_MOE_WRITE_CLAMPED_ACT") == NULL &&
getenv("DS4_METAL_GRAPH_DUMP_PREFIX") == NULL;
/*
* The MXFP4 32x32 specialization uses two SIMDgroups and 8 KiB of
* threadgroup memory, and exactly culls SIMDgroup 1 on at-most-16-row
* final work items. Balanced full-model A/B at both 2K and 8K makes
* it the resident pre-M5 default for large prefill.
*/
const bool resident_tile_types =
(gate_type == DS4_METAL_TENSOR_MXFP4 && down_type == DS4_METAL_TENSOR_MXFP4) ||
(v41_packed && /* Q4_K verified on the V4.1 shape only */
gate_type == DS4_METAL_TENSOR_Q4_K && down_type == DS4_METAL_TENSOR_Q4_K);
const bool use_pre_m5_mxfp4_mm_id_pair_swiglu_compact_tile_default =
ds4_gpu_device_is_pre_m5_apple_silicon() &&
!g_ssd_streaming_mode &&
n_tokens >= 2048u &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_PAIR_SWIGLU_COMPACT_TILE") == NULL;
const bool use_mxfp4_mm_id_pair_swiglu_compact_tile =
use_mm_id_pair_swiglu &&
resident_tile_types &&
g_tp_split_world == 1 &&
(use_pre_m5_mxfp4_mm_id_pair_swiglu_compact_tile_default ||
(g_test_flags & DS4_GPU_TEST_MXFP4_PAIR_COMPACT_TILE) != 0u);
/*
* Token-centric map construction for the compact resident pre-M5 MXFP4
* path. It emits the established work_count/uint2 trailer consumed
* by the existing padded direct launches and changes no arithmetic.
*/
const bool use_pre_m5_mxfp4_mm_id_map_scatter_default =
ds4_gpu_device_is_pre_m5_apple_silicon() &&
n_tokens >= 2048u &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_MAP_SCATTER") == NULL;
const bool use_mxfp4_mm_id_map_scatter =
use_mxfp4_mm_id_pair_swiglu_compact_tile &&
resident_tile_types &&
n_expert == 6 &&
!g_ssd_streaming_mode &&
g_tp_split_world == 1 &&
(use_pre_m5_mxfp4_mm_id_map_scatter_default ||
(g_test_flags & DS4_GPU_TEST_MXFP4_MAP_SCATTER) != 0u);
/*
* On a final occupied 32-row work tile with at most 16 rows, the
* second SIMDgroup pair has no valid MMA output. The specialized
* kernel still keeps every thread in staging, barriers, and the
* scalar SwiGLU epilogue. Keep the automatic path to the resident pre-M5
* shape and prefixes covered by the full-model A/B gate.
*/
const bool use_pre_m5_mxfp4_mm_id_pair_tail_simdgroup_cull_default =
ds4_gpu_device_is_pre_m5_apple_silicon() &&
!g_ssd_streaming_mode &&
n_tokens >= 2048u &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_PAIR_TAIL_SIMDGROUP_CULL") == NULL;
const bool use_mxfp4_mm_id_pair_tail_simdgroup_cull =
use_mm_id_pair_swiglu &&
resident_tile_types &&
g_tp_split_world == 1 &&
(use_pre_m5_mxfp4_mm_id_pair_tail_simdgroup_cull_default ||
(g_test_flags & DS4_GPU_TEST_MXFP4_PAIR_TAIL_CULL) != 0u);
const bool use_pre_m5_mxfp4_mm_id_down_tail_simdgroup_cull_default =
ds4_gpu_device_is_pre_m5_apple_silicon() &&
!g_ssd_streaming_mode &&
n_tokens >= 2048u &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_DOWN_TAIL_SIMDGROUP_CULL") == NULL;
const bool use_mxfp4_mm_id_down_tail_simdgroup_cull =
use_mm_id &&
request_mid_f16 &&
resident_tile_types &&
g_tp_split_world == 1 &&
(use_pre_m5_mxfp4_mm_id_down_tail_simdgroup_cull_default ||
(g_test_flags & DS4_GPU_TEST_MXFP4_DOWN_TAIL_CULL) != 0u);
/*
* Exact half-result dequantization table for the resident MXFP4 down
* projection. Balanced full-model A/B at 2K and 8K makes it the
* pre-M5 Apple-Silicon default for large prefill.
*/
const bool use_pre_m5_mxfp4_mm_id_down_half_lut_default =
ds4_gpu_device_is_pre_m5_apple_silicon() &&
n_tokens >= 2048u &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MOE_MM_ID_DOWN_HALF_LUT") == NULL;
const bool use_mxfp4_mm_id_down_half_lut =
use_mm_id &&
request_mid_f16 &&
gate_type == DS4_METAL_TENSOR_MXFP4 &&
down_type == DS4_METAL_TENSOR_MXFP4 &&
n_expert == 6 &&
!g_ssd_streaming_mode &&
g_tp_split_world == 1 &&
(use_pre_m5_mxfp4_mm_id_down_half_lut_default ||
(g_test_flags & DS4_GPU_TEST_MXFP4_DOWN_HALF_LUT) != 0u);
if (use_mm_id) {
gate_map_args =
ds4_gpu_make_mul_mm_id_map_args(expert_in_dim, n_total_expert, 1, n_expert, n_tokens);
gate_mm_args =
ds4_gpu_make_mul_mm_id_args(expert_in_dim, expert_mid_dim, n_total_expert,
gate_row_bytes, gate_expert_bytes,
1, n_expert, n_tokens);
down_mm_args =
ds4_gpu_make_mul_mm_id_args_src1_size(expert_mid_dim, out_dim, n_total_expert,
down_row_bytes, down_expert_bytes,
n_expert, n_expert, n_tokens,
request_mid_f16 ? sizeof(uint16_t) : sizeof(float));
gate_mm_args.tp_rank = g_tp_split_rank;
gate_mm_args.tp_world = tp_world_host;
gate_mm_args.tp_expert_base = tp_expert_base_host;
down_mm_args.tp_rank = g_tp_split_rank;
down_mm_args.tp_world = tp_world_host;
down_mm_args.tp_expert_base = tp_expert_base_host;
map_pipeline = ds4_gpu_get_pipeline(
use_mxfp4_mm_id_map_scatter ?
"kernel_mul_mm_id_map_scatter_work_ne20_6" :
ds4_gpu_mul_mm_id_map0_name(n_expert));
gate_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type);
up_mm_pipeline = ds4_gpu_routed_mm_pipeline(gate_type);
down_mm_pipeline = use_mxfp4_mm_id_down_half_lut ?
ds4_gpu_get_mul_mm_id_pipeline(
use_mxfp4_mm_id_down_tail_simdgroup_cull ?
"kernel_mul_mm_id_mxfp4_f16_half_lut_tail_cull" :
"kernel_mul_mm_id_mxfp4_f16_half_lut",
false) :
use_mxfp4_mm_id_down_tail_simdgroup_cull ?
ds4_gpu_get_mul_mm_id_pipeline(
gate_type == DS4_METAL_TENSOR_Q4_K ?
"kernel_mul_mm_id_q4_K_f16_tail_cull" :
"kernel_mul_mm_id_mxfp4_f16_tail_cull", false) :
request_mid_f16 ?
ds4_gpu_routed_mm_f16_rhs_pipeline(down_type) :
ds4_gpu_routed_mm_pipeline(down_type);
const int mpp_mask = ds4_gpu_routed_mm_mpp_mask();
if (mpp_mask && gate_type == DS4_METAL_TENSOR_IQ2_XXS) {
id<MTLComputePipelineState> mpp =
ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_iq2_xxs_f32_mpp", false);
if (mpp) {
if (mpp_mask & 1) gate_mm_pipeline = mpp;
if (mpp_mask & 2) up_mm_pipeline = mpp;
}
}
/* MXFP4 routed experts on the TensorOps pipeline (M5 nax units):
* same half weights x half activations with f32 accumulation as
* the simdgroup kernel, different accumulation order.
* DS4_METAL_DISABLE_MXFP4_MM_ID_MPP=1 keeps the simdgroup kernels. */
static int mxfp4_mpp = -1;
if (mxfp4_mpp < 0)
mxfp4_mpp = getenv("DS4_METAL_DISABLE_MXFP4_MM_ID_MPP") == NULL;
if (mpp_mask && mxfp4_mpp && gate_type == DS4_METAL_TENSOR_MXFP4) {
id<MTLComputePipelineState> mpp =
ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_mxfp4_f32_mpp", false);
if (mpp) {
if (mpp_mask & 1) gate_mm_pipeline = mpp;
if (mpp_mask & 2) up_mm_pipeline = mpp;
}
}
/* Q4_K routed gate/up on the TensorOps pipeline: same operand
* domains as the simdgroup kernel, different accumulation order.
* DS4_METAL_DISABLE_Q4_K_MM_ID_MPP=1 keeps the simdgroup kernels. */
static int q4_k_mpp = -1;
if (q4_k_mpp < 0)
q4_k_mpp = getenv("DS4_METAL_DISABLE_Q4_K_MM_ID_MPP") == NULL;
if (mpp_mask && q4_k_mpp && gate_type == DS4_METAL_TENSOR_Q4_K) {
id<MTLComputePipelineState> mpp =
ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f32_dbuf_mpp", false);
if (mpp) {
if (mpp_mask & 1) gate_mm_pipeline = mpp;
if (mpp_mask & 2) up_mm_pipeline = mpp;
mm_id_mpp_active = true;
}
}
if ((mpp_mask & 4) && mxfp4_mpp && request_mid_f16 &&
down_type == DS4_METAL_TENSOR_MXFP4) {
id<MTLComputePipelineState> mpp =
ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_mxfp4_f16_mpp", false);
if (mpp) {
down_mm_pipeline = mpp;
}
}
if ((mpp_mask & 4) && q4_k_mpp && request_mid_f16 &&
down_type == DS4_METAL_TENSOR_Q4_K) {
id<MTLComputePipelineState> mpp =
ds4_gpu_get_mul_mm_id_pipeline("kernel_mul_mm_id_q4_K_f16_dbuf_mpp", false);
if (mpp) {
down_mm_pipeline = mpp;
mm_id_mpp_active = true;
}
}
if ((mpp_mask & 4) && request_mid_f16 &&
(down_type == DS4_METAL_TENSOR_Q2_K || down_type == DS4_METAL_TENSOR_IQ2_XXS)) {
id<MTLComputePipelineState> mpp = ds4_gpu_get_mul_mm_id_pipeline(
down_type == DS4_METAL_TENSOR_Q2_K ?
"kernel_mul_mm_id_q2_K_f16_mpp" :
"kernel_mul_mm_id_iq2_xxs_f16_mpp", false);
if (mpp) {
down_mm_pipeline = mpp;
}
}
if (use_packed_mpp) {
gate_mm_pipeline = up_mm_pipeline =
ds4_gpu_get_pipeline(gate_type == DS4_METAL_TENSOR_MXFP4 ?
"kernel_mul_mm_id_mxfp4_mpp_packed" : packed_m32n128 ?
"kernel_mul_mm_id_iq2_xxs_mpp_packed_m32n128" : "kernel_mul_mm_id_iq2_xxs_mpp_packed");
down_mm_pipeline = ds4_gpu_get_pipeline(down_type == DS4_METAL_TENSOR_MXFP4 ?
"kernel_mul_mm_id_mxfp4_mpp_packed" : packed_m32n128 ?
"kernel_mul_mm_id_q2_K_mpp_packed_m32n128" : "kernel_mul_mm_id_q2_K_mpp_packed");
}
if (use_packed_mpp) mm_id_mpp_active = false;
if (use_mm_id_pair_swiglu) {
/* Exact half-domain block scaling for the resident MXFP4 pair
* tile: E8M0 is a power of two and every E2M1 magnitude is
* exact in binary16, so inside the guarded exponent band the
* half product is the identical once-rounded value while the
* per-element f32->f16 narrowing and the 64-byte float table
* disappear. Outside the band the kernel keeps the established
* float arithmetic. */
const bool use_mxfp4_mm_id_pair_half_scale =
gate_type == DS4_METAL_TENSOR_MXFP4 &&
ds4_gpu_device_is_pre_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_PRE_M5_MXFP4_MM_ID_PAIR_HALF_SCALE") == NULL &&
!g_quality_mode &&
!g_ssd_streaming_mode &&
g_tp_split_world == 1;
pair_swiglu_mm_pipeline =
ds4_gpu_get_pipeline(
gate_type == DS4_METAL_TENSOR_Q4_K ?
(use_mxfp4_mm_id_pair_swiglu_compact_tile ?
"kernel_mul_mm_id_q4_K_pair_swiglu_f16_compact_tail_cull" :
use_mxfp4_mm_id_pair_tail_simdgroup_cull ?
"kernel_mul_mm_id_q4_K_pair_swiglu_f16_tail_cull" :
"kernel_mul_mm_id_q4_K_pair_swiglu_f16") :
gate_type == DS4_METAL_TENSOR_MXFP4 ?
(use_mxfp4_mm_id_pair_swiglu_compact_tile ?
(use_mxfp4_mm_id_pair_half_scale ?
"kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull_half_scale" :
"kernel_mul_mm_id_mxfp4_pair_swiglu_f16_compact_tail_cull") :
use_mxfp4_mm_id_pair_tail_simdgroup_cull ?
(use_mxfp4_mm_id_pair_half_scale ?
"kernel_mul_mm_id_mxfp4_pair_swiglu_f16_tail_cull_half_scale" :
"kernel_mul_mm_id_mxfp4_pair_swiglu_f16_tail_cull") :
(use_mxfp4_mm_id_pair_half_scale ?
"kernel_mul_mm_id_mxfp4_pair_swiglu_f16_half_scale" :
"kernel_mul_mm_id_mxfp4_pair_swiglu_f16")) :
"kernel_mul_mm_id_iq2_xxs_pair_swiglu_f16");
}
if (!map_pipeline || !gate_mm_pipeline || !up_mm_pipeline || !down_mm_pipeline ||
(use_mm_id_pair_swiglu && !pair_swiglu_mm_pipeline)) {
return 0;
}
}
if (use_iq2_cached_batch && use_mm_id) {
const int mpp = ds4_gpu_routed_mm_mpp_mask();
gate_mm_pipeline = up_mm_pipeline = ds4_gpu_get_mul_mm_id_pipeline(
mpp ? "kernel_mul_mm_id_iq2_xxs_cached_f32_mpp" :
"kernel_mul_mm_id_iq2_xxs_cached_f32", false);
down_mm_pipeline = ds4_gpu_get_mul_mm_id_pipeline(
request_mid_f16 ? (mpp ? "kernel_mul_mm_id_iq2_xxs_cached_f16_mpp" :
"kernel_mul_mm_id_iq2_xxs_cached_f16") :
"kernel_mul_mm_id_iq2_xxs_cached_f32", false);
if (!gate_mm_pipeline || !down_mm_pipeline) return 0;
}
if (use_cached_batch) {
const int had_batch = g_batch_cb != nil;
if (had_batch && ds4_gpu_end_commands() == 0) {
return 0;
}
g_stream_prefill_batch_selected_addr_building++;
if (!ds4_gpu_stream_expert_cache_prepare_selected_batch(
model_map,
model_size,
layer_index,
selected,
n_tokens,
n_total_expert,
n_expert,
gate_offset,
up_offset,
down_offset,
gate_expert_bytes,
down_expert_bytes,
&stream_gate_addr_buf,
&stream_up_addr_buf,
&stream_down_addr_buf,
stream_resources,
&stream_resource_count,
&stream_unique,
&stream_overflow_gate,
&stream_overflow_up,
&stream_overflow_down)) {
g_stream_prefill_batch_selected_addr_building--;
return 0;
}
g_stream_prefill_batch_selected_addr_building--;
if (stream_unique == 0) {
ds4_gpu_stream_expert_cache_clear_layer(layer_index);
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr layer=%u "
"produced no resident experts\n",
layer_index);
return 0;
}
for (uint32_t i = 0; i < stream_resource_count; i++) {
ds4_gpu_stream_expert_cache_entry *entry = stream_resources[i];
if (!entry ||
!entry->valid ||
!entry->gate_buffer ||
!entry->up_buffer ||
!entry->down_buffer) {
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr layer=%u "
"lost resident expert resource %u/%u during preparation\n",
layer_index,
i,
stream_resource_count);
ds4_gpu_stream_expert_cache_clear_layer(layer_index);
return 0;
}
}
if (had_batch && ds4_gpu_begin_commands() == 0) {
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr layer=%u "
"failed to reopen command batch after preparation\n",
layer_index);
ds4_gpu_stream_expert_cache_clear_layer(layer_index);
return 0;
}
}
if (use_q4_batch_expert_table) {
gate_table = ds4_gpu_q4_expert_table(model_map,
model_size,
gate_offset,
gate_expert_bytes,
n_total_expert,
g_moe_table_q4_pair_gate_encoder);
up_table = ds4_gpu_q4_expert_table(model_map,
model_size,
up_offset,
gate_expert_bytes,
n_total_expert,
g_moe_table_q4_pair_up_encoder);
down_table = ds4_gpu_q4_expert_table(model_map,
model_size,
down_offset,
down_expert_bytes,
n_total_expert,
g_moe_table_q4_sum_down_encoder);
if (!gate_table || !up_table || !down_table) {
return 0;
}
q4_table_layer_residency =
ds4_gpu_q4_expert_layer_residency_set(gate_table,
up_table,
down_table,
q4_batch_expert_table_auto);
if ((q4_batch_table_queue_residency ||
getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL) &&
!q4_table_layer_residency &&
!ds4_gpu_q4_table_model_residency_enabled()) {
fprintf(stderr, "ds4: Metal Q4 batch expert table residency set is not available\n");
return 0;
}
} else if (use_cached_batch) {
if (n_expert > 1 && (!expertsbuf ||
ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) {
if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
&g_moe_down_scratch_bytes,
(NSUInteger)down_scratch_bytes,
"ds4_moe_down_scratch")) {
ds4_gpu_stream_expert_cache_clear_layer(layer_index);
return 0;
}
}
} else {
if (n_expert > 1 && (!expertsbuf ||
ds4_gpu_tensor_bytes(experts) < down_scratch_bytes)) {
if (!ds4_gpu_ensure_scratch_buffer(&g_moe_down_scratch_buffer,
&g_moe_down_scratch_bytes,
(NSUInteger)down_scratch_bytes,
"ds4_moe_down_scratch")) {
return 0;
}
}
gate_buf = ds4_gpu_wrap_model_range(model_map,
model_size,
gate_offset,
gate_tensor_bytes,
&gate_inner);
up_buf = ds4_gpu_wrap_model_range(model_map,
model_size,
up_offset,
gate_tensor_bytes,
&up_inner);
down_buf = ds4_gpu_wrap_model_range(model_map,
model_size,
down_offset,
down_tensor_bytes,
&down_inner);
if (!gate_buf || !up_buf || !down_buf) return 0;
if (getenv("DS4_GLM_TP_DEBUG") && layer_index == 3) {
fprintf(stderr,
"ds4: batch mv binds l=%u rank=%d base=%d gate=%llu+%llu "
"up=%llu down=%llu inner=%llu/%llu/%llu ne02=%d nei0=%d nr0=%d\n",
layer_index, g_tp_split_rank, tp_expert_base_host,
(unsigned long long)gate_offset,
(unsigned long long)gate_tensor_bytes,
(unsigned long long)up_offset,
(unsigned long long)down_offset,
(unsigned long long)gate_inner,
(unsigned long long)up_inner,
(unsigned long long)down_inner,
gate_args.ne02, gate_args.nei0, gate_args.nr0);
}
}
const bool q4_batch_table_boundary =
use_q4_batch_expert_table &&
g_batch_cb != nil &&
getenv("DS4_METAL_Q4_TABLE_RESIDENCY_SET") != NULL &&
!q4_batch_table_queue_residency &&
getenv("DS4_METAL_DISABLE_Q4_TABLE_BOUNDARY") == NULL;
if (q4_batch_table_boundary) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
if (use_q4_batch_expert_table && !ds4_gpu_use_model_residency_set(cb)) {
return 0;
}
if (!q4_batch_table_queue_residency &&
q4_table_layer_residency &&
[cb respondsToSelector:@selector(useResidencySet:)]) {
[cb useResidencySet:q4_table_layer_residency];
}
const bool moe_stage_profile =
g_batch_cb != nil &&
ds4_gpu_stage_profile_enabled_for_layer("DS4_METAL_MOE_STAGE_PROFILE",
"DS4_METAL_MOE_STAGE_PROFILE_LAYER",
layer_index);
const char *moe_stage_filter = getenv("DS4_METAL_MOE_STAGE_PROFILE_FILTER");
const char *moe_path =
use_q4_batch_expert_table ? "q4_table_pair_swiglu" :
use_iq2_batch_selected_addr ? "iq2_batch_stream_addr" :
use_iq2_cached_batch ? (use_mm_id ? "iq2_cached_mm" : "iq2_cached_mv") :
use_mm_id_pair_swiglu ? "mm_id_pair_swiglu" :
use_mm_id ? "mm_id" :
use_tiny_pair_swiglu ? "tiny_pair_swiglu" :
(use_tiny_pair_mv ? "tiny_pair_mv" : "mv");
double moe_stage_t0 = moe_stage_profile ? ds4_gpu_now_ms() : 0.0;
if (moe_stage_profile) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
moe_stage_t0 = ds4_gpu_now_ms();
}
if (use_iq2_cached_batch && stream_resource_count &&
!ds4_gpu_stream_expert_cache_mark_entries_inflight(
stream_resources, stream_resource_count, 0)) return 0;
#define DS4_METAL_PROFILE_MOE_STAGE(name) do { \
if (ok && moe_stage_profile) { \
if (ds4_gpu_end_commands() == 0) { \
ok = 0; \
} else { \
const char *stage_name = (name); \
const double now_ms = ds4_gpu_now_ms(); \
const int print_stage = \
!moe_stage_filter || !moe_stage_filter[0] || \
strstr(stage_name, moe_stage_filter) != NULL; \
if (print_stage) { \
fprintf(stderr, \
"ds4: Metal routed MoE stage layer=%u tokens=%u pairs=%u experts=%u " \
"gate=%s down=%s path=%s mid=%s %s=%.3f ms\n", \
layer_index, n_tokens, pair_rows, n_expert, \
ds4_gpu_metal_tensor_type_name(gate_type), \
ds4_gpu_metal_tensor_type_name(down_type), \
moe_path, \
request_mid_f16 ? "f16" : "f32", \
stage_name, now_ms - moe_stage_t0); \
} \
moe_stage_t0 = now_ms; \
if (ds4_gpu_begin_commands() == 0) { \
ok = 0; \
} else { \
cb = ds4_gpu_command_buffer(&owned); \
if (!cb) ok = 0; \
} \
} \
} \
} while (0)
const NSUInteger gate_smem = ds4_gpu_routed_mv_smem(gate_type);
const NSUInteger down_smem = ds4_gpu_routed_mv_smem(down_type);
const NSUInteger gate_nsg = ds4_gpu_routed_mv_nsg(gate_type);
const NSUInteger down_nsg = ds4_gpu_routed_mv_nsg(down_type);
const bool gate_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(gate_type);
const bool down_rows_per_group_is_nr0 = ds4_gpu_routed_mv_rows_per_group_is_nr0(down_type);
id<MTLComputePipelineState> down_sum6_pipeline = nil;
if (down_type == DS4_METAL_TENSOR_Q2_K) {
down_sum6_pipeline = g_moe_mul_mv_id_q2_k_sum6_pipeline;
} else if (down_type == DS4_METAL_TENSOR_Q4_K) {
down_sum6_pipeline = g_moe_mul_mv_id_q4_k_sum6_pipeline;
} else if (down_type == DS4_METAL_TENSOR_MXFP4) {
down_sum6_pipeline = use_tp_mxfp4_static_batch ?
g_moe_mul_mv_id_mxfp4_sum6_tp_full_rows_static_pipeline_nsg1 :
g_moe_mul_mv_id_mxfp4_sum6_pipeline;
}
const bool direct_down_sum =
!g_quality_mode &&
!use_q4_batch_expert_table &&
!use_mm_id &&
n_expert == 6 &&
(n_tokens <= 4u || v41_decode_batch || use_tp_mxfp4_static_batch) &&
down_sum6_pipeline != nil;
int ok = 0;
if (use_iq2_batch_selected_addr) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(
cb,
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline,
&gate_args,
&act_args,
stream_resources,
stream_resource_count,
stream_gate_addr_buf,
stream_up_addr_buf,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false,
stream_overflow_gate,
stream_overflow_up);
if (!ok) {
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr layer=%u "
"failed to encode gate/up path tokens=%u unique=%u\n",
layer_index,
n_tokens,
stream_unique);
}
} else if (use_q4_batch_expert_table) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mv_table_q4_pair_swiglu(cb,
g_moe_mul_mv_table_q4_k_pair_swiglu_pipeline,
&gate_args,
&act_args,
gate_table,
up_table,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
2,
false,
q4_batch_table_queue_residency);
} else if (use_mm_id) {
/*
* The routed pair ids are the same for gate, up, and down. Build
* the expert-major work map once, then reuse it for all three
* batched expert matmuls.
*/
ok = ds4_gpu_encode_mul_mm_id_map(cb,
map_pipeline,
&gate_map_args,
&gate_mm_args,
selectedbuf,
ds4_gpu_tensor_offset(selected));
if (ok && use_packed_mpp)
ok = ds4_gpu_encode_moe_packed_rhs(cb, &gate_mm_args,
xbuf, ds4_gpu_tensor_offset(x), false, packed_limit);
DS4_METAL_PROFILE_MOE_STAGE("map");
if (ok && use_mm_id_pair_swiglu) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(uint16_t),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mm_id_iq2_pair_swiglu_f16(cb,
pair_swiglu_mm_pipeline,
use_mxfp4_mm_id_pair_swiglu_compact_tile,
&gate_mm_args,
&act_args,
gate_buf,
(NSUInteger)gate_inner,
up_buf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
midbuf,
ds4_gpu_tensor_offset(mid),
weightsbuf,
ds4_gpu_tensor_offset(weights));
DS4_METAL_PROFILE_MOE_STAGE("gate_up_fused");
} else if (ok) {
ok = ds4_gpu_encode_mul_mm_id_mapped_tile_resources(cb,
gate_mm_pipeline,
&gate_mm_args,
use_iq2_cached_batch ? stream_gate_addr_buf : gate_buf,
(NSUInteger)gate_inner,
use_packed_mpp ? g_moe_packed_rhs_buffer : xbuf,
use_packed_mpp ? 0 : ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
packed_m32n128 ? 16384u : mm_id_mpp_active ? 12288u : 8192u,
use_iq2_cached_batch ? stream_resources : NULL,
stream_resource_count, 0, stream_overflow_gate,
packed_m32n128 ? 32u : 64u);
DS4_METAL_PROFILE_MOE_STAGE("gate");
}
if (ok && !use_mm_id_pair_swiglu) {
ok = ds4_gpu_encode_mul_mm_id_mapped_tile_resources(cb,
up_mm_pipeline,
&gate_mm_args,
use_iq2_cached_batch ? stream_up_addr_buf : up_buf,
(NSUInteger)up_inner,
use_packed_mpp ? g_moe_packed_rhs_buffer : xbuf,
use_packed_mpp ? 0 : ds4_gpu_tensor_offset(x),
upbuf,
ds4_gpu_tensor_offset(up),
packed_m32n128 ? 16384u : mm_id_mpp_active ? 12288u : 8192u,
use_iq2_cached_batch ? stream_resources : NULL,
stream_resource_count, 1, stream_overflow_up,
packed_m32n128 ? 32u : 64u);
DS4_METAL_PROFILE_MOE_STAGE("up");
}
} else if (use_iq2_cached_batch && use_tiny_pair_swiglu) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim, .rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float), .clamp_value = clamp,
};
ok = ds4_gpu_encode_mul_mv_addr_iq2_pair_swiglu(cb,
g_moe_mul_mv_addr_iq2_xxs_pair_swiglu_pipeline, &gate_args, &act_args,
stream_resources, stream_resource_count, stream_gate_addr_buf, stream_up_addr_buf,
xbuf, ds4_gpu_tensor_offset(x), gatebuf, ds4_gpu_tensor_offset(gate),
upbuf, ds4_gpu_tensor_offset(up), midbuf, ds4_gpu_tensor_offset(mid),
selectedbuf, ds4_gpu_tensor_offset(selected), weightsbuf, ds4_gpu_tensor_offset(weights),
gate_smem, 2, false, stream_overflow_gate, stream_overflow_up);
} else if (use_iq2_cached_batch) {
ok = ds4_gpu_encode_mul_mv_addr_iq2(cb,
g_moe_mul_mv_addr_iq2_xxs_pipeline, &gate_args,
stream_resources, stream_resource_count, stream_gate_addr_buf,
xbuf, ds4_gpu_tensor_offset(x), gatebuf, ds4_gpu_tensor_offset(gate),
selectedbuf, ds4_gpu_tensor_offset(selected), gate_smem, gate_nsg,
gate_rows_per_group_is_nr0, 0, stream_overflow_gate) &&
ds4_gpu_encode_mul_mv_addr_iq2(cb,
g_moe_mul_mv_addr_iq2_xxs_pipeline, &gate_args,
stream_resources, stream_resource_count, stream_up_addr_buf,
xbuf, ds4_gpu_tensor_offset(x), upbuf, ds4_gpu_tensor_offset(up),
selectedbuf, ds4_gpu_tensor_offset(selected), gate_smem, gate_nsg,
gate_rows_per_group_is_nr0, 1, stream_overflow_up);
} else if (use_tiny_pair_swiglu) {
ds4_gpu_dsv4_moe_swiglu_weight_args act_args = {
.width = expert_mid_dim,
.rows = pair_rows,
.gate_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.up_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.mid_row_stride = (uint64_t)expert_mid_dim * sizeof(float),
.weight_stride = sizeof(float),
.write_clamped = 0,
.clamp_value = clamp,
};
/* a few single-node rows: every distinct expert read once */
static int dedup_off = -1;
if (dedup_off < 0) dedup_off = getenv("DS4_METAL_DISABLE_V41_MOE_DEDUP") != NULL;
const bool dedup = gate_type == DS4_METAL_TENSOR_MXFP4 && n_tokens >= 2u && n_tokens <= 8u &&
!use_tp_mxfp4_static_batch && g_tp_split_world == 1 && !dedup_off;
ok = dedup ? ds4_gpu_encode_mul_mv_id_pair_swiglu_dedup(cb, &gate_args, &act_args,
gate_buf, (NSUInteger)gate_inner, up_buf, (NSUInteger)up_inner,
xbuf, ds4_gpu_tensor_offset(x), gatebuf, ds4_gpu_tensor_offset(gate),
upbuf, ds4_gpu_tensor_offset(up), midbuf, ds4_gpu_tensor_offset(mid),
selectedbuf, ds4_gpu_tensor_offset(selected), weightsbuf, ds4_gpu_tensor_offset(weights),
gate_smem, n_tokens) :
ds4_gpu_encode_mul_mv_id_pair_swiglu(cb,
tiny_pair_swiglu_pipeline,
&gate_args,
&act_args,
gate_buf,
(NSUInteger)gate_inner,
up_buf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
selectedbuf,
ds4_gpu_tensor_offset(selected),
weightsbuf,
ds4_gpu_tensor_offset(weights),
gate_smem,
use_tp_mxfp4_static_batch ? 1 : 2,
false);
} else if (use_tiny_pair_mv) {
id<MTLComputePipelineState> pair_pipeline =
gate_type == DS4_METAL_TENSOR_IQ2_XXS ?
g_moe_mul_mv_id_iq2_xxs_pair_pipeline :
g_moe_mul_mv_id_q4_k_pair_pipeline;
ok = ds4_gpu_encode_mul_mv_id_pair(cb,
pair_pipeline,
&gate_args,
gate_buf,
(NSUInteger)gate_inner,
up_buf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
2,
false);
} else {
ok = ds4_gpu_encode_mul_mv_id(cb,
gate_mv_pipeline,
&gate_args,
gate_buf,
(NSUInteger)gate_inner,
xbuf,
ds4_gpu_tensor_offset(x),
gatebuf,
ds4_gpu_tensor_offset(gate),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
gate_nsg,
gate_rows_per_group_is_nr0) &&
ds4_gpu_encode_mul_mv_id(cb,
gate_mv_pipeline,
&gate_args,
up_buf,
(NSUInteger)up_inner,
xbuf,
ds4_gpu_tensor_offset(x),
upbuf,
ds4_gpu_tensor_offset(up),
selectedbuf,
ds4_gpu_tensor_offset(selected),
gate_smem,
gate_nsg,
gate_rows_per_group_is_nr0);
}
DS4_METAL_PROFILE_MOE_STAGE("gate_up");
const bool use_fused_activation = !g_quality_mode && !use_q4_batch_expert_table;
const bool use_mid_f16 =
use_mm_id &&
use_fused_activation &&
request_mid_f16;
if (mid_is_f16) *mid_is_f16 = use_mid_f16;
if (ok && use_iq2_batch_selected_addr) {
/* The address-table pair kernel already wrote weighted SwiGLU rows into mid. */
} else if (ok && use_q4_batch_expert_table) {
/* The table pair kernel already wrote weighted SwiGLU rows into mid. */
} else if (ok && use_mm_id_pair_swiglu) {
/* The fused batch mm_id pair kernel already wrote weighted f16 SwiGLU rows into mid. */
} else if (ok && use_tiny_pair_swiglu) {
/* The fused tiny pair kernel already wrote weighted F32 SwiGLU rows into mid. */
} else if (ok && use_fused_activation) {
ok = ds4_gpu_encode_moe_swiglu_weight(cb,
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
weightsbuf,
ds4_gpu_tensor_offset(weights),
expert_mid_dim,
pair_rows,
clamp,
use_mid_f16);
} else if (ok && clamp > 1.0e-6f) {
ok = ds4_gpu_encode_unary_f32_rows(cb,
g_unary_clamp_pipeline,
gatebuf,
ds4_gpu_tensor_offset(gate),
gatebuf,
ds4_gpu_tensor_offset(gate),
expert_mid_dim,
pair_rows,
0,
-FLT_MAX,
clamp);
if (ok) {
ok = ds4_gpu_encode_unary_f32_rows(cb,
g_unary_silu_pipeline,
gatebuf,
ds4_gpu_tensor_offset(gate),
midbuf,
ds4_gpu_tensor_offset(mid),
expert_mid_dim,
pair_rows,
1,
0.0f,
0.0f);
}
if (ok) {
ok = ds4_gpu_encode_unary_f32_rows(cb,
g_unary_clamp_pipeline,
upbuf,
ds4_gpu_tensor_offset(up),
upbuf,
ds4_gpu_tensor_offset(up),
expert_mid_dim,
pair_rows,
0,
-clamp,
clamp);
}
if (ok) {
ds4_gpu_bin_args mul_args =
ds4_gpu_make_bin_same_rows_args(expert_mid_dim, pair_rows);
ok = ds4_gpu_encode_bin_f32_rows(cb,
g_mul_pipeline,
&mul_args,
midbuf,
ds4_gpu_tensor_offset(mid),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid));
}
} else if (ok) {
ok = ds4_gpu_encode_swiglu_flat(cb,
gatebuf,
ds4_gpu_tensor_offset(gate),
upbuf,
ds4_gpu_tensor_offset(up),
midbuf,
ds4_gpu_tensor_offset(mid),
(uint32_t)((uint64_t)pair_rows * expert_mid_dim));
}
if (ok && !use_fused_activation && !use_q4_batch_expert_table) {
ds4_gpu_bin_args weight_args =
ds4_gpu_make_bin_rowwise_scalar_args(expert_mid_dim, pair_rows);
ok = ds4_gpu_encode_bin_f32_rows(cb,
g_bin_mul_scalar_pipeline,
&weight_args,
midbuf,
ds4_gpu_tensor_offset(mid),
weightsbuf,
ds4_gpu_tensor_offset(weights),
midbuf,
ds4_gpu_tensor_offset(mid));
}
DS4_METAL_PROFILE_MOE_STAGE("activation_weight");
id<MTLBuffer> down_dst = n_expert == 1 ? outbuf : (expertsbuf ? expertsbuf : g_moe_down_scratch_buffer);
NSUInteger down_dst_off = n_expert == 1 ? ds4_gpu_tensor_offset(out) :
(expertsbuf ? ds4_gpu_tensor_offset(experts) : 0);
if (ok) {
if (use_iq2_batch_selected_addr) {
ok = ds4_gpu_encode_mul_mv_addr_q2_sum6(
cb,
g_moe_mul_mv_addr_q2_k_sum6_pipeline,
&down_args,
stream_resources,
stream_resource_count,
stream_down_addr_buf,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
2,
stream_overflow_down);
if (!ok) {
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr layer=%u "
"failed to encode down path tokens=%u unique=%u\n",
layer_index,
n_tokens,
stream_unique);
}
} else if (use_q4_batch_expert_table) {
ok = ds4_gpu_encode_mul_mv_table_q4_sum6(cb,
g_moe_mul_mv_table_q4_k_sum6_pipeline,
&down_args,
down_table,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
2,
q4_batch_table_queue_residency);
} else if (direct_down_sum) {
ok = ds4_gpu_encode_mul_mv_id_sum6(cb,
down_sum6_pipeline,
&down_args,
down_buf,
(NSUInteger)down_inner,
midbuf,
ds4_gpu_tensor_offset(mid),
outbuf,
ds4_gpu_tensor_offset(out),
selectedbuf,
ds4_gpu_tensor_offset(selected),
nil,
0,
down_smem,
use_tp_mxfp4_static_batch ? 1 : 2);
} else if (use_mm_id) {
if (use_packed_mpp)
ok = ds4_gpu_encode_moe_packed_rhs(cb, &down_mm_args,
midbuf, ds4_gpu_tensor_offset(mid), true, packed_limit);
if (ok) ok = ds4_gpu_encode_mul_mm_id_mapped_tile_resources(cb,
down_mm_pipeline,
&down_mm_args,
use_iq2_cached_batch ? stream_down_addr_buf : down_buf,
(NSUInteger)down_inner,
use_packed_mpp ? g_moe_packed_rhs_buffer : midbuf,
use_packed_mpp ? 0 : ds4_gpu_tensor_offset(mid),
down_dst,
down_dst_off,
packed_m32n128 ? 16384u : mm_id_mpp_active ? 12288u : 8192u,
use_iq2_cached_batch ? stream_resources : NULL,
stream_resource_count, 2, stream_overflow_down,
packed_m32n128 ? 32u : 64u);
} else if (use_iq2_cached_batch) {
ok = ds4_gpu_encode_mul_mv_addr_iq2(cb,
g_moe_mul_mv_addr_iq2_xxs_pipeline, &down_args,
stream_resources, stream_resource_count, stream_down_addr_buf,
midbuf, ds4_gpu_tensor_offset(mid), down_dst, down_dst_off,
selectedbuf, ds4_gpu_tensor_offset(selected), down_smem, down_nsg,
down_rows_per_group_is_nr0, 2, stream_overflow_down);
} else {
ok = ds4_gpu_encode_mul_mv_id(cb,
down_mv_pipeline,
&down_args,
down_buf,
(NSUInteger)down_inner,
midbuf,
ds4_gpu_tensor_offset(mid),
down_dst,
down_dst_off,
selectedbuf,
ds4_gpu_tensor_offset(selected),
down_smem,
down_nsg,
down_rows_per_group_is_nr0);
}
}
DS4_METAL_PROFILE_MOE_STAGE("down");
if (ok &&
n_expert > 1 &&
!direct_down_sum &&
!use_q4_batch_expert_table &&
!use_iq2_batch_selected_addr) {
ok = ds4_gpu_encode_moe_sum_experts(cb,
down_dst,
down_dst_off,
outbuf,
ds4_gpu_tensor_offset(out),
out_dim,
n_expert,
n_tokens);
}
DS4_METAL_PROFILE_MOE_STAGE("sum");
if (!ok) {
fprintf(stderr,
"ds4: Metal routed batch MoE failed before submit layer=%u tokens=%u "
"gate=%s down=%s path=%s\n",
layer_index,
n_tokens,
ds4_gpu_metal_tensor_type_name(gate_type),
ds4_gpu_metal_tensor_type_name(down_type),
moe_path);
return 0;
}
if (!ds4_gpu_finish_command_buffer(cb, owned, "routed batch MoE")) {
if (use_cached_batch) {
ds4_gpu_stream_expert_cache_clear_layer(layer_index);
}
return 0;
}
if (use_cached_batch) {
if (!owned) {
if (ds4_gpu_end_commands() == 0) {
ds4_gpu_stream_expert_cache_clear_layer(layer_index);
return 0;
}
if (ds4_gpu_begin_commands() == 0) {
fprintf(stderr,
"ds4: Metal streaming prefill batch selected addr layer=%u "
"failed to reopen command batch after execution\n",
layer_index);
return 0;
}
}
}
if (q4_batch_table_boundary) {
if (ds4_gpu_end_commands() == 0 || ds4_gpu_begin_commands() == 0) {
return 0;
}
}
#undef DS4_METAL_PROFILE_MOE_STAGE
}
return 1;
}
int ds4_gpu_hc_split_sinkhorn_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *mix,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (n_hc == 0 || n_hc > 16) return 0;
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t mix_bytes = mix_hc * sizeof(float);
const uint64_t scale_bytes = 3ull * sizeof(float);
@autoreleasepool {
id<MTLBuffer> mixbuf = ds4_gpu_tensor_buffer(mix);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t mix_tensor_bytes = ds4_gpu_tensor_bytes(mix);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
if (!mixbuf || !outbuf ||
mix_tensor_bytes < mix_bytes ||
out_tensor_bytes < mix_bytes) {
fprintf(stderr, "ds4: Metal HC split received undersized activation buffers\n");
return 0;
}
if (scale_offset > model_size || scale_bytes > model_size - scale_offset ||
base_offset > model_size || mix_bytes > model_size - base_offset) {
fprintf(stderr, "ds4: Metal HC split parameter range is outside the mapped model\n");
return 0;
}
uint64_t scale_inner = 0;
uint64_t base_inner = 0;
id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner);
id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner);
if (!scalebuf || !basebuf) return 0;
uint64_t n_rows64 = mix_tensor_bytes / mix_bytes;
const uint64_t out_rows64 = out_tensor_bytes / mix_bytes;
if (out_rows64 < n_rows64) n_rows64 = out_rows64;
if (n_rows64 == 0 || n_rows64 > UINT32_MAX) {
fprintf(stderr, "ds4: Metal HC split row count is outside supported range\n");
return 0;
}
ds4_gpu_hc_split_args args = {
.n_hc = (int32_t)n_hc,
.sinkhorn_iters = (int32_t)sinkhorn_iters,
.n_rows = (int64_t)n_rows64,
.mix_hc = (int64_t)mix_hc,
.nb01 = mix_bytes,
.nb1 = mix_bytes,
.eps = eps,
};
const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_rows64));
const NSUInteger n_tg = ((NSUInteger)n_rows64 + nth - 1u) / nth;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_hc_split_sinkhorn_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1];
[enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
[enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sinkhorn")) return 0;
}
return 1;
}
static int ds4_gpu_hc_weighted_sum_strided(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *weights,
uint64_t weight_offset,
uint64_t weight_row_stride,
uint32_t n_embd,
uint32_t n_hc,
int round_bf16,
const char *label) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !residual_hc || !weights || n_embd == 0 || n_hc == 0 ||
weight_row_stride < (uint64_t)n_hc * sizeof(float)) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> wbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes || out_tensor_bytes % out_row_bytes != 0) {
fprintf(stderr, "ds4: Metal HC weighted sum output size is not a whole token row\n");
return 0;
}
const uint64_t n_tokens64 = out_tensor_bytes / out_row_bytes;
if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
fprintf(stderr, "ds4: Metal HC weighted sum token count is outside supported range\n");
return 0;
}
const uint64_t x_row_values = (uint64_t)n_hc * n_embd;
if (x_row_values == 0 ||
x_row_values > UINT64_MAX / sizeof(float) ||
n_tokens64 > UINT64_MAX / (x_row_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / ((uint64_t)n_hc * sizeof(float))) {
fprintf(stderr, "ds4: Metal HC weighted sum activation size overflow\n");
return 0;
}
const uint64_t x_bytes = n_tokens64 * x_row_values * sizeof(float);
const uint64_t w_last = weight_offset +
(n_tokens64 - 1u) * weight_row_stride +
(uint64_t)n_hc * sizeof(float);
if (!xbuf || !wbuf || !outbuf ||
ds4_gpu_tensor_bytes(residual_hc) < x_bytes ||
ds4_gpu_tensor_bytes(weights) < w_last) {
fprintf(stderr, "ds4: Metal HC weighted sum received undersized activation buffers\n");
return 0;
}
ds4_gpu_hc_weighted_sum_args args = {
.n_embd = n_embd,
.n_hc = n_hc,
.n_tokens = (int64_t)n_tokens64,
.nb_x0 = sizeof(float),
.nb_x1 = (uint64_t)n_embd * sizeof(float),
.nb_x2 = (uint64_t)n_hc * n_embd * sizeof(float),
.nb_w0 = sizeof(float),
.nb_w1 = weight_row_stride,
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_embd * sizeof(float),
.round_bf16 = (uint64_t)(round_bf16 != 0),
};
const uint64_t n_elem = (uint64_t)n_embd * n_tokens64;
const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_hc_weighted_sum_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:1];
[enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(weights) + (NSUInteger)weight_offset atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, label)) return 0;
}
return 1;
}
int ds4_gpu_hc_weighted_sum_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *weights,
uint32_t n_embd,
uint32_t n_hc) {
return ds4_gpu_hc_weighted_sum_strided(out,
residual_hc,
weights,
0,
(uint64_t)n_hc * sizeof(float),
n_embd,
n_hc,
0,
"HC weighted sum");
}
int ds4_gpu_hc_weighted_sum_bf16_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *weights,
uint32_t n_embd,
uint32_t n_hc) {
return ds4_gpu_hc_weighted_sum_strided(out, residual_hc, weights, 0,
(uint64_t)n_hc * sizeof(float), n_embd, n_hc, 1,
"HC weighted sum");
}
int ds4_gpu_hc_weighted_sum_split_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
return ds4_gpu_hc_weighted_sum_strided(out,
residual_hc,
split,
0,
mix_hc * sizeof(float),
n_embd,
n_hc,
0,
"HC weighted sum split");
}
/* Release decode fused HC pre-sublayer operation. The graph driver owns the
* optional reference fallback so this function stays a direct fused dispatch. */
int ds4_gpu_hc_split_weighted_sum_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix,
const ds4_gpu_tensor *residual_hc,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !split || !mix || !residual_hc || !model_map ||
n_embd == 0 || n_hc == 0) {
return 0;
}
if (n_hc != 4) {
fprintf(stderr, "ds4: Metal fused HC split/sum is specialized for HC=4\n");
return 0;
}
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t mix_bytes = mix_hc * sizeof(float);
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t scale_bytes = 3ull * sizeof(float);
@autoreleasepool {
id<MTLBuffer> mixbuf = ds4_gpu_tensor_buffer(mix);
id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes ||
out_tensor_bytes % out_row_bytes != 0) {
fprintf(stderr, "ds4: Metal fused HC split/sum output size is not a whole token row\n");
return 0;
}
const uint64_t n_rows64 = out_tensor_bytes / out_row_bytes;
if (n_rows64 == 0 || n_rows64 > UINT32_MAX ||
n_rows64 > UINT64_MAX / mix_bytes ||
n_rows64 > UINT64_MAX / residual_row_bytes) {
fprintf(stderr, "ds4: Metal fused HC split/sum row count is outside supported range\n");
return 0;
}
const uint64_t mix_total_bytes = n_rows64 * mix_bytes;
const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes;
if (!mixbuf || !splitbuf || !xbuf || !outbuf ||
ds4_gpu_tensor_bytes(mix) < mix_total_bytes ||
ds4_gpu_tensor_bytes(split) < mix_total_bytes ||
ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes) {
fprintf(stderr, "ds4: Metal fused HC split/sum received undersized activation buffers\n");
return 0;
}
if (scale_offset > model_size || scale_bytes > model_size - scale_offset ||
base_offset > model_size || mix_bytes > model_size - base_offset) {
fprintf(stderr, "ds4: Metal fused HC split/sum parameter range is outside the mapped model\n");
return 0;
}
uint64_t scale_inner = 0;
uint64_t base_inner = 0;
id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner);
id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner);
if (!scalebuf || !basebuf) return 0;
ds4_gpu_hc_split_weighted_sum_args args = {
.n_embd = (int64_t)n_embd,
.n_hc = (int32_t)n_hc,
.sinkhorn_iters = (int32_t)sinkhorn_iters,
.n_rows = (int64_t)n_rows64,
.mix_hc = (int64_t)mix_hc,
.nb_mix1 = mix_bytes,
.nb_split1 = mix_bytes,
.nb_x0 = sizeof(float),
.nb_x1 = (uint64_t)n_embd * sizeof(float),
.nb_x2 = residual_row_bytes,
.nb0 = sizeof(float),
.nb1 = out_row_bytes,
.eps = eps,
};
NSUInteger nth = g_hc_split_weighted_sum_pipeline.maxTotalThreadsPerThreadgroup;
if (nth > 256u) nth = 256u;
if (nth > (NSUInteger)n_embd) nth = (NSUInteger)n_embd;
if (nth == 0) nth = 1u;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_hc_split_weighted_sum_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1];
[enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
[enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6];
[enc setThreadgroupMemoryLength:(NSUInteger)n_hc * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum fused")) return 0;
}
return 1;
}
/* HC-pre plus the immediately following weighted RMSNorm, specialized for
* DS4's HC=4 shape. Both decode and batched prefill use this implementation;
* the kernel preserves their established single-row and batched scale formulas. */
/* pre_w selects the DeepSeek V4.1 variant: one row, the weighted sum's pre
* weights from that tensor, bf16 rounding. */
static int ds4_gpu_hc_split_weighted_sum_norm_impl(
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *pre_w,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint64_t norm_weight_offset,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps,
float norm_eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !norm_out || !split || !mix || !residual_hc || !model_map ||
n_hc != 4 || (n_embd & 3u) != 0) {
return 0;
}
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t mix_bytes = mix_hc * sizeof(float);
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t scale_bytes = 3ull * sizeof(float);
@autoreleasepool {
id<MTLBuffer> mixbuf = ds4_gpu_tensor_buffer(mix);
id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> normbuf = ds4_gpu_tensor_buffer(norm_out);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
if (out_row_bytes == 0 || out_tensor_bytes < out_row_bytes ||
(!pre_w && out_tensor_bytes % out_row_bytes != 0)) {
fprintf(stderr, "ds4: Metal fused HC split/sum/norm output size is not a whole token row\n");
return 0;
}
const uint64_t n_rows64 = pre_w ? 1u : out_tensor_bytes / out_row_bytes;
if (n_rows64 == 0 || n_rows64 > UINT32_MAX ||
n_rows64 > UINT64_MAX / mix_bytes ||
n_rows64 > UINT64_MAX / residual_row_bytes) {
fprintf(stderr, "ds4: Metal fused HC split/sum/norm row count is outside supported range\n");
return 0;
}
const uint64_t mix_total_bytes = n_rows64 * mix_bytes;
const uint64_t residual_total_bytes = n_rows64 * residual_row_bytes;
const uint64_t out_total_bytes = n_rows64 * out_row_bytes;
if (!mixbuf || !splitbuf || !xbuf || !outbuf || !normbuf ||
ds4_gpu_tensor_bytes(mix) < mix_total_bytes ||
ds4_gpu_tensor_bytes(split) < mix_total_bytes ||
ds4_gpu_tensor_bytes(residual_hc) < residual_total_bytes ||
ds4_gpu_tensor_bytes(norm_out) < out_total_bytes) {
fprintf(stderr, "ds4: Metal fused HC split/sum/norm received undersized activation buffers\n");
return 0;
}
if (scale_offset > model_size || scale_bytes > model_size - scale_offset ||
base_offset > model_size || mix_bytes > model_size - base_offset ||
norm_weight_offset > model_size || out_row_bytes > model_size - norm_weight_offset) {
fprintf(stderr, "ds4: Metal fused HC split/sum/norm parameter range is outside the mapped model\n");
return 0;
}
uint64_t scale_inner = 0;
uint64_t base_inner = 0;
uint64_t norm_inner = 0;
id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, scale_bytes, &scale_inner);
id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner);
id<MTLBuffer> normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_weight_offset, out_row_bytes, &norm_inner);
if (!scalebuf || !basebuf || !normwbuf) return 0;
id<MTLComputePipelineState> pipeline = pre_w ?
ds4_gpu_get_pipeline("kernel_dsv41_hc_split_sum_norm") :
ds4_gpu_hot_pipeline(g_hc_split_weighted_sum_norm_pipeline,
"kernel_dsv4_hc_split_weighted_sum_norm4");
if (!pipeline || (pre_w && ds4_gpu_tensor_bytes(pre_w) < (uint64_t)n_hc * sizeof(float))) return 0;
ds4_gpu_hc_split_weighted_sum_norm_args args = {
.n_embd = (int64_t)n_embd,
.n_hc = (int32_t)n_hc,
.sinkhorn_iters = (int32_t)sinkhorn_iters,
.n_rows = (int64_t)n_rows64,
.mix_hc = (int64_t)mix_hc,
.nb_mix1 = mix_bytes,
.nb_split1 = mix_bytes,
.nb_x0 = sizeof(float),
.nb_x1 = (uint64_t)n_embd * sizeof(float),
.nb_x2 = residual_row_bytes,
.nb0 = sizeof(float),
.nb1 = out_row_bytes,
.nb_norm1 = out_row_bytes,
.eps = eps,
.norm_eps = norm_eps,
};
NSUInteger nth = ds4_gpu_rms_norm_threads(n_embd);
if (nth > pipeline.maxTotalThreadsPerThreadgroup) {
fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu threads but pipeline supports %lu\n",
(unsigned long)nth,
(unsigned long)pipeline.maxTotalThreadsPerThreadgroup);
return 0;
}
const NSUInteger shared_bytes = ((NSUInteger)n_embd + 4u + 32u) * sizeof(float);
const NSUInteger max_shared = [g_device maxThreadgroupMemoryLength];
if (max_shared != 0 && shared_bytes > max_shared) {
fprintf(stderr, "ds4: Metal fused HC split/sum/norm requires %lu bytes of threadgroup memory but device supports %lu\n",
(unsigned long)shared_bytes,
(unsigned long)max_shared);
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:1];
[enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
[enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:5];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:6];
[enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:7];
[enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:8];
if (pre_w) [enc setBuffer:ds4_gpu_tensor_buffer(pre_w) offset:ds4_gpu_tensor_offset(pre_w) atIndex:9];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)n_rows64, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC split/sum/norm fused")) return 0;
}
return 1;
}
int ds4_gpu_hc_split_weighted_sum_norm_tensor(
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix,
const ds4_gpu_tensor *residual_hc,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint64_t norm_weight_offset,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps,
float norm_eps) {
return ds4_gpu_hc_split_weighted_sum_norm_impl(out, norm_out, split, mix, residual_hc, NULL,
model_map, model_size, scale_offset, base_offset, norm_weight_offset,
n_embd, n_hc, sinkhorn_iters, eps, norm_eps);
}
int ds4_gpu_dsv41_hc_project(ds4_gpu_tensor *mix, const ds4_gpu_tensor *residual,
const void *model_map, uint64_t model_size, uint64_t fn_offset,
uint32_t n, uint32_t mix_dim, float eps) {
return ds4_gpu_hc_rms_norm_mix_f16_tensor(mix, residual, model_map, model_size, fn_offset,
n, mix_dim, eps);
}
int ds4_gpu_dsv41_hc_block_input_rows(ds4_gpu_tensor *mix, ds4_gpu_tensor *x, ds4_gpu_tensor *norm,
ds4_gpu_tensor *split, const ds4_gpu_tensor *stream,
const ds4_gpu_tensor *pre, const void *model_map, uint64_t model_size,
uint64_t fn_offset, uint64_t scale_offset, uint64_t base_offset,
uint64_t norm_offset, uint32_t n, uint32_t mix_dim, uint32_t n_embd,
uint32_t n_hc, uint32_t sinkhorn_iters, float hc_eps, float norm_eps,
uint32_t rows, uint32_t pre_stride) {
static id<MTLComputePipelineState> pipeline;
enum { max_rows = 8, row_words = 64 }; /* per row: the count and 32 partial sums, line-aligned */
if (!g_initialized && !ds4_gpu_init()) return 0;
/* the norm/mix half needs the 1024-thread norm and an empty matvec tail */
if (!mix || !x || !norm || !split || !stream || !pre || !model_map || n < 4096u ||
(n % 4096u) != 0u || mix_dim != 24u || n_hc != 4u || (n_embd & 3u) != 0u ||
n != n_hc * n_embd || !rows || rows > max_rows || pre_stride < n_hc) return 0;
@autoreleasepool {
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t mix_bytes = mix_hc * sizeof(float);
const uint64_t row_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t weight_bytes = (uint64_t)n * sizeof(uint16_t) * mix_dim;
if (fn_offset > model_size || weight_bytes > model_size - fn_offset ||
scale_offset > model_size || 3u * sizeof(float) > model_size - scale_offset ||
base_offset > model_size || mix_bytes > model_size - base_offset ||
norm_offset > model_size || row_bytes > model_size - norm_offset ||
ds4_gpu_tensor_bytes(stream) < (uint64_t)n * rows * sizeof(float) ||
ds4_gpu_tensor_bytes(mix) < mix_bytes * rows || ds4_gpu_tensor_bytes(split) < mix_bytes * rows ||
ds4_gpu_tensor_bytes(x) < row_bytes * rows || ds4_gpu_tensor_bytes(norm) < row_bytes * rows ||
ds4_gpu_tensor_bytes(pre) < ((uint64_t)(rows - 1u) * pre_stride + n_hc) * sizeof(float)) return 0;
uint64_t fn_inner = 0, scale_inner = 0, base_inner = 0, norm_inner = 0;
id<MTLBuffer> fnbuf = ds4_gpu_wrap_model_range(model_map, model_size, fn_offset, weight_bytes, &fn_inner);
id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size, scale_offset, 3u * sizeof(float), &scale_inner);
id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size, base_offset, mix_bytes, &base_inner);
id<MTLBuffer> normwbuf = ds4_gpu_wrap_model_range(model_map, model_size, norm_offset, row_bytes, &norm_inner);
if (!fnbuf || !scalebuf || !basebuf || !normwbuf) return 0;
if (!pipeline) pipeline = ds4_gpu_get_pipeline("kernel_dsv41_hc_block_input");
/* the arrival counters live with the other completion buffers, which a
* failed command buffer drops */
if (!g_dsv4_completion_cache) return 0;
id<MTLBuffer> completion = [g_dsv4_completion_cache objectForKey:@"v41_block_input"];
if (!completion) {
completion = [g_device newBufferWithLength:(NSUInteger)max_rows * row_words * sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (!completion) return 0;
memset([completion contents], 0, (size_t)max_rows * row_words * sizeof(uint32_t));
[g_dsv4_completion_cache setObject:completion forKey:@"v41_block_input"];
}
[g_transient_buffers addObject:completion];
const NSUInteger shared_bytes = 96u * sizeof(float); /* the norm/mix reductions */
if (!pipeline || !completion || pipeline.maxTotalThreadsPerThreadgroup < 256u ||
shared_bytes > [g_device maxThreadgroupMemoryLength]) return 0;
const struct { int32_t n, out_dim; float eps; } args = {(int32_t)n, (int32_t)mix_dim, norm_eps};
const ds4_gpu_hc_split_weighted_sum_norm_args split_args = {
.n_embd = (int64_t)n_embd,
.n_hc = (int32_t)n_hc,
.sinkhorn_iters = (int32_t)sinkhorn_iters,
.n_rows = 1,
.mix_hc = (int64_t)mix_hc,
.nb_mix1 = mix_bytes,
.nb_split1 = mix_bytes,
.nb_x0 = sizeof(float),
.nb_x1 = row_bytes,
.nb_x2 = (uint64_t)n_hc * row_bytes,
.nb0 = sizeof(float),
.nb1 = row_bytes,
.nb_norm1 = row_bytes,
.eps = hc_eps,
.norm_eps = norm_eps,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBytes:&split_args length:sizeof(split_args) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(stream) offset:ds4_gpu_tensor_offset(stream) atIndex:2];
[enc setBuffer:fnbuf offset:(NSUInteger)fn_inner atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(mix) offset:ds4_gpu_tensor_offset(mix) atIndex:4];
[enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:5];
[enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:6];
[enc setBuffer:ds4_gpu_tensor_buffer(split) offset:ds4_gpu_tensor_offset(split) atIndex:7];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:8];
[enc setBuffer:normwbuf offset:(NSUInteger)norm_inner atIndex:9];
[enc setBuffer:ds4_gpu_tensor_buffer(norm) offset:ds4_gpu_tensor_offset(norm) atIndex:10];
[enc setBuffer:ds4_gpu_tensor_buffer(pre) offset:ds4_gpu_tensor_offset(pre) atIndex:11];
[enc setBuffer:completion offset:0 atIndex:12];
[enc setBytes:&pre_stride length:sizeof(pre_stride) atIndex:13];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((mix_dim + 1u) / 2u + 4u, rows, 1)
threadsPerThreadgroup:MTLSizeMake(32, 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 HC block input");
}
}
int ds4_gpu_dsv41_hc_input(ds4_gpu_tensor *x, ds4_gpu_tensor *norm, ds4_gpu_tensor *split,
const ds4_gpu_tensor *mix, const ds4_gpu_tensor *residual,
const ds4_gpu_tensor *pre, const void *model_map, uint64_t model_size,
uint64_t scale_offset, uint64_t base_offset, uint64_t norm_offset,
uint32_t n_embd, uint32_t n_hc, uint32_t sinkhorn_iters,
float hc_eps, float norm_eps) {
return ds4_gpu_hc_split_weighted_sum_norm_impl(x, norm, split, mix, residual, pre,
model_map, model_size, scale_offset, base_offset, norm_offset,
n_embd, n_hc, sinkhorn_iters, hc_eps, norm_eps);
}
typedef struct {
int32_t n;
int32_t out_dim;
float eps;
} ds4_gpu_hc_norm_mix_args;
int ds4_gpu_hc_rms_norm_mix_f16_available(void) {
if (!g_initialized && !ds4_gpu_init()) return 0;
static int cached = -1;
if (cached < 0) {
cached = ds4_gpu_get_pipeline("kernel_dsv4_hc_rms_norm_mix_f16") != nil;
}
return cached;
}
/* Fused decode HC-pre: unweighted RMSNorm over the flattened HC row plus the
* F16 HC-mix matvec in one dispatch. The kernel reproduces the standalone
* kernels' reduction trees bit-exactly (see metal/dsv4_hc.metal), so this is
* purely a dispatch-count optimization. */
int ds4_gpu_hc_rms_norm_mix_f16_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t n,
uint32_t out_dim,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
/* the norm runs 1024 threads and the matvec tail is empty at these widths */
if (!out || !x || !model_map || n < 4096u || (n % 4096u) != 0u || out_dim != 24u) return 0;
@autoreleasepool {
const uint64_t row_bytes = (uint64_t)n * sizeof(uint16_t);
const uint64_t weight_bytes = row_bytes * out_dim;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal fused HC norm/mix weight range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
if (!xbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < (uint64_t)n * sizeof(float) ||
ds4_gpu_tensor_bytes(out) < (uint64_t)out_dim * sizeof(float)) {
fprintf(stderr, "ds4: Metal fused HC norm/mix received undersized activation buffers\n");
return 0;
}
uint64_t weight_inner = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
weight_offset, weight_bytes,
&weight_inner);
if (!wbuf) return 0;
const bool use_cluster2 = n == 16384u &&
ds4_gpu_device_is_m5_apple_silicon() &&
getenv("DS4_METAL_DISABLE_M5_HC_NORM_MIX_CLUSTER2") == NULL;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(
use_cluster2 ? "kernel_dsv4_hc_rms_norm_mix_f16_cluster2" :
"kernel_dsv4_hc_rms_norm_mix_f16");
if (!pipeline ||
(use_cluster2 && pipeline.maxTotalThreadsPerThreadgroup < 512u)) {
return 0;
}
const NSUInteger shared_floats =
use_cluster2 ? 32u + 32u * 2u * 2u : 32u + 32u * 2u;
ds4_gpu_hc_norm_mix_args args = {
.n = (int32_t)n,
.out_dim = (int32_t)out_dim,
.eps = eps,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:wbuf offset:(NSUInteger)weight_inner atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:shared_floats * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:
MTLSizeMake(use_cluster2 ? (out_dim + 3u) / 4u :
(out_dim + 1u) / 2u,
1, 1)
threadsPerThreadgroup:
MTLSizeMake(32, use_cluster2 ? 16 : 8, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "fused HC norm/mix")) return 0;
}
return 1;
}
int ds4_gpu_hc_rms_norm_mix_split_norm_f16_tensor(
ds4_gpu_tensor *mix,
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *residual_hc,
const void *model_map,
uint64_t model_size,
uint64_t mix_weight_offset,
uint64_t scale_offset,
uint64_t base_offset,
uint64_t norm_weight_offset,
uint32_t n,
uint32_t mix_dim,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps,
float hc_eps,
float norm_eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!mix || !out || !norm_out || !split || !residual_hc || !model_map ||
n != 16384u || mix_dim != 24u || n_embd != 4096u || n_hc != 4u) {
return 0;
}
@autoreleasepool {
const uint64_t mix_weight_bytes =
(uint64_t)n * (uint64_t)mix_dim * sizeof(uint16_t);
const uint64_t mix_bytes = (uint64_t)mix_dim * sizeof(float);
const uint64_t residual_bytes = (uint64_t)n * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t scale_bytes = 3u * sizeof(float);
if (mix_weight_offset > model_size ||
mix_weight_bytes > model_size - mix_weight_offset ||
scale_offset > model_size || scale_bytes > model_size - scale_offset ||
base_offset > model_size || mix_bytes > model_size - base_offset ||
norm_weight_offset > model_size ||
out_bytes > model_size - norm_weight_offset) {
fprintf(stderr,
"ds4: Metal compound HC producer parameter range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> mixbuf = ds4_gpu_tensor_buffer(mix);
id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> normbuf = ds4_gpu_tensor_buffer(norm_out);
if (!xbuf || !mixbuf || !splitbuf || !outbuf || !normbuf ||
ds4_gpu_tensor_bytes(residual_hc) < residual_bytes ||
ds4_gpu_tensor_bytes(mix) < mix_bytes ||
ds4_gpu_tensor_bytes(split) < mix_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(norm_out) < out_bytes) {
fprintf(stderr,
"ds4: Metal compound HC producer received undersized activation buffers\n");
return 0;
}
uint64_t mix_weight_inner = 0;
uint64_t scale_inner = 0;
uint64_t base_inner = 0;
uint64_t norm_inner = 0;
id<MTLBuffer> mix_weight = ds4_gpu_wrap_model_range(
model_map, model_size, mix_weight_offset, mix_weight_bytes,
&mix_weight_inner);
id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(
model_map, model_size, scale_offset, scale_bytes, &scale_inner);
id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(
model_map, model_size, base_offset, mix_bytes, &base_inner);
id<MTLBuffer> norm_weight = ds4_gpu_wrap_model_range(
model_map, model_size, norm_weight_offset, out_bytes, &norm_inner);
if (!mix_weight || !scalebuf || !basebuf || !norm_weight) return 0;
if (!g_dsv4_hc_producer_pre_norm_pipeline) {
g_dsv4_hc_producer_pre_norm_pipeline = ds4_gpu_get_pipeline(
"kernel_dsv4_hc_rms_norm_mix_f16_cluster2_pre_norm");
}
id<MTLComputePipelineState> producer =
g_dsv4_hc_producer_pre_norm_pipeline;
if (!producer || producer.maxTotalThreadsPerThreadgroup < 512u) {
return 0;
}
if (!g_dsv4_completion_cache) return 0;
const NSUInteger mix_offset = ds4_gpu_tensor_offset(mix);
id<MTLBuffer> completion = nil;
if (g_dsv4_hc_producer_last_mix_buffer == mixbuf &&
g_dsv4_hc_producer_last_mix_offset == mix_offset) {
completion = g_dsv4_hc_producer_last_completion;
} else {
NSString *completion_key = [NSString stringWithFormat:@"%p:%llu",
(void *)mixbuf, (unsigned long long)mix_offset];
completion = [g_dsv4_completion_cache
objectForKey:completion_key];
if (!completion) {
completion = [g_device newBufferWithLength:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (!completion) return 0;
*((uint32_t *)[completion contents]) = 0u;
[g_dsv4_completion_cache
setObject:completion forKey:completion_key];
}
g_dsv4_hc_producer_last_mix_buffer = mixbuf;
g_dsv4_hc_producer_last_mix_offset = mix_offset;
g_dsv4_hc_producer_last_completion = completion;
}
[g_transient_buffers addObject:completion];
ds4_gpu_hc_norm_mix_args mix_args = {
.n = (int32_t)n,
.out_dim = (int32_t)mix_dim,
.eps = eps,
};
ds4_gpu_hc_split_weighted_sum_norm_args split_args = {
.n_embd = (int64_t)n_embd,
.n_hc = (int32_t)n_hc,
.sinkhorn_iters = (int32_t)sinkhorn_iters,
.n_rows = 1,
.mix_hc = (int64_t)mix_dim,
.nb_mix1 = mix_bytes,
.nb_split1 = mix_bytes,
.nb_x0 = sizeof(float),
.nb_x1 = (uint64_t)n_embd * sizeof(float),
.nb_x2 = residual_bytes,
.nb0 = sizeof(float),
.nb1 = out_bytes,
.nb_norm1 = out_bytes,
.eps = hc_eps,
.norm_eps = norm_eps,
};
const NSUInteger shared_floats = 32u + 4u * 32u + 4u + 32u;
const NSUInteger shared_bytes = shared_floats * sizeof(float);
const NSUInteger max_shared = [g_device maxThreadgroupMemoryLength];
if (max_shared != 0 && shared_bytes > max_shared) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:producer];
[enc setBytes:&mix_args length:sizeof(mix_args) atIndex:0];
[enc setBytes:&split_args length:sizeof(split_args) atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
[enc setBuffer:mix_weight offset:(NSUInteger)mix_weight_inner atIndex:3];
[enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:4];
[enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:5];
[enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:6];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:7];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:8];
[enc setBuffer:norm_weight offset:(NSUInteger)norm_inner atIndex:9];
[enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:10];
[enc setBuffer:completion offset:0 atIndex:11];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(6, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 16, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "compound HC producer/pre-norm")) {
return -1;
}
}
return 1;
}
int ds4_gpu_hc_expand_add_rms_norm_mix_split_norm_f16_tensor(
ds4_gpu_tensor *mix,
ds4_gpu_tensor *out,
ds4_gpu_tensor *norm_out,
ds4_gpu_tensor *split,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_prev,
const ds4_gpu_tensor *post,
const ds4_gpu_tensor *comb,
const void *model_map,
uint64_t model_size,
uint64_t mix_weight_offset,
uint64_t scale_offset,
uint64_t base_offset,
uint64_t norm_weight_offset,
uint32_t n,
uint32_t mix_dim,
uint32_t n_embd,
uint32_t n_hc,
uint32_t sinkhorn_iters,
float eps,
float hc_eps,
float norm_eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!mix || !out || !norm_out || !split || !residual_hc || !model_map ||
!block_out || !block_add || !residual_prev || !post || !comb ||
n != 16384u || mix_dim != 24u || n_embd != 4096u || n_hc != 4u) {
return 0;
}
@autoreleasepool {
const uint64_t mix_weight_bytes =
(uint64_t)n * (uint64_t)mix_dim * sizeof(uint16_t);
const uint64_t mix_bytes = (uint64_t)mix_dim * sizeof(float);
const uint64_t residual_bytes = (uint64_t)n * sizeof(float);
const uint64_t out_bytes = (uint64_t)n_embd * sizeof(float);
const uint64_t scale_bytes = 3u * sizeof(float);
if (mix_weight_offset > model_size ||
mix_weight_bytes > model_size - mix_weight_offset ||
scale_offset > model_size || scale_bytes > model_size - scale_offset ||
base_offset > model_size || mix_bytes > model_size - base_offset ||
norm_weight_offset > model_size ||
out_bytes > model_size - norm_weight_offset) {
fprintf(stderr,
"ds4: Metal compound HC producer parameter range is outside the mapped model\n");
return 0;
}
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> mixbuf = ds4_gpu_tensor_buffer(mix);
id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> normbuf = ds4_gpu_tensor_buffer(norm_out);
if (!xbuf || !mixbuf || !splitbuf || !outbuf || !normbuf ||
ds4_gpu_tensor_bytes(residual_hc) < residual_bytes ||
ds4_gpu_tensor_bytes(mix) < mix_bytes ||
ds4_gpu_tensor_bytes(split) < mix_bytes ||
ds4_gpu_tensor_bytes(out) < out_bytes ||
ds4_gpu_tensor_bytes(norm_out) < out_bytes) {
fprintf(stderr,
"ds4: Metal compound HC producer received undersized activation buffers\n");
return 0;
}
uint64_t mix_weight_inner = 0;
uint64_t scale_inner = 0;
uint64_t base_inner = 0;
uint64_t norm_inner = 0;
id<MTLBuffer> mix_weight = ds4_gpu_wrap_model_range(
model_map, model_size, mix_weight_offset, mix_weight_bytes,
&mix_weight_inner);
id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(
model_map, model_size, scale_offset, scale_bytes, &scale_inner);
id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(
model_map, model_size, base_offset, mix_bytes, &base_inner);
id<MTLBuffer> norm_weight = ds4_gpu_wrap_model_range(
model_map, model_size, norm_weight_offset, out_bytes, &norm_inner);
if (!mix_weight || !scalebuf || !basebuf || !norm_weight) return 0;
if (!g_dsv4_hc_expand_producer_pre_norm_pipeline) {
g_dsv4_hc_expand_producer_pre_norm_pipeline = ds4_gpu_get_pipeline(
"kernel_dsv4_hc_expand4_rms_norm_mix_f16_cluster2_pre_norm");
}
id<MTLComputePipelineState> producer =
g_dsv4_hc_expand_producer_pre_norm_pipeline;
id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
id<MTLBuffer> addbuf = ds4_gpu_tensor_buffer(block_add);
id<MTLBuffer> prevbuf = ds4_gpu_tensor_buffer(residual_prev);
id<MTLBuffer> postbuf = ds4_gpu_tensor_buffer(post);
id<MTLBuffer> combbuf = ds4_gpu_tensor_buffer(comb);
if (!blockbuf || !addbuf || !prevbuf || !postbuf || !combbuf ||
ds4_gpu_tensor_bytes(block_out) < out_bytes ||
ds4_gpu_tensor_bytes(block_add) < out_bytes ||
ds4_gpu_tensor_bytes(residual_prev) < residual_bytes ||
ds4_gpu_tensor_bytes(post) < (uint64_t)n_hc * sizeof(float) ||
ds4_gpu_tensor_bytes(comb) < (uint64_t)n_hc * n_hc * sizeof(float)) {
fprintf(stderr,
"ds4: Metal fused HC expand/producer received undersized expand buffers\n");
return 0;
}
if (!producer || producer.maxTotalThreadsPerThreadgroup < 512u) {
return 0;
}
if (!g_dsv4_completion_cache) return 0;
const NSUInteger mix_offset = ds4_gpu_tensor_offset(mix);
id<MTLBuffer> completion = nil;
if (g_dsv4_hc_producer_last_mix_buffer == mixbuf &&
g_dsv4_hc_producer_last_mix_offset == mix_offset) {
completion = g_dsv4_hc_producer_last_completion;
} else {
NSString *completion_key = [NSString stringWithFormat:@"%p:%llu",
(void *)mixbuf, (unsigned long long)mix_offset];
completion = [g_dsv4_completion_cache
objectForKey:completion_key];
if (!completion) {
completion = [g_device newBufferWithLength:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (!completion) return 0;
*((uint32_t *)[completion contents]) = 0u;
[g_dsv4_completion_cache
setObject:completion forKey:completion_key];
}
g_dsv4_hc_producer_last_mix_buffer = mixbuf;
g_dsv4_hc_producer_last_mix_offset = mix_offset;
g_dsv4_hc_producer_last_completion = completion;
}
[g_transient_buffers addObject:completion];
/* Monotonic cross-threadgroup barrier counter for this mix tensor,
* plus the host-side dispatch count that sets the target. */
NSString *barrier_key = [NSString stringWithFormat:@"%p:%llu",
(void *)mixbuf, (unsigned long long)mix_offset];
if (!g_dsv4_hc_barrier_cache) {
g_dsv4_hc_barrier_cache = [NSMutableDictionary dictionary];
g_dsv4_hc_barrier_gen = [NSMutableDictionary dictionary];
}
id<MTLBuffer> barrier = [g_dsv4_hc_barrier_cache objectForKey:barrier_key];
if (!barrier) {
barrier = [g_device newBufferWithLength:sizeof(uint32_t)
options:MTLResourceStorageModeShared];
if (!barrier) return 0;
*((uint32_t *)[barrier contents]) = 0u;
[g_dsv4_hc_barrier_cache setObject:barrier forKey:barrier_key];
}
const uint32_t gen =
(uint32_t)[[g_dsv4_hc_barrier_gen objectForKey:barrier_key] unsignedIntValue] + 1u;
[g_dsv4_hc_barrier_gen setObject:@(gen) forKey:barrier_key];
const uint32_t barrier_target = 6u * gen;
[g_transient_buffers addObject:barrier];
ds4_gpu_hc_expand_args ex = {
.n_embd = n_embd,
.n_hc = n_hc,
.n_tokens = 1,
.nb_block0 = sizeof(float),
.nb_block1 = (uint64_t)n_embd * sizeof(float),
.nb_add0 = sizeof(float),
.nb_add1 = (uint64_t)n_embd * sizeof(float),
.nb_res0 = sizeof(float),
.nb_res1 = (uint64_t)n_embd * sizeof(float),
.nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
.nb_post0 = sizeof(float),
.nb_post1 = (uint64_t)n_hc * sizeof(float),
.nb_comb0 = sizeof(float),
.nb_comb1 = (uint64_t)n_hc * sizeof(float),
.nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float),
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_embd * sizeof(float),
.nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
.has_add = 1,
};
ds4_gpu_hc_norm_mix_args mix_args = {
.n = (int32_t)n,
.out_dim = (int32_t)mix_dim,
.eps = eps,
};
ds4_gpu_hc_split_weighted_sum_norm_args split_args = {
.n_embd = (int64_t)n_embd,
.n_hc = (int32_t)n_hc,
.sinkhorn_iters = (int32_t)sinkhorn_iters,
.n_rows = 1,
.mix_hc = (int64_t)mix_dim,
.nb_mix1 = mix_bytes,
.nb_split1 = mix_bytes,
.nb_x0 = sizeof(float),
.nb_x1 = (uint64_t)n_embd * sizeof(float),
.nb_x2 = residual_bytes,
.nb0 = sizeof(float),
.nb1 = out_bytes,
.nb_norm1 = out_bytes,
.eps = hc_eps,
.norm_eps = norm_eps,
};
const NSUInteger shared_floats = 32u + 4u * 32u + 4u + 32u;
const NSUInteger shared_bytes = shared_floats * sizeof(float);
const NSUInteger max_shared = [g_device maxThreadgroupMemoryLength];
if (max_shared != 0 && shared_bytes > max_shared) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:producer];
[enc setBytes:&mix_args length:sizeof(mix_args) atIndex:0];
[enc setBytes:&split_args length:sizeof(split_args) atIndex:1];
[enc setBytes:&ex length:sizeof(ex) atIndex:2];
[enc setBytes:&barrier_target length:sizeof(barrier_target) atIndex:3];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:4];
[enc setBuffer:mix_weight offset:(NSUInteger)mix_weight_inner atIndex:5];
[enc setBuffer:mixbuf offset:ds4_gpu_tensor_offset(mix) atIndex:6];
[enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:7];
[enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:8];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) atIndex:9];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:10];
[enc setBuffer:norm_weight offset:(NSUInteger)norm_inner atIndex:11];
[enc setBuffer:normbuf offset:ds4_gpu_tensor_offset(norm_out) atIndex:12];
[enc setBuffer:completion offset:0 atIndex:13];
[enc setBuffer:barrier offset:0 atIndex:14];
[enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:15];
[enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:16];
[enc setBuffer:prevbuf offset:ds4_gpu_tensor_offset(residual_prev) atIndex:17];
[enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:18];
[enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:19];
[enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(6, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 16, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "fused HC expand/producer/pre-norm")) {
return -1;
}
}
return 1;
}
int ds4_gpu_output_hc_weights_tensor(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *pre,
const void *model_map,
uint64_t model_size,
uint64_t scale_offset,
uint64_t base_offset,
uint32_t n_hc,
float eps) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !pre || !model_map || n_hc == 0) return 0;
@autoreleasepool {
if ((n_hc % 4u) != 0) {
fprintf(stderr, "ds4: Metal output HC weights requires a multiple-of-4 HC width\n");
return 0;
}
id<MTLBuffer> prebuf = ds4_gpu_tensor_buffer(pre);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out);
if (row_bytes == 0 || out_tensor_bytes < row_bytes || out_tensor_bytes % row_bytes != 0) {
fprintf(stderr, "ds4: Metal output HC weights size is not a whole token row\n");
return 0;
}
const uint64_t n_tokens64 = out_tensor_bytes / row_bytes;
if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX ||
n_tokens64 > UINT64_MAX / row_bytes) {
fprintf(stderr, "ds4: Metal output HC weights token count is outside supported range\n");
return 0;
}
const uint64_t bytes = n_tokens64 * row_bytes;
if (!prebuf || !outbuf ||
ds4_gpu_tensor_bytes(pre) < bytes ||
ds4_gpu_tensor_bytes(out) < bytes) {
fprintf(stderr, "ds4: Metal output HC weights received undersized buffers\n");
return 0;
}
uint64_t scale_inner = 0;
uint64_t base_inner = 0;
id<MTLBuffer> scalebuf = ds4_gpu_wrap_model_range(model_map, model_size,
scale_offset, sizeof(float),
&scale_inner);
id<MTLBuffer> basebuf = ds4_gpu_wrap_model_range(model_map, model_size,
base_offset, row_bytes,
&base_inner);
if (!scalebuf || !basebuf) return 0;
const bool require_weights4 =
getenv("DS4_METAL_REQUIRE_OUTPUT_HC_WEIGHTS4") != NULL;
const bool weights4_shape = n_hc == 4u && n_tokens64 == 1u;
const bool use_weights4 =
weights4_shape && !g_quality_mode &&
g_output_hc_weights4_pipeline != nil &&
(ds4_gpu_device_name_contains("M3") ||
(g_test_flags & DS4_GPU_TEST_OUTPUT_HC_WEIGHTS4) != 0u) &&
g_output_hc_weights4_pipeline.maxTotalThreadsPerThreadgroup >= 2u;
if (require_weights4 && weights4_shape && !use_weights4) {
fprintf(stderr,
"ds4: required Metal output HC weights4 kernel was not selected\n");
return 0;
}
if (use_weights4) {
ds4_gpu_output_hc_weights4_args args = {
.post_scale = 1.0f,
.eps = eps,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_output_hc_weights4_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1];
[enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
[enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:3];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1)
threadsPerThreadgroup:MTLSizeMake(2, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(
cb, owned, "output HC weights4")) {
return 0;
}
return 1;
}
const uint32_t n_tokens = (uint32_t)n_tokens64;
ds4_gpu_bin_args mul_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, 1);
ds4_gpu_bin_args add_args = ds4_gpu_make_bin_rows_args(n_hc, n_tokens, n_hc);
ds4_gpu_unary_args sigmoid_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 0.0f, 0.0f);
ds4_gpu_unary_args scale_args = ds4_gpu_make_unary_rows_args(n_hc, n_tokens, 1, 1.0f, eps);
NSUInteger mul_nth_max = g_bin_mul_scalar_pipeline.maxTotalThreadsPerThreadgroup;
if (mul_nth_max > 256u) mul_nth_max = 256u;
NSUInteger mul_nth = 1u;
while (2u * mul_nth < (NSUInteger)mul_args.ne0 && mul_nth < mul_nth_max) {
mul_nth *= 2u;
}
NSUInteger add_nth_max = g_add_pipeline.maxTotalThreadsPerThreadgroup;
if (add_nth_max > 256u) add_nth_max = 256u;
NSUInteger add_nth = 1u;
while (2u * add_nth < (NSUInteger)add_args.ne0 && add_nth < add_nth_max) {
add_nth *= 2u;
}
NSUInteger unary_nth_max = g_unary_sigmoid_pipeline.maxTotalThreadsPerThreadgroup;
if (unary_nth_max > 256u) unary_nth_max = 256u;
NSUInteger unary_nth = (NSUInteger)sigmoid_args.ne00;
if (unary_nth > unary_nth_max) unary_nth = unary_nth_max;
if (unary_nth == 0) unary_nth = 1u;
const NSUInteger unary_nk0 = ((NSUInteger)sigmoid_args.ne00 + unary_nth - 1u) / unary_nth;
const NSUInteger out_offset = ds4_gpu_tensor_offset(out);
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:g_bin_mul_scalar_pipeline];
[enc setBytes:&mul_args length:sizeof(mul_args) atIndex:0];
[enc setBuffer:prebuf offset:ds4_gpu_tensor_offset(pre) atIndex:1];
[enc setBuffer:scalebuf offset:(NSUInteger)scale_inner atIndex:2];
[enc setBuffer:outbuf offset:out_offset atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)mul_args.ne01,
(NSUInteger)mul_args.ne02,
(NSUInteger)mul_args.ne03)
threadsPerThreadgroup:MTLSizeMake(mul_nth, 1, 1)];
[enc setComputePipelineState:g_add_pipeline];
[enc setBytes:&add_args length:sizeof(add_args) atIndex:0];
[enc setBuffer:outbuf offset:out_offset atIndex:1];
[enc setBuffer:basebuf offset:(NSUInteger)base_inner atIndex:2];
[enc setBuffer:outbuf offset:out_offset atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((NSUInteger)add_args.ne01,
(NSUInteger)add_args.ne02,
(NSUInteger)add_args.ne03)
threadsPerThreadgroup:MTLSizeMake(add_nth, 1, 1)];
[enc setComputePipelineState:g_unary_sigmoid_pipeline];
[enc setBytes:&sigmoid_args length:sizeof(sigmoid_args) atIndex:0];
[enc setBuffer:outbuf offset:out_offset atIndex:1];
[enc setBuffer:outbuf offset:out_offset atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)sigmoid_args.ne01,
(NSUInteger)sigmoid_args.ne02,
(NSUInteger)sigmoid_args.ne03)
threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)];
[enc setComputePipelineState:g_unary_scale_pipeline];
[enc setBytes:&scale_args length:sizeof(scale_args) atIndex:0];
[enc setBuffer:outbuf offset:out_offset atIndex:1];
[enc setBuffer:outbuf offset:out_offset atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(unary_nk0 * (NSUInteger)scale_args.ne01,
(NSUInteger)scale_args.ne02,
(NSUInteger)scale_args.ne03)
threadsPerThreadgroup:MTLSizeMake(unary_nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "output HC weights")) return 0;
}
return 1;
}
int ds4_gpu_hc_expand_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *post,
const ds4_gpu_tensor *comb,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (n_embd == 0 || n_hc == 0) return 0;
@autoreleasepool {
id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> postbuf = ds4_gpu_tensor_buffer(post);
id<MTLBuffer> combbuf = ds4_gpu_tensor_buffer(comb);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc);
if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) {
fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n");
return 0;
}
const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes;
if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n");
return 0;
}
const uint64_t block_values = (uint64_t)n_embd;
const uint64_t hc_values = (uint64_t)n_hc * n_embd;
const uint64_t comb_values = (uint64_t)n_hc * n_hc;
if (hc_values == 0 ||
hc_values > UINT64_MAX / sizeof(float) ||
comb_values > UINT64_MAX / sizeof(float) ||
n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) {
fprintf(stderr, "ds4: Metal HC expand activation size overflow\n");
return 0;
}
const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float);
const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float);
const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float);
const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float);
if (!blockbuf || !resbuf || !postbuf || !combbuf || !outbuf ||
ds4_gpu_tensor_bytes(block_out) < block_bytes ||
ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
ds4_gpu_tensor_bytes(post) < post_bytes ||
ds4_gpu_tensor_bytes(comb) < comb_bytes) {
fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n");
return 0;
}
ds4_gpu_hc_expand_args args = {
.n_embd = n_embd,
.n_hc = n_hc,
.n_tokens = (int64_t)n_tokens64,
.nb_block0 = sizeof(float),
.nb_block1 = (uint64_t)n_embd * sizeof(float),
.nb_add0 = sizeof(float),
.nb_add1 = (uint64_t)n_embd * sizeof(float),
.nb_res0 = sizeof(float),
.nb_res1 = (uint64_t)n_embd * sizeof(float),
.nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
.nb_post0 = sizeof(float),
.nb_post1 = (uint64_t)n_hc * sizeof(float),
.nb_comb0 = sizeof(float),
.nb_comb1 = (uint64_t)n_hc * sizeof(float),
.nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float),
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_embd * sizeof(float),
.nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
.has_add = 0,
};
id<MTLComputePipelineState> expand_pipeline = g_hc_expand_pipeline;
uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64;
if (n_hc == 4) {
expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline,
"kernel_dsv4_hc_expand4");
n_elem = (uint64_t)n_embd * n_tokens64;
}
if (!expand_pipeline) return 0;
const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:expand_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1];
[enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
[enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3];
[enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4];
[enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:5];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand")) return 0;
}
return 1;
}
/* Expand of the SUM of two block vectors — the TP attention combine folded
* into the HC expand (the kernel's has_add path adds them element-wise
* before the post/comb mixing, canonical rank order preserved by argument
* position). */
int ds4_gpu_hc_expand_add_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *post,
const ds4_gpu_tensor *comb,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (n_embd == 0 || n_hc != 4) return 0;
@autoreleasepool {
id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
id<MTLBuffer> addbuf = ds4_gpu_tensor_buffer(block_add);
id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> postbuf = ds4_gpu_tensor_buffer(post);
id<MTLBuffer> combbuf = ds4_gpu_tensor_buffer(comb);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc);
if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) {
fprintf(stderr, "ds4: Metal HC expand output size is not a whole HC token row\n");
return 0;
}
const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes;
if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
fprintf(stderr, "ds4: Metal HC expand token count is outside supported range\n");
return 0;
}
const uint64_t block_values = (uint64_t)n_embd;
const uint64_t hc_values = (uint64_t)n_hc * n_embd;
const uint64_t comb_values = (uint64_t)n_hc * n_hc;
if (hc_values == 0 ||
hc_values > UINT64_MAX / sizeof(float) ||
comb_values > UINT64_MAX / sizeof(float) ||
n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / (comb_values * sizeof(float))) {
fprintf(stderr, "ds4: Metal HC expand activation size overflow\n");
return 0;
}
const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float);
const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float);
const uint64_t post_bytes = n_tokens64 * (uint64_t)n_hc * sizeof(float);
const uint64_t comb_bytes = n_tokens64 * comb_values * sizeof(float);
if (!blockbuf || !addbuf || !resbuf || !postbuf || !combbuf || !outbuf ||
ds4_gpu_tensor_bytes(block_out) < block_bytes ||
ds4_gpu_tensor_bytes(block_add) < block_bytes ||
ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
ds4_gpu_tensor_bytes(post) < post_bytes ||
ds4_gpu_tensor_bytes(comb) < comb_bytes) {
fprintf(stderr, "ds4: Metal HC expand received undersized activation buffers\n");
return 0;
}
ds4_gpu_hc_expand_args args = {
.n_embd = n_embd,
.n_hc = n_hc,
.n_tokens = (int64_t)n_tokens64,
.nb_block0 = sizeof(float),
.nb_block1 = (uint64_t)n_embd * sizeof(float),
.nb_add0 = sizeof(float),
.nb_add1 = (uint64_t)n_embd * sizeof(float),
.nb_res0 = sizeof(float),
.nb_res1 = (uint64_t)n_embd * sizeof(float),
.nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
.nb_post0 = sizeof(float),
.nb_post1 = (uint64_t)n_hc * sizeof(float),
.nb_comb0 = sizeof(float),
.nb_comb1 = (uint64_t)n_hc * sizeof(float),
.nb_comb2 = (uint64_t)n_hc * n_hc * sizeof(float),
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_embd * sizeof(float),
.nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
.has_add = 1,
};
id<MTLComputePipelineState> expand_pipeline = g_hc_expand_pipeline;
uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64;
if (n_hc == 4) {
expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline,
"kernel_dsv4_hc_expand4");
n_elem = (uint64_t)n_embd * n_tokens64;
}
if (!expand_pipeline) return 0;
const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:expand_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1];
[enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
[enc setBuffer:postbuf offset:ds4_gpu_tensor_offset(post) atIndex:3];
[enc setBuffer:combbuf offset:ds4_gpu_tensor_offset(comb) atIndex:4];
[enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add")) return 0;
}
return 1;
}
static int ds4_gpu_hc_expand_split_impl(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc,
int round_bf16) {
if (!g_initialized && !ds4_gpu_init()) {
fprintf(stderr, "ds4: Metal HC expand split could not initialize the backend\n");
return 0;
}
if (!out_hc || !block_out || !residual_hc || !split || n_embd == 0 || n_hc == 0) {
fprintf(stderr, "ds4: Metal HC expand split received invalid arguments\n");
return 0;
}
@autoreleasepool {
id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc);
if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) {
fprintf(stderr, "ds4: Metal HC expand split output size is not a whole HC token row\n");
return 0;
}
const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes;
if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
fprintf(stderr, "ds4: Metal HC expand split token count is outside supported range\n");
return 0;
}
const uint64_t block_values = (uint64_t)n_embd;
const uint64_t hc_values = (uint64_t)n_hc * n_embd;
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
if (hc_values == 0 ||
hc_values > UINT64_MAX / sizeof(float) ||
mix_hc > UINT64_MAX / sizeof(float) ||
n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) {
fprintf(stderr, "ds4: Metal HC expand split activation size overflow\n");
return 0;
}
const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float);
const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float);
const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float);
if (!blockbuf || !resbuf || !splitbuf || !outbuf ||
ds4_gpu_tensor_bytes(block_out) < block_bytes ||
ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
ds4_gpu_tensor_bytes(split) < split_bytes) {
fprintf(stderr, "ds4: Metal HC expand split received undersized activation buffers\n");
return 0;
}
ds4_gpu_hc_expand_args args = {
.n_embd = n_embd,
.n_hc = n_hc,
.n_tokens = (int64_t)n_tokens64,
.nb_block0 = sizeof(float),
.nb_block1 = (uint64_t)n_embd * sizeof(float),
.nb_add0 = sizeof(float),
.nb_add1 = (uint64_t)n_embd * sizeof(float),
.nb_res0 = sizeof(float),
.nb_res1 = (uint64_t)n_embd * sizeof(float),
.nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
.nb_post0 = sizeof(float),
.nb_post1 = mix_hc * sizeof(float),
.nb_comb0 = sizeof(float),
.nb_comb1 = (uint64_t)n_hc * sizeof(float),
.nb_comb2 = mix_hc * sizeof(float),
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_embd * sizeof(float),
.nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
.has_add = block_add != NULL,
.round_bf16 = round_bf16 != 0,
.round_block_in = block_add != NULL && round_bf16 != 0,
};
id<MTLComputePipelineState> expand_pipeline = g_hc_expand_pipeline;
uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64;
if (n_hc == 4) {
expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline,
"kernel_dsv4_hc_expand4");
n_elem = (uint64_t)n_embd * n_tokens64;
}
if (!expand_pipeline) {
fprintf(stderr, "ds4: Metal HC expand split pipeline is unavailable\n");
return 0;
}
const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) {
fprintf(stderr, "ds4: Metal HC expand split could not obtain a command buffer\n");
return 0;
}
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:expand_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1];
[enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4];
[enc setBuffer:block_add ? ds4_gpu_tensor_buffer(block_add) : blockbuf
offset:block_add ? ds4_gpu_tensor_offset(block_add) : ds4_gpu_tensor_offset(block_out) atIndex:5];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand split")) return 0;
}
return 1;
}
int ds4_gpu_hc_expand_split_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
return ds4_gpu_hc_expand_split_impl(out_hc, block_out, NULL, residual_hc, split, n_embd, n_hc, 0);
}
int ds4_gpu_hc_expand_split_bf16_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
return ds4_gpu_hc_expand_split_impl(out_hc, block_out, NULL, residual_hc, split, n_embd, n_hc, 1);
}
/* The block is block_out + block_add rounded to bf16 before the expansion. */
int ds4_gpu_hc_expand_split_add_bf16_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (!block_add || ds4_gpu_tensor_bytes(block_add) < ds4_gpu_tensor_bytes(block_out)) return 0;
return ds4_gpu_hc_expand_split_impl(out_hc, block_out, block_add, residual_hc, split, n_embd, n_hc, 1);
}
int ds4_gpu_hc_expand_split_half_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out_h,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
(void)out_hc; (void)block_out_h; (void)residual_hc; (void)split;
(void)n_embd; (void)n_hc;
return 0;
}
int ds4_gpu_hc_expand_add_split_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out_hc || !block_out || !block_add || !residual_hc || !split || n_embd == 0 || n_hc == 0) return 0;
@autoreleasepool {
id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
id<MTLBuffer> addbuf = ds4_gpu_tensor_buffer(block_add);
id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
const uint64_t hc_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t out_tensor_bytes = ds4_gpu_tensor_bytes(out_hc);
if (hc_row_bytes == 0 || out_tensor_bytes < hc_row_bytes || out_tensor_bytes % hc_row_bytes != 0) {
fprintf(stderr, "ds4: Metal HC expand add split output size is not a whole HC token row\n");
return 0;
}
const uint64_t n_tokens64 = out_tensor_bytes / hc_row_bytes;
if (n_tokens64 == 0 || n_tokens64 > UINT32_MAX) {
fprintf(stderr, "ds4: Metal HC expand add split token count is outside supported range\n");
return 0;
}
const uint64_t block_values = (uint64_t)n_embd;
const uint64_t hc_values = (uint64_t)n_hc * n_embd;
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
if (hc_values == 0 ||
hc_values > UINT64_MAX / sizeof(float) ||
mix_hc > UINT64_MAX / sizeof(float) ||
n_tokens64 > UINT64_MAX / (block_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / (hc_values * sizeof(float)) ||
n_tokens64 > UINT64_MAX / (mix_hc * sizeof(float))) {
fprintf(stderr, "ds4: Metal HC expand add split activation size overflow\n");
return 0;
}
const uint64_t block_bytes = n_tokens64 * block_values * sizeof(float);
const uint64_t hc_bytes = n_tokens64 * hc_values * sizeof(float);
const uint64_t split_bytes = n_tokens64 * mix_hc * sizeof(float);
if (!blockbuf || !addbuf || !resbuf || !splitbuf || !outbuf ||
ds4_gpu_tensor_bytes(block_out) < block_bytes ||
ds4_gpu_tensor_bytes(block_add) < block_bytes ||
ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
ds4_gpu_tensor_bytes(split) < split_bytes) {
fprintf(stderr, "ds4: Metal HC expand add split received undersized activation buffers\n");
return 0;
}
ds4_gpu_hc_expand_args args = {
.n_embd = n_embd,
.n_hc = n_hc,
.n_tokens = (int64_t)n_tokens64,
.nb_block0 = sizeof(float),
.nb_block1 = (uint64_t)n_embd * sizeof(float),
.nb_add0 = sizeof(float),
.nb_add1 = (uint64_t)n_embd * sizeof(float),
.nb_res0 = sizeof(float),
.nb_res1 = (uint64_t)n_embd * sizeof(float),
.nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
.nb_post0 = sizeof(float),
.nb_post1 = mix_hc * sizeof(float),
.nb_comb0 = sizeof(float),
.nb_comb1 = (uint64_t)n_hc * sizeof(float),
.nb_comb2 = mix_hc * sizeof(float),
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_embd * sizeof(float),
.nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
.has_add = 1,
};
id<MTLComputePipelineState> expand_pipeline = g_hc_expand_pipeline;
uint64_t n_elem = (uint64_t)n_embd * n_hc * n_tokens64;
if (n_hc == 4) {
expand_pipeline = ds4_gpu_hot_pipeline(g_dsv4_hc_expand4_pipeline,
"kernel_dsv4_hc_expand4");
n_elem = (uint64_t)n_embd * n_tokens64;
}
if (!expand_pipeline) return 0;
const NSUInteger nth = MIN((NSUInteger)256, MAX((NSUInteger)1, (NSUInteger)n_elem));
const NSUInteger n_tg = ((NSUInteger)n_elem + nth - 1u) / nth;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:expand_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:1];
[enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:2];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:3];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:4];
[enc setBuffer:addbuf offset:ds4_gpu_tensor_offset(block_add) atIndex:5];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:6];
[enc dispatchThreadgroups:MTLSizeMake(n_tg, 1, 1)
threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "HC expand add split")) return 0;
}
return 1;
}
int ds4_gpu_hc_expand_add_split_half_add_tensor(
ds4_gpu_tensor *out_hc,
const ds4_gpu_tensor *block_out,
const ds4_gpu_tensor *block_add_h,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
(void)out_hc; (void)block_out; (void)block_add_h; (void)residual_hc;
(void)split; (void)n_embd; (void)n_hc;
return 0;
}
int ds4_gpu_shared_down_hc_expand_q8_0_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *shared_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *shared_mid,
const ds4_gpu_tensor *routed_out,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out_hc || !shared_out || !model_map || !shared_mid || !routed_out ||
!residual_hc || !split || n_embd == 0 || n_hc == 0 ||
n_hc != 4 || out_dim != n_embd || (in_dim & 31u) != 0 ||
in_dim > UINT32_MAX || out_dim > UINT32_MAX) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> midbuf = ds4_gpu_tensor_buffer(shared_mid);
id<MTLBuffer> sharedbuf = ds4_gpu_tensor_buffer(shared_out);
id<MTLBuffer> routedbuf = ds4_gpu_tensor_buffer(routed_out);
id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
const uint64_t row_bytes = (in_dim / 32u) * 34u;
const uint64_t weight_bytes = out_dim * row_bytes;
const uint64_t shared_mid_bytes = in_dim * sizeof(float);
const uint64_t embd_bytes = out_dim * sizeof(float);
const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t split_bytes = mix_hc * sizeof(float);
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal shared-down HC fusion weight range is outside the mapped model\n");
return 0;
}
if (!midbuf || !sharedbuf || !routedbuf || !resbuf || !splitbuf || !outbuf ||
ds4_gpu_tensor_bytes(shared_mid) < shared_mid_bytes ||
ds4_gpu_tensor_bytes(shared_out) < embd_bytes ||
ds4_gpu_tensor_bytes(routed_out) < embd_bytes ||
ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
ds4_gpu_tensor_bytes(split) < split_bytes ||
ds4_gpu_tensor_bytes(out_hc) < hc_bytes) {
fprintf(stderr, "ds4: Metal shared-down HC fusion received undersized buffers\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
weight_offset, weight_bytes,
&inner_offset);
if (!wbuf) return 0;
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
mv_args.nr0 = mv_dispatch.nr0;
ds4_gpu_hc_expand_args hc_args = {
.n_embd = n_embd,
.n_hc = n_hc,
.n_tokens = 1,
.nb_block0 = sizeof(float),
.nb_block1 = (uint64_t)n_embd * sizeof(float),
.nb_add0 = sizeof(float),
.nb_add1 = (uint64_t)n_embd * sizeof(float),
.nb_res0 = sizeof(float),
.nb_res1 = (uint64_t)n_embd * sizeof(float),
.nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
.nb_post0 = sizeof(float),
.nb_post1 = mix_hc * sizeof(float),
.nb_comb0 = sizeof(float),
.nb_comb1 = (uint64_t)n_hc * sizeof(float),
.nb_comb2 = mix_hc * sizeof(float),
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_embd * sizeof(float),
.nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
.has_add = 1,
};
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_shared_down_hc_expand4_q8_0",
mv_dispatch.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2];
[enc setBuffer:midbuf offset:ds4_gpu_tensor_offset(shared_mid) atIndex:3];
[enc setBuffer:sharedbuf offset:ds4_gpu_tensor_offset(shared_out) atIndex:4];
[enc setBuffer:routedbuf offset:ds4_gpu_tensor_offset(routed_out) atIndex:5];
[enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:6];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)n_hc * sizeof(float) atIndex:7];
[enc setBuffer:splitbuf offset:ds4_gpu_tensor_offset(split) + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:8];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:9];
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
(NSUInteger)mv_dispatch.nr0,
1,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "shared-down HC expand fused")) return 0;
}
return 1;
}
int ds4_gpu_matmul_q8_0_hc_expand_tensor(
ds4_gpu_tensor *out_hc,
ds4_gpu_tensor *block_out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t in_dim,
uint64_t out_dim,
const ds4_gpu_tensor *x,
const ds4_gpu_tensor *residual_hc,
const ds4_gpu_tensor *split,
uint32_t n_embd,
uint32_t n_hc) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out_hc || !block_out || !model_map || !x || !residual_hc || !split ||
n_embd == 0 || n_hc == 0 || n_hc != 4 || out_dim != n_embd ||
(in_dim & 31u) != 0 || in_dim > UINT32_MAX || out_dim > UINT32_MAX) {
return 0;
}
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> blockbuf = ds4_gpu_tensor_buffer(block_out);
id<MTLBuffer> resbuf = ds4_gpu_tensor_buffer(residual_hc);
id<MTLBuffer> splitbuf = ds4_gpu_tensor_buffer(split);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out_hc);
const uint64_t row_bytes = (in_dim / 32u) * 34u;
const uint64_t weight_bytes = out_dim * row_bytes;
const uint64_t x_bytes = in_dim * sizeof(float);
const uint64_t embd_bytes = out_dim * sizeof(float);
const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
const uint64_t split_bytes = mix_hc * sizeof(float);
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) {
fprintf(stderr, "ds4: Metal Q8 HC fusion weight range is outside the mapped model\n");
return 0;
}
if (!xbuf || !blockbuf || !resbuf || !splitbuf || !outbuf ||
ds4_gpu_tensor_bytes(x) < x_bytes ||
ds4_gpu_tensor_bytes(block_out) < embd_bytes ||
ds4_gpu_tensor_bytes(residual_hc) < hc_bytes ||
ds4_gpu_tensor_bytes(split) < split_bytes ||
ds4_gpu_tensor_bytes(out_hc) < hc_bytes) {
fprintf(stderr, "ds4: Metal Q8 HC fusion received undersized buffers\n");
return 0;
}
uint64_t inner_offset = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size,
weight_offset, weight_bytes,
&inner_offset);
if (!wbuf) return 0;
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
mv_args.nr0 = mv_dispatch.nr0;
ds4_gpu_hc_expand_args hc_args = {
.n_embd = n_embd,
.n_hc = n_hc,
.n_tokens = 1,
.nb_block0 = sizeof(float),
.nb_block1 = (uint64_t)n_embd * sizeof(float),
.nb_add0 = sizeof(float),
.nb_add1 = (uint64_t)n_embd * sizeof(float),
.nb_res0 = sizeof(float),
.nb_res1 = (uint64_t)n_embd * sizeof(float),
.nb_res2 = (uint64_t)n_hc * n_embd * sizeof(float),
.nb_post0 = sizeof(float),
.nb_post1 = mix_hc * sizeof(float),
.nb_comb0 = sizeof(float),
.nb_comb1 = (uint64_t)n_hc * sizeof(float),
.nb_comb2 = mix_hc * sizeof(float),
.nb0 = sizeof(float),
.nb1 = (uint64_t)n_embd * sizeof(float),
.nb2 = (uint64_t)n_hc * n_embd * sizeof(float),
.has_add = 0,
};
const NSUInteger split_offset = ds4_gpu_tensor_offset(split);
const bool vec_hc =
(split_offset & 15u) == 0u &&
getenv("DS4_METAL_DISABLE_M5_Q8_HC_VEC") == NULL &&
ds4_gpu_device_is_m5_apple_silicon();
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(
vec_hc ? "kernel_dsv4_q8_hc_expand4_q8_0_vec_hc" :
"kernel_dsv4_q8_hc_expand4_q8_0",
mv_dispatch.nsg);
if (!pipeline && vec_hc) {
pipeline = ds4_gpu_get_mul_mv_pipeline(
"kernel_dsv4_q8_hc_expand4_q8_0", mv_dispatch.nsg);
}
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBytes:&hc_args length:sizeof(hc_args) atIndex:1];
[enc setBuffer:wbuf offset:(NSUInteger)inner_offset atIndex:2];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:3];
[enc setBuffer:blockbuf offset:ds4_gpu_tensor_offset(block_out) atIndex:4];
[enc setBuffer:resbuf offset:ds4_gpu_tensor_offset(residual_hc) atIndex:5];
[enc setBuffer:splitbuf offset:split_offset + (NSUInteger)n_hc * sizeof(float) atIndex:6];
[enc setBuffer:splitbuf offset:split_offset + (NSUInteger)(2u * n_hc) * sizeof(float) atIndex:7];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out_hc) atIndex:8];
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) /
(NSUInteger)mv_dispatch.nr0,
1,
1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "Q8 HC expand fused")) return 0;
}
return 1;
}
/* Kimi Delta Attention primitive shared with the GLM-5.3 graph. */
static int glm53_gpu_mul_u64(uint64_t a, uint64_t b, uint64_t *out) {
if (!out || (a != 0 && b > UINT64_MAX / a)) return 0;
*out = a * b;
return 1;
}
static int glm53_gpu_tensor_has(
const ds4_gpu_tensor *tensor,
uint64_t elements,
uint64_t element_size) {
uint64_t bytes = 0;
return tensor &&
glm53_gpu_mul_u64(elements, element_size, &bytes) &&
ds4_gpu_tensor_buffer(tensor) != nil &&
ds4_gpu_tensor_bytes(tensor) >= bytes;
}
static id<MTLBuffer> glm53_gpu_weight_buffer(
const void *model_map,
uint64_t model_size,
uint64_t offset,
uint64_t bytes,
uint64_t *inner_offset,
const char *label) {
if (!model_map || offset > model_size || bytes > model_size - offset) {
fprintf(stderr, "ds4: GLM-5.3 %s range is outside the mapped model\n",
label ? label : "weight");
return nil;
}
return ds4_gpu_wrap_model_range(model_map, model_size, offset, bytes,
inner_offset);
}
typedef struct {
uint32_t in_dim;
uint32_t out_dim;
uint32_t n_rows;
} glm53_gpu_bf16_matmul_args;
int ds4_gpu_glm53_embedding_bf16(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
const ds4_gpu_tensor *token_ids,
uint32_t n_tokens,
uint32_t n_embd,
uint32_t n_vocab) {
if (!g_initialized && !ds4_gpu_init()) return 0;
uint64_t weights = 0, output = 0;
if (n_tokens == 0 || n_embd == 0 || n_vocab == 0 ||
!glm53_gpu_mul_u64(n_embd, n_vocab, &weights) ||
!glm53_gpu_mul_u64(n_embd, n_tokens, &output) ||
!glm53_gpu_tensor_has(token_ids, n_tokens, sizeof(int32_t)) ||
!glm53_gpu_tensor_has(out, output, sizeof(float))) {
fprintf(stderr, "ds4: GLM-5.3 BF16 embedding received invalid buffers\n");
return 0;
}
@autoreleasepool {
uint64_t inner = 0;
id<MTLBuffer> weightbuf = glm53_gpu_weight_buffer(
model_map, model_size, weight_offset, weights * sizeof(uint16_t),
&inner, "BF16 embedding");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_embedding_bf16");
if (!weightbuf || !pipeline) return 0;
glm53_gpu_bf16_matmul_args args = {
.in_dim = n_embd,
.out_dim = n_vocab,
.n_rows = n_tokens,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weightbuf offset:(NSUInteger)inner atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(token_ids)
offset:ds4_gpu_tensor_offset(token_ids) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreads:MTLSizeMake(n_embd, n_tokens, 1)
threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)n_embd, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 BF16 embedding");
}
}
static uint32_t glm53_gpu_bf16_mv_nsg(void) {
return ds4_gpu_device_name_contains("M3 Ultra") &&
getenv("DS4_METAL_DISABLE_M3_ULTRA_GLM53_DECODE") == NULL &&
getenv("DS4_METAL_DISABLE_M3_ULTRA_GLM53_BF16_NSG4") == NULL
? 4u : 8u;
}
int ds4_gpu_glm53_matmul_bf16(
ds4_gpu_tensor *out,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x,
uint32_t n_rows) {
if (!g_initialized && !ds4_gpu_init()) return 0;
uint64_t weights = 0, input = 0, output = 0;
if (in_dim == 0 || out_dim == 0 || n_rows == 0 ||
!glm53_gpu_mul_u64(in_dim, out_dim, &weights) ||
!glm53_gpu_mul_u64(in_dim, n_rows, &input) ||
!glm53_gpu_mul_u64(out_dim, n_rows, &output) ||
!glm53_gpu_tensor_has(x, input, sizeof(float)) ||
!glm53_gpu_tensor_has(out, output, sizeof(float))) {
fprintf(stderr, "ds4: GLM-5.3 BF16 matmul received invalid buffers\n");
return 0;
}
@autoreleasepool {
uint64_t inner = 0;
id<MTLBuffer> weightbuf = glm53_gpu_weight_buffer(
model_map, model_size, weight_offset, weights * sizeof(uint16_t),
&inner, "BF16 matrix");
if (!weightbuf) return 0;
const bool use_mv = n_rows <= 8u;
const bool bc_inp = (in_dim % 32u) != 0u;
const bool bc_out = (out_dim % 64u) != 0u || (n_rows % 32u) != 0u;
id<MTLComputePipelineState> pipeline = use_mv
? ds4_gpu_get_pipeline("kernel_glm53_mul_mv_bf16_f32")
: ds4_gpu_get_mul_mm_pipeline(
"kernel_glm53_mul_mm_bf16_f32", bc_inp, bc_out);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
if (use_mv) {
const uint32_t nsg = glm53_gpu_bf16_mv_nsg();
glm53_gpu_bf16_matmul_args args = {
.in_dim = in_dim,
.out_dim = out_dim,
.n_rows = n_rows,
};
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weightbuf offset:(NSUInteger)inner atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreadgroups:MTLSizeMake((out_dim + nsg - 1u) / nsg,
n_rows, 1)
threadsPerThreadgroup:MTLSizeMake(32u * nsg, 1, 1)];
} else {
ds4_gpu_mul_mm_args args = ds4_gpu_make_mm_args(
in_dim, out_dim, n_rows, (uint64_t)in_dim * sizeof(uint16_t));
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weightbuf offset:(NSUInteger)inner atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:(bc_out ? 8192u : 6144u) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake((n_rows + 31u) / 32u,
(out_dim + 63u) / 64u, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 BF16 matmul");
}
}
int ds4_gpu_glm53_matmul_bf16_qkv(
ds4_gpu_tensor *out_q,
ds4_gpu_tensor *out_k,
ds4_gpu_tensor *out_v,
const void *model_map,
uint64_t model_size,
uint64_t weight_q_offset,
uint64_t weight_k_offset,
uint64_t weight_v_offset,
uint32_t in_dim,
uint32_t out_dim,
const ds4_gpu_tensor *x) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!ds4_gpu_device_name_contains("M3 Ultra")) return 0;
uint64_t weights = 0;
if (in_dim == 0 || out_dim == 0 ||
!glm53_gpu_mul_u64(in_dim, out_dim, &weights) ||
!glm53_gpu_tensor_has(x, in_dim, sizeof(float)) ||
!glm53_gpu_tensor_has(out_q, out_dim, sizeof(float)) ||
!glm53_gpu_tensor_has(out_k, out_dim, sizeof(float)) ||
!glm53_gpu_tensor_has(out_v, out_dim, sizeof(float))) {
return 0;
}
@autoreleasepool {
const uint64_t weight_bytes = weights * sizeof(uint16_t);
uint64_t inner_q = 0, inner_k = 0, inner_v = 0;
id<MTLBuffer> weight_q = glm53_gpu_weight_buffer(
model_map, model_size, weight_q_offset, weight_bytes,
&inner_q, "BF16 Q matrix");
id<MTLBuffer> weight_k = glm53_gpu_weight_buffer(
model_map, model_size, weight_k_offset, weight_bytes,
&inner_k, "BF16 K matrix");
id<MTLBuffer> weight_v = glm53_gpu_weight_buffer(
model_map, model_size, weight_v_offset, weight_bytes,
&inner_v, "BF16 V matrix");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_mul_mv_bf16_f32_qkv");
if (!weight_q || !weight_k || !weight_v || !pipeline) return 0;
const uint32_t nsg = glm53_gpu_bf16_mv_nsg();
glm53_gpu_bf16_matmul_args args = {
.in_dim = in_dim,
.out_dim = out_dim,
.n_rows = 1u,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:weight_q offset:(NSUInteger)inner_q atIndex:1];
[enc setBuffer:weight_k offset:(NSUInteger)inner_k atIndex:2];
[enc setBuffer:weight_v offset:(NSUInteger)inner_v atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(out_q)
offset:ds4_gpu_tensor_offset(out_q) atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(out_k)
offset:ds4_gpu_tensor_offset(out_k) atIndex:6];
[enc setBuffer:ds4_gpu_tensor_buffer(out_v)
offset:ds4_gpu_tensor_offset(out_v) atIndex:7];
[enc dispatchThreadgroups:MTLSizeMake((out_dim + nsg - 1u) / nsg,
1u, 3u)
threadsPerThreadgroup:MTLSizeMake(32u * nsg, 1u, 1u)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned,
"GLM-5.3 BF16 QKV matmul");
}
}
typedef struct {
uint32_t width;
uint32_t rows;
float eps;
} glm53_vision_rows_args;
typedef struct {
uint32_t rows;
uint32_t grid_h;
uint32_t grid_w;
float eps;
} glm53_vision_qkv_args;
typedef struct {
uint32_t rows;
float scale;
} glm53_vision_attention_args;
typedef struct {
uint32_t dst_row;
uint32_t image_row;
uint32_t rows;
uint32_t total_rows;
uint32_t width;
uint32_t hc;
} glm53_vision_scatter_args;
static int glm53_vision_dispatch_rows(
const char *kernel,
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint32_t width,
uint32_t rows,
float eps,
const char *label) {
uint64_t inner = 0;
id<MTLBuffer> weight = glm53_gpu_weight_buffer(
model_map, model_size, weight_offset,
(uint64_t)width * sizeof(uint16_t), &inner, label);
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(kernel);
if (!weight || !pipeline) return 0;
glm53_vision_rows_args args = { width, rows, eps };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:weight offset:(NSUInteger)inner atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:8u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, label);
}
static int glm53_vision_dispatch_bias(
const char *kernel,
ds4_gpu_tensor *x,
const ds4_gpu_tensor *residual,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint32_t width,
uint32_t rows,
const char *label) {
uint64_t inner = 0;
id<MTLBuffer> bias = glm53_gpu_weight_buffer(
model_map, model_size, bias_offset,
(uint64_t)width * sizeof(uint16_t), &inner, label);
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(kernel);
if (!bias || !pipeline) return 0;
glm53_vision_rows_args args = { width, rows, 0.0f };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:bias offset:(NSUInteger)inner atIndex:2];
if (residual) {
[enc setBuffer:ds4_gpu_tensor_buffer(residual)
offset:ds4_gpu_tensor_offset(residual) atIndex:3];
}
[enc dispatchThreads:MTLSizeMake(width, rows, 1)
threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)width, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, label);
}
static int glm53_vision_dispatch_qkv(
ds4_gpu_tensor *q,
ds4_gpu_tensor *k,
ds4_gpu_tensor *v,
const ds4_gpu_tensor *qkv,
const void *model_map,
uint64_t model_size,
const ds4_glm53_vision_layer_weights *w,
uint32_t rows,
uint32_t grid_h,
uint32_t grid_w) {
uint64_t bias_inner = 0, q_inner = 0, k_inner = 0;
id<MTLBuffer> bias = glm53_gpu_weight_buffer(
model_map, model_size, w->qkv_bias, 3072u * sizeof(uint16_t),
&bias_inner, "vision QKV bias");
id<MTLBuffer> qw = glm53_gpu_weight_buffer(
model_map, model_size, w->q_norm, 64u * sizeof(uint16_t),
&q_inner, "vision Q norm");
id<MTLBuffer> kw = glm53_gpu_weight_buffer(
model_map, model_size, w->k_norm, 64u * sizeof(uint16_t),
&k_inner, "vision K norm");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_vision_qkv_rope");
if (!bias || !qw || !kw || !pipeline) return 0;
glm53_vision_qkv_args args = { rows, grid_h, grid_w, 1.0e-5f };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(qkv)
offset:ds4_gpu_tensor_offset(qkv) atIndex:1];
[enc setBuffer:bias offset:(NSUInteger)bias_inner atIndex:2];
[enc setBuffer:qw offset:(NSUInteger)q_inner atIndex:3];
[enc setBuffer:kw offset:(NSUInteger)k_inner atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(q)
offset:ds4_gpu_tensor_offset(q) atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(k)
offset:ds4_gpu_tensor_offset(k) atIndex:6];
[enc setBuffer:ds4_gpu_tensor_buffer(v)
offset:ds4_gpu_tensor_offset(v) atIndex:7];
[enc dispatchThreadgroups:MTLSizeMake(rows, 16u, 1)
threadsPerThreadgroup:MTLSizeMake(32u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision QKV");
}
static int glm53_vision_dispatch_attention(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *k,
const ds4_gpu_tensor *v,
uint32_t rows) {
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_vision_attention");
if (!pipeline) return 0;
glm53_vision_attention_args args = { rows, 0.125f };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(q)
offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(k)
offset:ds4_gpu_tensor_offset(k) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(v)
offset:ds4_gpu_tensor_offset(v) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:4];
[enc dispatchThreadgroups:MTLSizeMake(rows, 16u, 1)
threadsPerThreadgroup:MTLSizeMake(32u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision attention");
}
static int glm53_vision_dispatch_swiglu_bias(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *gate,
const ds4_gpu_tensor *up,
const void *model_map,
uint64_t model_size,
uint64_t gate_bias_offset,
uint64_t up_bias_offset,
uint32_t width,
uint32_t rows) {
uint64_t gate_inner = 0, up_inner = 0;
id<MTLBuffer> gate_bias = glm53_gpu_weight_buffer(
model_map, model_size, gate_bias_offset,
(uint64_t)width * sizeof(uint16_t), &gate_inner,
"vision gate bias");
id<MTLBuffer> up_bias = glm53_gpu_weight_buffer(
model_map, model_size, up_bias_offset,
(uint64_t)width * sizeof(uint16_t), &up_inner,
"vision up bias");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_vision_swiglu_bias");
if (!gate_bias || !up_bias || !pipeline) return 0;
glm53_vision_rows_args args = { width, rows, 0.0f };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(gate)
offset:ds4_gpu_tensor_offset(gate) atIndex:1];
[enc setBuffer:gate_bias offset:(NSUInteger)gate_inner atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(up)
offset:ds4_gpu_tensor_offset(up) atIndex:3];
[enc setBuffer:up_bias offset:(NSUInteger)up_inner atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:5];
[enc dispatchThreads:MTLSizeMake(width, rows, 1)
threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)width, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision SwiGLU");
}
static int glm53_vision_dispatch_reorder(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
uint32_t rows) {
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_vision_downsample_reorder");
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:1];
[enc dispatchThreads:MTLSizeMake(4096u, rows, 1)
threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision downsample reorder");
}
static int glm53_vision_dispatch_layernorm_gelu(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t weight_offset,
uint64_t bias_offset,
uint32_t rows) {
uint64_t weight_inner = 0, bias_inner = 0;
id<MTLBuffer> weight = glm53_gpu_weight_buffer(
model_map, model_size, weight_offset, 4096u * sizeof(uint16_t),
&weight_inner, "vision merger norm");
id<MTLBuffer> bias = glm53_gpu_weight_buffer(
model_map, model_size, bias_offset, 4096u * sizeof(uint16_t),
&bias_inner, "vision merger norm bias");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_vision_layernorm_gelu");
if (!weight || !bias || !pipeline) return 0;
glm53_vision_rows_args args = { 4096u, rows, 1.0e-5f };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:weight offset:(NSUInteger)weight_inner atIndex:2];
[enc setBuffer:bias offset:(NSUInteger)bias_inner atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:4];
[enc setThreadgroupMemoryLength:8u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision merger norm");
}
int ds4_gpu_glm53_vision_encode(
float *out,
const float *patches,
uint32_t grid_h,
uint32_t grid_w,
const void *model_map,
uint64_t model_size,
const ds4_glm53_vision_weights *weights) {
if (!out || !patches || !model_map || !weights || grid_h == 0 || grid_w == 0 ||
(grid_h & 1u) != 0 || (grid_w & 1u) != 0 ||
grid_h > UINT32_MAX / grid_w || ds4_gpu_commands_active()) {
fprintf(stderr, "ds4: invalid GLM-5.3 vision encoder input\n");
return 0;
}
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t rows = grid_h * grid_w;
const uint32_t merged_rows = rows / 4u;
const uint64_t row1024 = (uint64_t)rows * 1024u;
const uint64_t row3072 = (uint64_t)rows * 3072u;
const uint64_t row4096 = (uint64_t)rows * 4096u;
const uint64_t merged10240 = (uint64_t)merged_rows * 10240u;
if (row4096 > SIZE_MAX / sizeof(float) ||
merged10240 > SIZE_MAX / sizeof(float)) return 0;
ds4_gpu_tensor *patch = NULL, *a = NULL, *b = NULL, *qkv = NULL;
ds4_gpu_tensor *q = NULL, *k = NULL, *v = NULL, *attn = NULL;
ds4_gpu_tensor *gate = NULL, *up = NULL, *mid = NULL;
int ok = 0;
#define VISION_ALLOC(name_, count_) do { \
name_ = ds4_gpu_tensor_alloc((count_) * sizeof(float)); \
if (!(name_)) goto cleanup; \
} while (0)
VISION_ALLOC(patch, (uint64_t)rows * 1176u);
VISION_ALLOC(a, row1024);
VISION_ALLOC(b, row1024);
VISION_ALLOC(qkv, row3072);
VISION_ALLOC(q, row1024);
VISION_ALLOC(k, row1024);
VISION_ALLOC(v, row1024);
VISION_ALLOC(attn, row1024);
VISION_ALLOC(gate, row4096);
VISION_ALLOC(up, row4096);
VISION_ALLOC(mid, row4096);
#undef VISION_ALLOC
if (!ds4_gpu_tensor_write(patch, 0, patches,
(uint64_t)rows * 1176u * sizeof(float))) goto cleanup;
if (!ds4_gpu_begin_commands()) goto cleanup;
ok = ds4_gpu_glm53_matmul_bf16(a, model_map, model_size,
weights->patch_weight,
1176u, 1024u, patch, rows);
if (ok) ok = glm53_vision_dispatch_bias(
"kernel_glm53_vision_add_bias", a, NULL, model_map, model_size,
weights->patch_bias, 1024u, rows, "vision patch bias");
ds4_gpu_tensor *cur = a;
ds4_gpu_tensor *tmp = b;
for (uint32_t il = 0; ok && il < DS4_GLM53_VISION_LAYERS; il++) {
const ds4_glm53_vision_layer_weights *w = &weights->layer[il];
ok = glm53_vision_dispatch_rows(
"kernel_glm53_vision_rms_bf16", tmp, cur,
model_map, model_size, w->norm1, 1024u, rows, 1.0e-5f,
"vision norm1");
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
qkv, model_map, model_size, w->qkv_weight,
1024u, 3072u, tmp, rows);
if (ok) ok = glm53_vision_dispatch_qkv(
q, k, v, qkv, model_map, model_size, w, rows, grid_h, grid_w);
if (ok) ok = glm53_vision_dispatch_attention(attn, q, k, v, rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, w->attn_proj_weight,
1024u, 1024u, attn, rows);
if (ok) ok = glm53_vision_dispatch_bias(
"kernel_glm53_vision_bias_residual", tmp, cur,
model_map, model_size, w->attn_proj_bias,
1024u, rows, "vision attention bias");
ds4_gpu_tensor *swap = cur; cur = tmp; tmp = swap;
if (ok) ok = glm53_vision_dispatch_rows(
"kernel_glm53_vision_rms_bf16", tmp, cur,
model_map, model_size, w->norm2, 1024u, rows, 1.0e-5f,
"vision norm2");
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
gate, model_map, model_size, w->gate_weight,
1024u, 4096u, tmp, rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
up, model_map, model_size, w->up_weight,
1024u, 4096u, tmp, rows);
if (ok) ok = glm53_vision_dispatch_swiglu_bias(
mid, gate, up, model_map, model_size,
w->gate_bias, w->up_bias, 4096u, rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, w->down_weight,
4096u, 1024u, mid, rows);
if (ok) ok = glm53_vision_dispatch_bias(
"kernel_glm53_vision_bias_residual", tmp, cur,
model_map, model_size, w->down_bias,
1024u, rows, "vision down bias");
swap = cur; cur = tmp; tmp = swap;
}
if (ok) ok = glm53_vision_dispatch_rows(
"kernel_glm53_vision_rms_bf16", tmp, cur,
model_map, model_size, weights->post_norm,
1024u, rows, 1.0e-5f, "vision post norm");
if (ok) ok = glm53_vision_dispatch_reorder(cur, tmp, merged_rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, weights->downsample_weight,
4096u, 4096u, cur, merged_rows);
if (ok) ok = glm53_vision_dispatch_bias(
"kernel_glm53_vision_add_bias", tmp, NULL,
model_map, model_size, weights->downsample_bias,
4096u, merged_rows, "vision downsample bias");
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
cur, model_map, model_size, weights->merger_proj,
4096u, 4096u, tmp, merged_rows);
if (ok) ok = glm53_vision_dispatch_layernorm_gelu(
tmp, cur, model_map, model_size,
weights->merger_norm, weights->merger_norm_bias, merged_rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
gate, model_map, model_size, weights->merger_gate,
4096u, 10240u, tmp, merged_rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
up, model_map, model_size, weights->merger_up,
4096u, 10240u, tmp, merged_rows);
if (ok) ok = ds4_gpu_swiglu_tensor(
mid, gate, up, merged_rows * 10240u, 10.0f, 1.0f);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
cur, model_map, model_size, weights->merger_down,
10240u, 4096u, mid, merged_rows);
if (ok) ok = ds4_gpu_end_commands();
else (void)ds4_gpu_end_commands();
if (ok) ok = ds4_gpu_tensor_read(
cur, 0, out, (uint64_t)merged_rows * 4096u * sizeof(float));
cleanup:
ds4_gpu_tensor_free(mid);
ds4_gpu_tensor_free(up);
ds4_gpu_tensor_free(gate);
ds4_gpu_tensor_free(attn);
ds4_gpu_tensor_free(v);
ds4_gpu_tensor_free(k);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(qkv);
ds4_gpu_tensor_free(b);
ds4_gpu_tensor_free(a);
ds4_gpu_tensor_free(patch);
return ok;
}
typedef struct {
uint32_t width;
uint32_t rows;
} deepseek4_vision_rows_args;
typedef struct {
uint32_t rows;
uint32_t grid_width;
} deepseek4_vision_qkv_args;
typedef struct {
uint32_t grid_height;
uint32_t grid_width;
uint32_t output_rows;
} deepseek4_vision_align_args;
static int deepseek4_vision_dispatch_qkv(
ds4_gpu_tensor *q,
ds4_gpu_tensor *k,
ds4_gpu_tensor *v,
const ds4_gpu_tensor *qkv,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint32_t rows,
uint32_t grid_width) {
uint64_t bias_inner = 0;
id<MTLBuffer> bias = glm53_gpu_weight_buffer(
model_map, model_size, bias_offset, 3072u * sizeof(uint16_t),
&bias_inner, "DeepSeek vision QKV bias");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_deepseek4_vision_qkv_rope");
if (!bias || !pipeline) return 0;
deepseek4_vision_qkv_args args = { rows, grid_width };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(qkv)
offset:ds4_gpu_tensor_offset(qkv) atIndex:1];
[enc setBuffer:bias offset:(NSUInteger)bias_inner atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(q)
offset:ds4_gpu_tensor_offset(q) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(k)
offset:ds4_gpu_tensor_offset(k) atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(v)
offset:ds4_gpu_tensor_offset(v) atIndex:5];
[enc dispatchThreadgroups:MTLSizeMake(rows, 16u, 1)
threadsPerThreadgroup:MTLSizeMake(32u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned,
"DeepSeek vision QKV");
}
static int deepseek4_vision_dispatch_two_rows(
const char *kernel,
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
uint32_t width,
uint32_t rows,
const char *label) {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(kernel);
if (!pipeline) return 0;
deepseek4_vision_rows_args args = { width, rows };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc dispatchThreads:MTLSizeMake(width, rows, 1)
threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)width, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, label);
}
static int deepseek4_vision_dispatch_round(
ds4_gpu_tensor *x,
uint32_t width,
uint32_t rows) {
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_deepseek4_vision_round_bf16");
if (!pipeline) return 0;
deepseek4_vision_rows_args args = { width, rows };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc dispatchThreads:MTLSizeMake(width, rows, 1)
threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)width, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned,
"DeepSeek vision BF16 round");
}
static int deepseek4_vision_dispatch_aligner_reorder(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
uint32_t grid_height,
uint32_t grid_width,
uint32_t output_rows) {
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_deepseek4_vision_aligner_reorder");
if (!pipeline) return 0;
deepseek4_vision_align_args args = {
grid_height, grid_width, output_rows
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:2];
[enc dispatchThreads:MTLSizeMake(9216u, output_rows, 1)
threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned,
"DeepSeek vision aligner reorder");
}
static int deepseek4_vision_dispatch_gelu_bias(
ds4_gpu_tensor *out,
const ds4_gpu_tensor *x,
const void *model_map,
uint64_t model_size,
uint64_t bias_offset,
uint32_t rows,
uint32_t width) {
uint64_t bias_inner = 0;
id<MTLBuffer> bias = glm53_gpu_weight_buffer(
model_map, model_size, bias_offset, (uint64_t)width * sizeof(uint16_t),
&bias_inner, "DeepSeek vision aligner bias");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_deepseek4_vision_gelu_bias");
if (!bias || !pipeline) return 0;
deepseek4_vision_rows_args args = { width, rows };
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x)
offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:bias offset:(NSUInteger)bias_inner atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc dispatchThreads:MTLSizeMake(width, rows, 1)
threadsPerThreadgroup:MTLSizeMake(256u, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned,
"DeepSeek vision aligner GELU");
}
static int deepseek4_vision_debug_dump(
const char *stage,
const ds4_gpu_tensor *tensor,
uint64_t values) {
const char *prefix = getenv("DS4_DEEPSEEK4_VISION_DEBUG_PREFIX");
if (!prefix || !prefix[0]) return 1;
if (!stage || !tensor || values > SIZE_MAX / sizeof(float)) return 0;
const bool resume = ds4_gpu_commands_active();
if (resume && !ds4_gpu_end_commands()) return 0;
float *data = malloc((size_t)values * sizeof(float));
if (!data) return 0;
int ok = ds4_gpu_tensor_read(
tensor, 0, data, values * sizeof(float)) != 0;
char path[PATH_MAX];
int length = snprintf(path, sizeof(path), "%s.%s.f32", prefix, stage);
if (ok && (length < 0 || (size_t)length >= sizeof(path))) ok = 0;
FILE *fp = ok ? fopen(path, "wb") : NULL;
if (ok && (!fp || fwrite(data, sizeof(float), (size_t)values, fp) != values))
ok = 0;
if (fp && fclose(fp) != 0) ok = 0;
free(data);
if (resume && !ds4_gpu_begin_commands()) ok = 0;
return ok;
}
int ds4_gpu_deepseek4_vision_encode(
float *out,
const float *patches,
uint32_t grid_h,
uint32_t grid_w,
const void *model_map,
uint64_t model_size,
const ds4_deepseek4_vision_weights *weights) {
if (!out || !patches || !model_map || !weights || grid_h == 0u ||
grid_w == 0u || grid_h > UINT32_MAX / grid_w ||
ds4_gpu_commands_active()) {
fprintf(stderr, "ds4: invalid DeepSeek vision encoder input\n");
return 0;
}
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t projection = weights->projection_dim ? weights->projection_dim : 4096u;
if (projection != 4096u && projection != 5120u) return 0;
const uint32_t rows = grid_h * grid_w;
const uint32_t aligned_rows =
((grid_h + 2u) / 3u) * ((grid_w + 2u) / 3u);
const uint64_t row1024 = (uint64_t)rows * 1024u;
const uint64_t row2816 = (uint64_t)rows * 2816u;
const uint64_t row3072 = (uint64_t)rows * 3072u;
const uint64_t row5632 = (uint64_t)rows * 5632u;
const uint64_t aligned_projection = (uint64_t)aligned_rows * projection;
const uint64_t aligned9216 = (uint64_t)aligned_rows * 9216u;
if (row5632 > SIZE_MAX / sizeof(float) ||
aligned9216 > SIZE_MAX / sizeof(float)) return 0;
ds4_gpu_tensor *patch = NULL, *a = NULL, *b = NULL, *qkv = NULL;
ds4_gpu_tensor *q = NULL, *k = NULL, *v = NULL, *attn = NULL;
ds4_gpu_tensor *mlp_w1 = NULL, *mlp_mid = NULL;
ds4_gpu_tensor *align_in = NULL, *align_a = NULL, *align_b = NULL;
int ok = 0;
#define DEEPSEEK4_VISION_ALLOC(name_, count_) do { \
name_ = ds4_gpu_tensor_alloc((count_) * sizeof(float)); \
if (!(name_)) goto cleanup; \
} while (0)
DEEPSEEK4_VISION_ALLOC(patch, (uint64_t)rows * 588u);
DEEPSEEK4_VISION_ALLOC(a, row1024);
DEEPSEEK4_VISION_ALLOC(b, row1024);
DEEPSEEK4_VISION_ALLOC(qkv, row3072);
DEEPSEEK4_VISION_ALLOC(q, row1024);
DEEPSEEK4_VISION_ALLOC(k, row1024);
DEEPSEEK4_VISION_ALLOC(v, row1024);
DEEPSEEK4_VISION_ALLOC(attn, row1024);
DEEPSEEK4_VISION_ALLOC(mlp_w1, row5632);
DEEPSEEK4_VISION_ALLOC(mlp_mid, row2816);
DEEPSEEK4_VISION_ALLOC(align_in, aligned9216);
DEEPSEEK4_VISION_ALLOC(align_a, aligned_projection);
DEEPSEEK4_VISION_ALLOC(align_b, aligned_projection);
#undef DEEPSEEK4_VISION_ALLOC
if (!ds4_gpu_tensor_write(patch, 0, patches,
(uint64_t)rows * 588u * sizeof(float)) ||
!ds4_gpu_begin_commands()) goto cleanup;
ok = projection != 5120u || deepseek4_vision_dispatch_round(patch, 588u, rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
a, model_map, model_size, weights->patch_weight,
588u, 1024u, patch, rows);
if (ok) ok = glm53_vision_dispatch_bias(
"kernel_glm53_vision_add_bias", a, NULL,
model_map, model_size, weights->patch_bias,
1024u, rows, "DeepSeek vision patch bias");
if (ok) ok = deepseek4_vision_dispatch_round(a, 1024u, rows);
if (ok) ok = deepseek4_vision_debug_dump(
"patch_embed", a, row1024);
ds4_gpu_tensor *cur = a;
ds4_gpu_tensor *tmp = b;
for (uint32_t il = 0; ok && il < DS4_DEEPSEEK4_VISION_LAYERS; il++) {
const ds4_deepseek4_vision_layer_weights *w = &weights->layer[il];
ok = glm53_vision_dispatch_rows(
"kernel_glm53_vision_rms_bf16", tmp, cur,
model_map, model_size, w->norm1,
1024u, rows, 1.0e-6f, "DeepSeek vision norm1");
if (ok) ok = deepseek4_vision_dispatch_round(tmp, 1024u, rows);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_norm1", tmp, row1024);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
qkv, model_map, model_size, w->qkv_weight,
1024u, 3072u, tmp, rows);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_qkv", qkv, row3072);
if (ok) ok = deepseek4_vision_dispatch_qkv(
q, k, v, qkv, model_map, model_size,
w->qkv_bias, rows, grid_w);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_q", q, row1024);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_k", k, row1024);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_v", v, row1024);
if (ok) ok = glm53_vision_dispatch_attention(attn, q, k, v, rows);
if (ok) ok = deepseek4_vision_dispatch_round(attn, 1024u, rows);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_attention", attn, row1024);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, w->attn_proj_weight,
1024u, 1024u, attn, rows);
if (ok) ok = glm53_vision_dispatch_bias(
projection == 5120u ? "kernel_deepseek41_vision_bias_residual" :
"kernel_glm53_vision_bias_residual", tmp, cur,
model_map, model_size, w->attn_proj_bias,
1024u, rows, "DeepSeek vision attention bias");
if (ok) ok = deepseek4_vision_dispatch_round(tmp, 1024u, rows);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_attn_residual", tmp, row1024);
ds4_gpu_tensor *swap = cur; cur = tmp; tmp = swap;
if (ok) ok = glm53_vision_dispatch_rows(
"kernel_glm53_vision_rms_bf16", tmp, cur,
model_map, model_size, w->norm2,
1024u, rows, 1.0e-6f, "DeepSeek vision norm2");
if (ok) ok = deepseek4_vision_dispatch_round(tmp, 1024u, rows);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_norm2", tmp, row1024);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
mlp_w1, model_map, model_size, w->mlp_w1,
1024u, 5632u, tmp, rows);
if (ok) ok = deepseek4_vision_dispatch_round(mlp_w1, 5632u, rows);
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_mlp_w1", mlp_w1, row5632);
if (ok) ok = deepseek4_vision_dispatch_two_rows(
projection == 5120u ? "kernel_deepseek41_vision_swiglu_split" :
"kernel_deepseek4_vision_swiglu_split",
mlp_w1, mlp_mid, 2816u, rows,
"DeepSeek vision SwiGLU");
if (ok && il == 0u) ok = deepseek4_vision_debug_dump(
"block0_mlp_mid", mlp_mid, row2816);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
tmp, model_map, model_size, w->mlp_w2,
2816u, 1024u, mlp_mid, rows);
if (ok) ok = deepseek4_vision_dispatch_round(tmp, 1024u, rows);
if (ok) ok = deepseek4_vision_dispatch_two_rows(
"kernel_deepseek4_vision_add_residual",
tmp, cur, 1024u, rows,
"DeepSeek vision MLP residual");
swap = cur; cur = tmp; tmp = swap;
if (ok && (il == 0u || il == 15u || il == 31u)) {
char stage[24];
snprintf(stage, sizeof(stage), "block%u", il);
ok = deepseek4_vision_debug_dump(stage, cur, row1024);
}
}
if (ok) ok = glm53_vision_dispatch_rows(
"kernel_glm53_vision_rms_bf16", tmp, cur,
model_map, model_size, weights->post_norm,
1024u, rows, 1.0e-6f, "DeepSeek vision final norm");
if (ok) ok = deepseek4_vision_dispatch_round(tmp, 1024u, rows);
if (ok) ok = deepseek4_vision_debug_dump(
"vision_norm", tmp, row1024);
if (ok) ok = deepseek4_vision_dispatch_aligner_reorder(
align_in, tmp, grid_h, grid_w, aligned_rows);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
align_a, model_map, model_size, weights->aligner_w1,
9216u, projection, align_in, aligned_rows);
if (ok) ok = deepseek4_vision_dispatch_gelu_bias(
align_b, align_a, model_map, model_size,
weights->aligner_w1_bias, aligned_rows, projection);
if (ok) ok = ds4_gpu_glm53_matmul_bf16(
align_a, model_map, model_size, weights->aligner_w2,
projection, projection, align_b, aligned_rows);
if (ok) ok = glm53_vision_dispatch_bias(
"kernel_glm53_vision_add_bias", align_a, NULL,
model_map, model_size, weights->aligner_w2_bias,
projection, aligned_rows, "DeepSeek vision aligner output bias");
if (ok) ok = deepseek4_vision_dispatch_round(
align_a, projection, aligned_rows);
if (ok) ok = deepseek4_vision_debug_dump(
"aligned", align_a, aligned_projection);
if (ok) ok = ds4_gpu_end_commands();
else (void)ds4_gpu_end_commands();
if (ok) ok = ds4_gpu_tensor_read(
align_a, 0, out, aligned_projection * sizeof(float));
cleanup:
ds4_gpu_tensor_free(align_b);
ds4_gpu_tensor_free(align_a);
ds4_gpu_tensor_free(align_in);
ds4_gpu_tensor_free(mlp_mid);
ds4_gpu_tensor_free(mlp_w1);
ds4_gpu_tensor_free(attn);
ds4_gpu_tensor_free(v);
ds4_gpu_tensor_free(k);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(qkv);
ds4_gpu_tensor_free(b);
ds4_gpu_tensor_free(a);
ds4_gpu_tensor_free(patch);
return ok;
}
int ds4_gpu_glm53_scatter_image_hc(
ds4_gpu_tensor *hc,
const ds4_gpu_tensor *image,
uint32_t dst_row,
uint32_t image_row,
uint32_t rows,
uint32_t total_rows,
uint32_t n_embd,
uint32_t n_hc) {
if (!hc || !image || rows == 0 || n_embd == 0 || n_hc == 0 ||
dst_row > total_rows || rows > total_rows - dst_row ||
!ds4_gpu_commands_active()) return 0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_vision_scatter_hc");
if (!pipeline) return 0;
glm53_vision_scatter_args args = {
dst_row, image_row, rows, total_rows, n_embd, n_hc
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(hc)
offset:ds4_gpu_tensor_offset(hc) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(image)
offset:ds4_gpu_tensor_offset(image) atIndex:2];
[enc dispatchThreads:MTLSizeMake(n_embd, rows * n_hc, 1)
threadsPerThreadgroup:MTLSizeMake(MIN((NSUInteger)n_embd, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "GLM-5.3 vision scatter");
}
typedef struct {
uint32_t n_heads;
uint32_t n_rows;
float lower_bound;
float norm_eps;
} glm53_gpu_kda_args;
int ds4_gpu_glm53_kda_decode(
ds4_gpu_tensor *out,
ds4_gpu_tensor *conv_state,
ds4_gpu_tensor *recurrent_state,
const ds4_gpu_tensor *q,
const ds4_gpu_tensor *k,
const ds4_gpu_tensor *v,
const ds4_gpu_tensor *raw_gate,
const ds4_gpu_tensor *raw_beta,
const ds4_gpu_tensor *output_gate,
const void *model_map,
uint64_t model_size,
uint64_t q_conv_offset,
uint64_t k_conv_offset,
uint64_t v_conv_offset,
uint64_t a_log_offset,
uint64_t dt_bias_offset,
uint64_t output_norm_offset,
uint32_t n_heads,
uint32_t n_rows,
float gate_lower_bound,
float norm_eps) {
enum { GLM53_KDA_DIM = 128, GLM53_KDA_HISTORY = 3 };
if (!g_initialized && !ds4_gpu_init()) return 0;
uint64_t projection = 0, activation_elements = 0;
uint64_t conv_elements = 0, state_elements = 0;
if (n_heads == 0 || n_rows == 0 || gate_lower_bound >= 0.0f ||
!glm53_gpu_mul_u64(n_heads, GLM53_KDA_DIM, &projection) ||
!glm53_gpu_mul_u64(projection, n_rows, &activation_elements) ||
!glm53_gpu_mul_u64(activation_elements,
3u * GLM53_KDA_HISTORY, &conv_elements) ||
!glm53_gpu_mul_u64(activation_elements, GLM53_KDA_DIM, &state_elements) ||
!glm53_gpu_tensor_has(q, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(k, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(v, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(raw_gate, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(output_gate, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(raw_beta,
(uint64_t)n_rows * n_heads, sizeof(float)) ||
!glm53_gpu_tensor_has(out, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(conv_state, conv_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(recurrent_state, state_elements, sizeof(float))) {
fprintf(stderr, "ds4: GLM-5.3 KDA decode received invalid buffers\n");
return 0;
}
uint64_t conv_bytes = 0, dt_bytes = 0;
if (!glm53_gpu_mul_u64(projection, 4u * sizeof(float), &conv_bytes) ||
!glm53_gpu_mul_u64(projection, sizeof(float), &dt_bytes)) {
return 0;
}
const uint64_t a_log_bytes = (uint64_t)n_heads * sizeof(float);
const uint64_t norm_bytes = GLM53_KDA_DIM * sizeof(float);
@autoreleasepool {
uint64_t qw_inner = 0, kw_inner = 0, vw_inner = 0;
uint64_t a_inner = 0, dt_inner = 0, norm_inner = 0;
id<MTLBuffer> qw = glm53_gpu_weight_buffer(
model_map, model_size, q_conv_offset, conv_bytes,
&qw_inner, "KDA Q convolution");
id<MTLBuffer> kw = glm53_gpu_weight_buffer(
model_map, model_size, k_conv_offset, conv_bytes,
&kw_inner, "KDA K convolution");
id<MTLBuffer> vw = glm53_gpu_weight_buffer(
model_map, model_size, v_conv_offset, conv_bytes,
&vw_inner, "KDA V convolution");
id<MTLBuffer> a_log = glm53_gpu_weight_buffer(
model_map, model_size, a_log_offset, a_log_bytes,
&a_inner, "KDA A_log");
id<MTLBuffer> dt_bias = glm53_gpu_weight_buffer(
model_map, model_size, dt_bias_offset, dt_bytes,
&dt_inner, "KDA dt bias");
id<MTLBuffer> output_norm = glm53_gpu_weight_buffer(
model_map, model_size, output_norm_offset, norm_bytes,
&norm_inner, "KDA output norm");
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_pipeline("kernel_glm53_kda_decode");
if (!qw || !kw || !vw || !a_log || !dt_bias || !output_norm ||
!pipeline) {
return 0;
}
glm53_gpu_kda_args args = {
.n_heads = n_heads,
.n_rows = n_rows,
.lower_bound = gate_lower_bound,
.norm_eps = norm_eps,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(q)
offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(k)
offset:ds4_gpu_tensor_offset(k) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(v)
offset:ds4_gpu_tensor_offset(v) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(raw_gate)
offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(raw_beta)
offset:ds4_gpu_tensor_offset(raw_beta) atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(output_gate)
offset:ds4_gpu_tensor_offset(output_gate) atIndex:6];
[enc setBuffer:qw offset:(NSUInteger)qw_inner atIndex:7];
[enc setBuffer:kw offset:(NSUInteger)kw_inner atIndex:8];
[enc setBuffer:vw offset:(NSUInteger)vw_inner atIndex:9];
[enc setBuffer:a_log offset:(NSUInteger)a_inner atIndex:10];
[enc setBuffer:dt_bias offset:(NSUInteger)dt_inner atIndex:11];
[enc setBuffer:output_norm offset:(NSUInteger)norm_inner atIndex:12];
[enc setBuffer:ds4_gpu_tensor_buffer(conv_state)
offset:ds4_gpu_tensor_offset(conv_state) atIndex:13];
[enc setBuffer:ds4_gpu_tensor_buffer(recurrent_state)
offset:ds4_gpu_tensor_offset(recurrent_state) atIndex:14];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:15];
[enc setThreadgroupMemoryLength:656u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_rows, n_heads, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(
cb, owned, "GLM-5.3 fused KDA decode");
}
}
int ds4_gpu_glm53_kda_prefill(
ds4_gpu_tensor *out,
ds4_gpu_tensor *conv_state,
ds4_gpu_tensor *recurrent_state,
ds4_gpu_tensor *q,
ds4_gpu_tensor *k,
ds4_gpu_tensor *v,
ds4_gpu_tensor *raw_gate,
const ds4_gpu_tensor *raw_beta,
const ds4_gpu_tensor *output_gate,
const void *model_map,
uint64_t model_size,
uint64_t q_conv_offset,
uint64_t k_conv_offset,
uint64_t v_conv_offset,
uint64_t a_log_offset,
uint64_t dt_bias_offset,
uint64_t output_norm_offset,
uint32_t n_heads,
uint32_t n_tokens,
float gate_lower_bound,
float norm_eps) {
enum { GLM53_KDA_DIM = 128, GLM53_KDA_HISTORY = 3 };
if (!g_initialized && !ds4_gpu_init()) return 0;
uint64_t projection = 0, activation_elements = 0;
uint64_t conv_elements = 0, state_elements = 0;
if (n_heads == 0 || n_tokens == 0 || gate_lower_bound >= 0.0f ||
!glm53_gpu_mul_u64(n_heads, GLM53_KDA_DIM, &projection) ||
!glm53_gpu_mul_u64(projection, n_tokens, &activation_elements) ||
!glm53_gpu_mul_u64(projection,
3u * GLM53_KDA_HISTORY, &conv_elements) ||
!glm53_gpu_mul_u64(projection, GLM53_KDA_DIM, &state_elements) ||
!glm53_gpu_tensor_has(q, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(k, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(v, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(raw_gate, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(output_gate, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(raw_beta,
(uint64_t)n_tokens * n_heads, sizeof(float)) ||
!glm53_gpu_tensor_has(out, activation_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(conv_state, conv_elements, sizeof(float)) ||
!glm53_gpu_tensor_has(recurrent_state, state_elements, sizeof(float))) {
fprintf(stderr, "ds4: GLM-5.3 KDA prefill received invalid buffers\n");
return 0;
}
uint64_t conv_bytes = 0, dt_bytes = 0;
if (!glm53_gpu_mul_u64(projection, 4u * sizeof(float), &conv_bytes) ||
!glm53_gpu_mul_u64(projection, sizeof(float), &dt_bytes)) {
return 0;
}
const uint64_t a_log_bytes = (uint64_t)n_heads * sizeof(float);
const uint64_t norm_bytes = GLM53_KDA_DIM * sizeof(float);
@autoreleasepool {
uint64_t qw_inner = 0, kw_inner = 0, vw_inner = 0;
uint64_t a_inner = 0, dt_inner = 0, norm_inner = 0;
id<MTLBuffer> qw = glm53_gpu_weight_buffer(
model_map, model_size, q_conv_offset, conv_bytes,
&qw_inner, "KDA Q convolution");
id<MTLBuffer> kw = glm53_gpu_weight_buffer(
model_map, model_size, k_conv_offset, conv_bytes,
&kw_inner, "KDA K convolution");
id<MTLBuffer> vw = glm53_gpu_weight_buffer(
model_map, model_size, v_conv_offset, conv_bytes,
&vw_inner, "KDA V convolution");
id<MTLBuffer> a_log = glm53_gpu_weight_buffer(
model_map, model_size, a_log_offset, a_log_bytes,
&a_inner, "KDA A_log");
id<MTLBuffer> dt_bias = glm53_gpu_weight_buffer(
model_map, model_size, dt_bias_offset, dt_bytes,
&dt_inner, "KDA dt bias");
id<MTLBuffer> output_norm = glm53_gpu_weight_buffer(
model_map, model_size, output_norm_offset, norm_bytes,
&norm_inner, "KDA output norm");
id<MTLComputePipelineState> prep_pipeline =
ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_prepare");
id<MTLComputePipelineState> recurrence_pipeline =
ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_recurrence");
id<MTLComputePipelineState> output_pipeline =
ds4_gpu_get_pipeline("kernel_glm53_kda_prefill_output");
if (!qw || !kw || !vw || !a_log || !dt_bias || !output_norm ||
!prep_pipeline || !recurrence_pipeline || !output_pipeline) {
return 0;
}
glm53_gpu_kda_args args = {
.n_heads = n_heads,
.n_rows = n_tokens,
.lower_bound = gate_lower_bound,
.norm_eps = norm_eps,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:prep_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(q)
offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(k)
offset:ds4_gpu_tensor_offset(k) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(v)
offset:ds4_gpu_tensor_offset(v) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(raw_gate)
offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4];
[enc setBuffer:qw offset:(NSUInteger)qw_inner atIndex:5];
[enc setBuffer:kw offset:(NSUInteger)kw_inner atIndex:6];
[enc setBuffer:vw offset:(NSUInteger)vw_inner atIndex:7];
[enc setBuffer:a_log offset:(NSUInteger)a_inner atIndex:8];
[enc setBuffer:dt_bias offset:(NSUInteger)dt_inner atIndex:9];
[enc setBuffer:ds4_gpu_tensor_buffer(conv_state)
offset:ds4_gpu_tensor_offset(conv_state) atIndex:10];
[enc setThreadgroupMemoryLength:264u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_heads, 1, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
[enc setComputePipelineState:recurrence_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(q)
offset:ds4_gpu_tensor_offset(q) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(k)
offset:ds4_gpu_tensor_offset(k) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(v)
offset:ds4_gpu_tensor_offset(v) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(raw_gate)
offset:ds4_gpu_tensor_offset(raw_gate) atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(raw_beta)
offset:ds4_gpu_tensor_offset(raw_beta) atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(recurrent_state)
offset:ds4_gpu_tensor_offset(recurrent_state) atIndex:6];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:7];
[enc dispatchThreadgroups:MTLSizeMake(n_heads, 32, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
[enc setComputePipelineState:output_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(out)
offset:ds4_gpu_tensor_offset(out) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(output_gate)
offset:ds4_gpu_tensor_offset(output_gate) atIndex:2];
[enc setBuffer:output_norm offset:(NSUInteger)norm_inner atIndex:3];
[enc setThreadgroupMemoryLength:4u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_tokens, n_heads, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(
cb, owned, "GLM-5.3 KDA layer-major prefill");
}
}
static bool dsv41_tensor_has_floats(const ds4_gpu_tensor *tensor, uint64_t count) {
return tensor && count <= UINT64_MAX / sizeof(float) &&
ds4_gpu_tensor_bytes(tensor) >= count * sizeof(float);
}
/* Frequency rounding errors accumulate into phase errors at long contexts.
* Preserve the reference's pow/reciprocal and YaRN operation order here. */
#pragma float_control(precise, on, push)
static int ds4_gpu_dsv41_rope_impl(ds4_gpu_tensor *x, uint32_t width, uint32_t heads,
uint32_t rows, uint32_t start, uint32_t stride,
bool compressed, bool inverse, uint32_t round_all);
int ds4_gpu_dsv41_rope_stride(ds4_gpu_tensor *x, uint32_t width, uint32_t heads,
uint32_t rows, uint32_t start, uint32_t stride,
bool compressed, bool inverse) {
return ds4_gpu_dsv41_rope_impl(x, width, heads, rows, start, stride, compressed, inverse, 0);
}
/* Every element is rounded to bf16 before the rotation, as a separate rounding
* pass followed by the rope would do. */
int ds4_gpu_dsv41_rope_bf16(ds4_gpu_tensor *x, uint32_t width, uint32_t heads,
uint32_t rows, uint32_t start, bool compressed, bool inverse) {
return ds4_gpu_dsv41_rope_impl(x, width, heads, rows, start, 1, compressed, inverse, 1);
}
/* The rope frequencies of the raw (0) and compressed (1) positions. */
static const float *ds4_gpu_dsv41_rope_frequencies(bool compressed) {
static float frequencies[2][32];
static dispatch_once_t once;
dispatch_once(&once, ^{
for (int kind = 0; kind < 2; kind++) {
const float base = kind ? 160000.0f : 10000.0f;
const float low = (float)floor(64.0 * log(65536.0 / (32.0 * 2.0 * M_PI)) / (2.0 * log(base)));
const float high = (float)ceil(64.0 * log(65536.0 / (2.0 * M_PI)) / (2.0 * log(base)));
for (int i = 0; i < 32; i++) {
const float denominator = powf(base, (float)i / 32.0f);
float f = 1.0f / denominator;
if (kind) {
const float ramp = fminf(1.0f, fmaxf(0.0f, (i - low) / (high - low)));
const float smooth = 1.0f - ramp;
const float interpolate = (f / 16.0f) * (1.0f - smooth);
const float extrapolate = f * smooth;
f = interpolate + extrapolate;
}
frequencies[kind][i] = f;
}
}
});
return frequencies[compressed ? 1 : 0];
}
static int ds4_gpu_dsv41_rope_impl(ds4_gpu_tensor *x, uint32_t width, uint32_t heads,
uint32_t rows, uint32_t start, uint32_t stride,
bool compressed, bool inverse, uint32_t round_all) {
if (width < 64 || !heads || !rows || rows > 1048576 || !stride ||
(uint64_t)start + (uint64_t)(rows - 1u) * stride >= 1048576u ||
(uint64_t)heads * rows > UINT64_MAX / width ||
!dsv41_tensor_has_floats(x, (uint64_t)width * heads * rows)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_rope");
if (!pipeline) return 0;
struct {
uint32_t width, heads, rows, start, inverse, stride;
float frequencies[32];
uint32_t round_all;
} args = {width, heads, rows, start, inverse, stride, {0}, round_all};
memcpy(args.frequencies, ds4_gpu_dsv41_rope_frequencies(compressed), sizeof(args.frequencies));
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc dispatchThreadgroups:MTLSizeMake(heads, rows, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 unit-magnitude RoPE");
}
}
#pragma float_control(pop)
int ds4_gpu_dsv41_rope(ds4_gpu_tensor *x, uint32_t width, uint32_t heads,
uint32_t rows, uint32_t start, bool compressed, bool inverse) {
return ds4_gpu_dsv41_rope_stride(x, width, heads, rows, start, 1, compressed, inverse);
}
int ds4_gpu_dsv41_router_rows(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs, const ds4_gpu_tensor *logits,
const void *model_map, uint64_t model_size, uint64_t bias_offset,
uint32_t n_expert, uint32_t top_k, float scale, uint32_t rows) {
if (!selected || !weights || !probs || !logits || !n_expert || !top_k || top_k > n_expert ||
top_k > 32u || !rows || rows > 65535u) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_router_one");
if (!pipeline) return 0;
NSUInteger nth = 1;
while (nth < n_expert) nth *= 2u;
if (nth > pipeline.maxTotalThreadsPerThreadgroup) return 0;
const uint64_t bias_bytes = (uint64_t)n_expert * sizeof(float);
if (bias_offset > model_size || bias_bytes > model_size - bias_offset) return 0;
uint64_t bias_inner = 0;
id<MTLBuffer> biasbuf = ds4_gpu_wrap_model_range(model_map, model_size, bias_offset,
bias_bytes, &bias_inner);
if (!biasbuf || ds4_gpu_tensor_bytes(logits) < bias_bytes * rows ||
ds4_gpu_tensor_bytes(probs) < bias_bytes * rows ||
ds4_gpu_tensor_bytes(selected) < (uint64_t)top_k * rows * sizeof(int32_t) ||
ds4_gpu_tensor_bytes(weights) < (uint64_t)top_k * rows * sizeof(float)) return 0;
const struct { uint32_t n_expert, top_k, has_bias; float scale; } args =
{n_expert, top_k, 1u, scale};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(logits) offset:ds4_gpu_tensor_offset(logits) atIndex:1];
[enc setBuffer:biasbuf offset:(NSUInteger)bias_inner atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(probs) offset:ds4_gpu_tensor_offset(probs) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(selected) offset:ds4_gpu_tensor_offset(selected) atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(weights) offset:ds4_gpu_tensor_offset(weights) atIndex:5];
[enc setThreadgroupMemoryLength:((nth * 12u) + 15u + 128u) & ~(NSUInteger)15u atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, rows, 1) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 router");
}
}
int ds4_gpu_dsv41_router_one(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights,
ds4_gpu_tensor *probs, const ds4_gpu_tensor *logits,
const void *model_map, uint64_t model_size, uint64_t bias_offset,
uint32_t n_expert, uint32_t top_k, float scale) {
return ds4_gpu_dsv41_router_rows(selected, weights, probs, logits, model_map, model_size,
bias_offset, n_expert, top_k, scale, 1u);
}
/* A few rows through the single-row F32 matvec kernel (one grid row each), so
* every row reduces exactly as a decode step does. */
int ds4_gpu_dsv41_matmul_f32_rows(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x, uint32_t rows) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !x || !model_map || !in_dim || !out_dim || !rows || rows > 65535u ||
ds4_gpu_tensor_bytes(x) < (uint64_t)rows * in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out) < (uint64_t)rows * out_dim * sizeof(float)) return 0;
@autoreleasepool {
const uint64_t row_bytes = (uint64_t)in_dim * sizeof(float);
const uint64_t weight_bytes = row_bytes * out_dim;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0;
uint64_t inner = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner);
if (!wbuf) return 0;
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_f32_mv_args(in_dim, out_dim, rows);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_plain_mv_dispatch(in_dim, 1);
if (ds4_gpu_plain_mv_single_row(out_dim)) {
mv_dispatch.nr0 = 1;
mv_dispatch.smem = 32u * sizeof(float);
}
mv_args.nr0 = mv_dispatch.nr0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out) offset:ds4_gpu_tensor_offset(out) atIndex:3];
if (mv_dispatch.smem) [enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0,
rows, 1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 F32 matvec rows");
}
}
int ds4_gpu_dsv41_quantize(ds4_gpu_tensor *x, uint32_t width, uint32_t rows,
ds4_v41_activation_format format) {
const uint32_t block = format == DS4_V41_FP4_E4M3 ? 16u : 32u;
if (!width || !rows || format < DS4_V41_BF16 || format > DS4_V41_FP4_E4M3 ||
(format != DS4_V41_BF16 && width % block) ||
!dsv41_tensor_has_floats(x, (uint64_t)width * rows)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
const uint64_t count = (uint64_t)width * rows;
const bool linear = format == DS4_V41_BF16 && count <= UINT32_MAX &&
ds4_gpu_tensor_offset(x) % 16u == 0 &&
!getenv("DS4_METAL_DISABLE_V41_LINEAR_BF16");
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(linear ?
"kernel_dsv41_bf16_linear" : "kernel_dsv41_quantize");
if (!pipeline) return 0;
const uint32_t args[] = {width, rows, (uint32_t)format};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
if (linear) [enc setBytes:&count length:sizeof(count) atIndex:0];
else [enc setBytes:args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:1];
if (linear) {
[enc dispatchThreadgroups:MTLSizeMake((count + 1023u) / 1024u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
} else {
[enc dispatchThreadgroups:MTLSizeMake(((uint64_t)width + block - 1u) / block, rows, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
}
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 activation quantization");
}
}
typedef struct {
uint32_t width, start, inverse;
float frequencies[32];
} ds4_gpu_mv_rope_args;
static ds4_gpu_mv_rope_args ds4_gpu_dsv41_mv_rope_args(uint32_t width, uint32_t start,
bool compressed, bool inverse) {
ds4_gpu_mv_rope_args rope = {width, start, inverse ? 1u : 0u, {0}};
memcpy(rope.frequencies, ds4_gpu_dsv41_rope_frequencies(compressed), sizeof(rope.frequencies));
return rope;
}
int ds4_gpu_dsv41_project_q(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x, uint32_t pos, bool compressed) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!out || !x || !model_map || !in_dim || (in_dim & 31u) || !out_dim || out_dim % 512u ||
out_dim > 65536u || pos >= 1048576u ||
ds4_gpu_tensor_bytes(x) < (uint64_t)in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out) < (uint64_t)out_dim * sizeof(float)) return 0;
@autoreleasepool {
const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u;
const uint64_t weight_bytes = row_bytes * out_dim;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0;
uint64_t inner = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner);
if (!wbuf) return 0;
ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch dispatch = ds4_gpu_make_q8_0_mv_dispatch();
if (strcmp(dispatch.function_name, "kernel_mul_mv_q8_0_f32") != 0 || dispatch.nr0 != 2) return 0;
args.nr0 = dispatch.nr0;
const ds4_gpu_mv_rope_args rope = ds4_gpu_dsv41_mv_rope_args(512u, pos, compressed, false);
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_mul_mv_q8_0_f32_bf16_rope", dispatch.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out) offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setBytes:&rope length:sizeof(rope) atIndex:4];
[enc setThreadgroupMemoryLength:dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + 1u) / 2u, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 q projection with rope");
}
}
int ds4_gpu_dsv41_matmul_expand_rows(ds4_gpu_tensor *out_hc, const void *model_map, uint64_t model_size,
uint64_t weight_offset, uint32_t in_dim, uint32_t out_dim,
const ds4_gpu_tensor *x, const ds4_gpu_tensor *add,
const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split,
uint32_t n_hc, uint32_t rows) {
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
if (!out_hc || !x || !residual_hc || !split || !model_map || n_hc != 4 || !in_dim ||
(in_dim & 31u) || !out_dim || out_dim > 65536u || !rows || rows > 8u ||
ds4_gpu_tensor_bytes(x) < (uint64_t)rows * in_dim * sizeof(float) ||
(add && ds4_gpu_tensor_bytes(add) < (uint64_t)rows * out_dim * sizeof(float)) ||
ds4_gpu_tensor_bytes(residual_hc) < (uint64_t)rows * n_hc * out_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out_hc) < (uint64_t)rows * n_hc * out_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(split) < (uint64_t)rows * mix_hc * sizeof(float)) return 0;
@autoreleasepool {
const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u;
const uint64_t weight_bytes = row_bytes * out_dim;
if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0;
uint64_t inner = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size, weight_offset, weight_bytes, &inner);
if (!wbuf) return 0;
ds4_gpu_q8_0_matvec_args args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch dispatch = ds4_gpu_make_q8_0_mv_dispatch();
if (strcmp(dispatch.function_name, "kernel_mul_mv_q8_0_f32") != 0 || dispatch.nr0 != 2) return 0;
args.nr0 = dispatch.nr0;
args.ne11 = (int32_t)rows;
args.nb12 = (uint64_t)rows * in_dim * sizeof(float);
args.ne1 = (int32_t)rows;
const struct { uint32_t n_embd, has_add, split_stride; } hc = { out_dim, add != NULL, (uint32_t)mix_hc };
const ds4_gpu_tensor *add_or_x = add ? add : x;
const NSUInteger split_offset = ds4_gpu_tensor_offset(split);
if (rows > 1u) { dispatch.nr0 = 4; args.nr0 = 4; }
const int mma_off = ds4_gpu_v41_rows_mma_off();
const int mma = in_dim >= 4096u ? 8 : 4;
const bool mma_kernel = !mma_off && rows >= (in_dim >= 8192u ? 4u : 3u) &&
(in_dim % (mma == 8 ? 256u : 128u)) == 0 && (out_dim % 8u) == 0;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_mul_mv_pipeline(
mma_kernel ? (mma == 8 ? "kernel_mul_mv_q8_0_f32_bf16io_hc_expand4_mma_rows8" :
"kernel_mul_mv_q8_0_f32_bf16io_hc_expand4_mma_rows") :
rows > 1u ? "kernel_mul_mv_q8_0_f32_bf16io_hc_expand4_rows_seq4" : "kernel_mul_mv_q8_0_f32_bf16io_hc_expand4",
dispatch.nsg);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out_hc) offset:ds4_gpu_tensor_offset(out_hc) atIndex:3];
[enc setBytes:&hc length:sizeof(hc) atIndex:4];
[enc setBuffer:ds4_gpu_tensor_buffer(add_or_x) offset:ds4_gpu_tensor_offset(add_or_x) atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(residual_hc) offset:ds4_gpu_tensor_offset(residual_hc) atIndex:6];
[enc setBuffer:ds4_gpu_tensor_buffer(split) offset:split_offset + n_hc * sizeof(float) atIndex:7];
[enc setBuffer:ds4_gpu_tensor_buffer(split) offset:split_offset + 2u * n_hc * sizeof(float) atIndex:8];
[enc setThreadgroupMemoryLength:(mma_kernel ? (NSUInteger)mma * (2u * 8u * 32u + 64u) * sizeof(float) :
rows > 1u ? 32u * sizeof(float) * 4u : dispatch.smem) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(mma_kernel ? (NSUInteger)out_dim / 8u :
((NSUInteger)out_dim + (NSUInteger)dispatch.nr0 - 1u) / (NSUInteger)dispatch.nr0, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, mma_kernel ? (NSUInteger)mma : (NSUInteger)dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 projection with HC expand");
}
}
int ds4_gpu_dsv41_attention_low(ds4_gpu_tensor *low, const void *model_map, uint64_t model_size,
uint64_t out_a_offset, uint32_t group_dim, uint32_t rank,
uint32_t n_groups, ds4_gpu_tensor *heads, uint32_t pos, bool compressed) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!low || !heads || !model_map || !group_dim || group_dim % 512u || !rank || !n_groups ||
pos >= 1048576u || (uint64_t)n_groups * rank > UINT32_MAX) return 0;
@autoreleasepool {
const uint64_t row_a_bytes = (uint64_t)(group_dim / 32u) * 34u;
const uint64_t out_a_bytes = (uint64_t)n_groups * rank * row_a_bytes;
if (out_a_offset > model_size || out_a_bytes > model_size - out_a_offset ||
ds4_gpu_tensor_bytes(heads) < (uint64_t)n_groups * group_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(low) < (uint64_t)n_groups * rank * sizeof(float)) return 0;
uint64_t inner = 0;
id<MTLBuffer> wbuf = ds4_gpu_wrap_model_range(model_map, model_size, out_a_offset, out_a_bytes, &inner);
if (!wbuf) return 0;
ds4_gpu_mul_mv_id_args args = {
.nei0 = (int32_t)n_groups,
.nei1 = 1,
.nbi1 = 0,
.ne00 = (int32_t)group_dim,
.ne01 = (int32_t)rank,
.ne02 = (int32_t)n_groups,
.nb00 = 34,
.nb01 = row_a_bytes,
.nb02 = (uint64_t)rank * row_a_bytes,
.ne10 = (int32_t)group_dim,
.ne11 = (int32_t)n_groups,
.ne12 = 1,
.ne13 = 1,
.nb10 = sizeof(float),
.nb11 = (uint64_t)group_dim * sizeof(float),
.nb12 = (uint64_t)n_groups * group_dim * sizeof(float),
.ne0 = (int32_t)rank,
.ne1 = (int32_t)n_groups,
.nb1 = (uint64_t)rank * sizeof(float),
.nr0 = 2,
};
const ds4_gpu_mv_rope_args rope = ds4_gpu_dsv41_mv_rope_args(512u, pos, compressed, true);
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline("kernel_dsv4_attn_out_low_q8_0_f32_rope", 4);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:wbuf offset:(NSUInteger)inner atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(heads) offset:ds4_gpu_tensor_offset(heads) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(low) offset:ds4_gpu_tensor_offset(low) atIndex:3];
[enc setBytes:&rope length:sizeof(rope) atIndex:4];
[enc setThreadgroupMemoryLength:32u * 2u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)rank + 1u) / 2u, 1, n_groups)
threadsPerThreadgroup:MTLSizeMake(32, 4, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 attention low projection with rope");
}
}
int ds4_gpu_dsv41_norm_pair_rows(ds4_gpu_tensor *out0, const ds4_gpu_tensor *x0, uint64_t weight0_offset, uint32_t n0,
ds4_gpu_tensor *out1, const ds4_gpu_tensor *x1, uint64_t weight1_offset, uint32_t n1,
const void *model_map, uint64_t model_size, float eps, uint32_t rows) {
if (!rows || rows > 65535u) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
const NSUInteger ntg0 = ds4_gpu_rms_norm_threads(n0), ntg1 = ds4_gpu_rms_norm_threads(n1);
if (!out0 || !x0 || !out1 || !x1 || !model_map || !n0 || !n1 || (n0 & 3u) || (n1 & 3u) ||
ntg0 % 32u || ntg1 % 32u || ntg0 + ntg1 > 1024u ||
!dsv41_tensor_has_floats(x0, (uint64_t)rows * n0) || !dsv41_tensor_has_floats(out0, (uint64_t)rows * n0) ||
!dsv41_tensor_has_floats(x1, (uint64_t)rows * n1) || !dsv41_tensor_has_floats(out1, (uint64_t)rows * n1) ||
weight0_offset > model_size || (uint64_t)n0 * sizeof(float) > model_size - weight0_offset ||
weight1_offset > model_size || (uint64_t)n1 * sizeof(float) > model_size - weight1_offset) return 0;
@autoreleasepool {
uint64_t inner0 = 0, inner1 = 0;
id<MTLBuffer> w0 = ds4_gpu_wrap_model_range(model_map, model_size, weight0_offset, (uint64_t)n0 * sizeof(float), &inner0);
id<MTLBuffer> w1 = ds4_gpu_wrap_model_range(model_map, model_size, weight1_offset, (uint64_t)n1 * sizeof(float), &inner1);
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_norm_pair");
if (!w0 || !w1 || !pipeline || ntg0 + ntg1 > pipeline.maxTotalThreadsPerThreadgroup) return 0;
const struct { uint32_t n0, n1, ntg0, ntg1; float eps; } args = {n0, n1, (uint32_t)ntg0, (uint32_t)ntg1, eps};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x0) offset:ds4_gpu_tensor_offset(x0) atIndex:1];
[enc setBuffer:w0 offset:(NSUInteger)inner0 atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(out0) offset:ds4_gpu_tensor_offset(out0) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(x1) offset:ds4_gpu_tensor_offset(x1) atIndex:4];
[enc setBuffer:w1 offset:(NSUInteger)inner1 atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(out1) offset:ds4_gpu_tensor_offset(out1) atIndex:6];
[enc setThreadgroupMemoryLength:64u * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(1, rows, 1)
threadsPerThreadgroup:MTLSizeMake(ntg0 + ntg1, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 paired norm");
}
}
int ds4_gpu_dsv41_rope_quantize(const ds4_gpu_tensor *x, ds4_gpu_tensor *dst, uint64_t dst_offset,
uint32_t width, uint32_t start, bool compressed,
ds4_v41_activation_format format) {
const uint32_t block = format == DS4_V41_FP4_E4M3 ? 16u : 32u;
static int off = -1;
if (off < 0) off = getenv("DS4_METAL_DISABLE_V41_ROPE_QUANTIZE") != NULL;
if (off || width < 64u || width % block || format < DS4_V41_FP8_E8M0 || format > DS4_V41_FP4_E4M3 ||
start >= 1048576u || dst_offset % 4u || !dsv41_tensor_has_floats(x, width) ||
!dst || ds4_gpu_tensor_bytes(dst) < dst_offset + (uint64_t)width * sizeof(float)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_rope_quantize");
if (!pipeline) return 0;
struct {
uint32_t width, start, inverse, mode;
float frequencies[32];
} args = {width, start, 0u, (uint32_t)format, {0}};
memcpy(args.frequencies, ds4_gpu_dsv41_rope_frequencies(compressed), sizeof(args.frequencies));
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(dst) offset:ds4_gpu_tensor_offset(dst) + (NSUInteger)dst_offset atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(width / block, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 rope, quantization and store");
}
}
int ds4_gpu_dsv41_engram_add(ds4_gpu_tensor *residual,
const ds4_gpu_tensor *kv,
const ds4_gpu_tensor *q_weight,
const ds4_gpu_tensor *k_weight,
const ds4_gpu_tensor *mask,
uint32_t width, uint32_t rows, float eps) {
const uint64_t count = (uint64_t)width * rows;
if (!width || !rows || !isfinite(eps) || eps <= 0 || count > UINT64_MAX / 5u ||
!dsv41_tensor_has_floats(residual, count * 4u) ||
!dsv41_tensor_has_floats(kv, count * 5u) ||
!dsv41_tensor_has_floats(q_weight, (uint64_t)width * 4u) ||
!dsv41_tensor_has_floats(k_weight, (uint64_t)width * 4u) ||
(mask && ds4_gpu_tensor_bytes(mask) < rows)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_engram_add");
if (!pipeline) return 0;
const struct { uint32_t width, rows; float eps; uint32_t masked; } args =
{width, rows, eps, mask != NULL};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
const ds4_gpu_tensor *inputs[] = {residual, kv, q_weight, k_weight, mask ? mask : residual};
for (NSUInteger i = 0; i < 5; i++)
[enc setBuffer:ds4_gpu_tensor_buffer(inputs[i])
offset:ds4_gpu_tensor_offset(inputs[i]) atIndex:i + 1];
[enc dispatchThreadgroups:MTLSizeMake(rows, 4, 1)
threadsPerThreadgroup:MTLSizeMake(32, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 Engram gate");
}
}
int ds4_gpu_dsv41_carry_copy(ds4_gpu_tensor *packed, uint32_t row_offset,
ds4_gpu_tensor *plain, uint32_t width, uint32_t rows,
uint32_t format, bool pack) {
if (!width || !rows || rows > UINT32_MAX - row_offset ||
format > DS4_V41_CARRY_MASK || packed == plain) return 0;
const uint32_t words = format == DS4_V41_CARRY_BF16 ?
((uint64_t)width + 1u) / 2u : ((uint64_t)width + 31u) / 32u;
if (!dsv41_tensor_has_floats(packed, (uint64_t)(row_offset + rows) * words) ||
!dsv41_tensor_has_floats(plain, (uint64_t)rows * width)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_carry_copy");
if (!pipeline) return 0;
const uint32_t args[] = {width, rows, words, format, pack};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(packed) offset:ds4_gpu_tensor_offset(packed) +
(NSUInteger)row_offset * words * sizeof(uint32_t) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(plain) offset:ds4_gpu_tensor_offset(plain) atIndex:2];
[enc dispatchThreadgroups:MTLSizeMake(((uint64_t)width + 127u) / 128u, rows, 1)
threadsPerThreadgroup:MTLSizeMake(128, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 compact prefill carry");
}
}
int ds4_gpu_dsv41_pool2(ds4_gpu_tensor *out,
const ds4_gpu_tensor *kv, const ds4_gpu_tensor *scores,
ds4_gpu_tensor *previous_kv, ds4_gpu_tensor *previous_scores,
uint32_t width, uint32_t rows, uint32_t start) {
const uint64_t count = (uint64_t)width * rows;
const uint32_t pairs = (uint32_t)(((uint64_t)rows + (start & 1u)) / 2u);
if (!width || !rows || rows > UINT32_MAX - start ||
!dsv41_tensor_has_floats(kv, count) || !dsv41_tensor_has_floats(scores, count) ||
!dsv41_tensor_has_floats(previous_kv, width) ||
!dsv41_tensor_has_floats(previous_scores, width) ||
(pairs && !dsv41_tensor_has_floats(out, (uint64_t)width * pairs))) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
if (pairs) {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_pool2");
if (!pipeline) return 0;
const uint32_t args[] = {width, pairs, start & 1u};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(args) atIndex:0];
const ds4_gpu_tensor *inputs[] = {out, kv, scores, previous_kv, previous_scores};
for (NSUInteger i = 0; i < 5; i++)
[enc setBuffer:ds4_gpu_tensor_buffer(inputs[i])
offset:ds4_gpu_tensor_offset(inputs[i]) atIndex:i + 1];
[enc dispatchThreads:MTLSizeMake(width, pairs, 1)
threadsPerThreadgroup:MTLSizeMake(MIN(width, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, "V4.1 KV pair pooling")) return 0;
}
/* Keep the last even-position input, as the row-at-a-time path does.
* It is live only for an odd frontier, but matching it also makes
* snapshots independent of how the prefix was partitioned. */
const uint32_t last_even = (start + rows - 1u) & ~1u;
if (last_even >= start) {
const uint64_t bytes = (uint64_t)width * sizeof(float);
const uint64_t offset = (last_even - start) * bytes;
if (!ds4_gpu_tensor_copy(previous_kv, 0, kv, offset, bytes) ||
!ds4_gpu_tensor_copy(previous_scores, 0, scores, offset, bytes)) return 0;
}
}
return 1;
}
static int dsv41_candidates(ds4_gpu_tensor *out, const ds4_gpu_tensor *scores,
const ds4_gpu_tensor *mask, uint32_t width,
uint32_t rows, uint32_t start, uint32_t ratio,
uint32_t mask_stride) {
if (!width || width > UINT32_MAX - 7u || !rows || !ratio ||
rows > UINT32_MAX - start) return 0;
const uint32_t blocks = (width + 7u) / 8u;
const uint32_t output_width = mask ? width : blocks;
if (!mask_stride) mask_stride = blocks;
if (mask_stride < blocks ||
!dsv41_tensor_has_floats(scores, (uint64_t)width * rows) ||
!dsv41_tensor_has_floats(out, (uint64_t)output_width * rows) ||
(mask && !dsv41_tensor_has_floats(mask, (uint64_t)mask_stride * (rows - 1u) + blocks)))
return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(mask ?
"kernel_dsv41_candidate_filter" : "kernel_dsv41_candidate_blocks");
if (!pipeline) return 0;
const uint32_t args[] = {width, rows, start, ratio, mask_stride};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(args) atIndex:0];
const ds4_gpu_tensor *inputs[] = {scores, out, mask ? mask : scores};
for (NSUInteger i = 0; i < 3; i++)
[enc setBuffer:ds4_gpu_tensor_buffer(inputs[i])
offset:ds4_gpu_tensor_offset(inputs[i]) atIndex:i + 1];
[enc dispatchThreads:MTLSizeMake(output_width, rows, 1)
threadsPerThreadgroup:MTLSizeMake(MIN(output_width, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 candidate selection");
}
}
int ds4_gpu_dsv41_candidate_blocks(ds4_gpu_tensor *blocks,
const ds4_gpu_tensor *scores,
uint32_t width, uint32_t rows,
uint32_t start, uint32_t ratio) {
return dsv41_candidates(blocks, scores, NULL, width, rows, start, ratio, 0);
}
int ds4_gpu_dsv41_candidate_filter(ds4_gpu_tensor *scores,
const ds4_gpu_tensor *block_mask,
uint32_t width, uint32_t rows,
uint32_t start, uint32_t ratio,
uint32_t mask_stride) {
if (!block_mask) return 0;
return dsv41_candidates(scores, scores, block_mask, width, rows, start, ratio, mask_stride);
}
/* Row t sees (start + t + 8) / 8 blocks: the causal sort with start + 7 and
* ratio 8 keeps each row's own width, padding and tie order. */
int ds4_gpu_dsv41_candidate_topk_batch(ds4_gpu_tensor *selected,
const ds4_gpu_tensor *blocks,
uint32_t width, uint32_t rows, uint32_t start) {
if (!rows || start > UINT32_MAX - 8u || rows > UINT32_MAX - 8u - start ||
width > INT32_MAX || rows > INT32_MAX ||
(start + rows + 7u) / 8u > width || (start + 8u) / 8u <= 2048u) return 0;
return ds4_gpu_indexer_topk_tensor_impl(selected, blocks, width, rows, 2048u, start + 7u, 8u);
}
int ds4_gpu_dsv41_candidate_mask_batch(ds4_gpu_tensor *mask,
const ds4_gpu_tensor *selected,
uint32_t width, uint32_t rows,
uint32_t mask_stride) {
return ds4_gpu_dsv4_topk_mask_tensor_impl(mask, selected, width, rows, 2048u, mask_stride);
}
int ds4_gpu_dsv41_hc_mean(uint32_t rows, uint32_t dim, uint32_t hc,
const ds4_gpu_tensor *stream, ds4_gpu_tensor *out,
uint32_t out_stride, uint32_t out_off) {
if (!rows || !dim || !hc || out_stride < dim || out_off > out_stride - dim ||
!dsv41_tensor_has_floats(stream, (uint64_t)rows * hc * dim) ||
!dsv41_tensor_has_floats(out, (uint64_t)rows * out_stride)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv41_hc_mean");
if (!pipeline) return 0;
const uint32_t args[] = {dim, hc, out_stride, out_off};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(stream) offset:ds4_gpu_tensor_offset(stream) atIndex:1];
[enc setBuffer:ds4_gpu_tensor_buffer(out) offset:ds4_gpu_tensor_offset(out) atIndex:2];
[enc dispatchThreads:MTLSizeMake(dim, rows, 1)
threadsPerThreadgroup:MTLSizeMake(MIN(dim, 256u), 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 stream mean");
}
}
int ds4_gpu_dsv41_markov_chain(uint32_t block, uint32_t vocab, uint32_t rank, uint32_t dim,
const ds4_gpu_tensor *logits, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size,
uint64_t embed_offset, uint64_t head_offset, int f16,
const ds4_gpu_tensor *conf_proj, ds4_gpu_tensor *tokens,
ds4_gpu_tensor *conf, ds4_gpu_tensor *parts, uint32_t n_parts) {
if (!block || !vocab || !rank || (rank % 32u) || rank > 512u || !dim || !n_parts || n_parts > 4096u ||
!dsv41_tensor_has_floats(logits, (uint64_t)block * vocab) ||
!dsv41_tensor_has_floats(x, (uint64_t)block * dim) ||
!dsv41_tensor_has_floats(conf_proj, (uint64_t)dim + rank) ||
!dsv41_tensor_has_floats(tokens, (uint64_t)block + 1u) ||
!dsv41_tensor_has_floats(conf, block) ||
!dsv41_tensor_has_floats(parts, (uint64_t)n_parts * 2u)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
const uint64_t bytes = (uint64_t)vocab * rank * (f16 ? 2u : 4u);
uint64_t e_inner = 0, h_inner = 0;
id<MTLBuffer> ebuf = ds4_gpu_wrap_model_range(model_map, model_size, embed_offset, bytes, &e_inner);
id<MTLBuffer> hbuf = ds4_gpu_wrap_model_range(model_map, model_size, head_offset, bytes, &h_inner);
id<MTLComputePipelineState> part = ds4_gpu_get_pipeline("kernel_dsv41_markov_part");
id<MTLComputePipelineState> final = ds4_gpu_get_pipeline("kernel_dsv41_markov_final");
if (!ebuf || !hbuf || !part || !final) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLBuffer> pbuf = ds4_gpu_tensor_buffer(parts);
const NSUInteger p_val = ds4_gpu_tensor_offset(parts);
const NSUInteger p_idx = p_val + (NSUInteger)n_parts * sizeof(float);
for (uint32_t step = 0; step < block; step++) {
const uint32_t args[] = {vocab, rank, dim, step, n_parts, f16 ? 1u : 0u};
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:part];
[enc setBytes:args length:sizeof(args) atIndex:0];
[enc setBuffer:ds4_gpu_tensor_buffer(logits) offset:ds4_gpu_tensor_offset(logits) atIndex:1];
[enc setBuffer:ebuf offset:(NSUInteger)e_inner atIndex:2];
[enc setBuffer:hbuf offset:(NSUInteger)h_inner atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(tokens) offset:ds4_gpu_tensor_offset(tokens) atIndex:4];
[enc setBuffer:pbuf offset:p_val atIndex:5];
[enc setBuffer:pbuf offset:p_idx atIndex:6];
[enc setThreadgroupMemoryLength:(NSUInteger)rank * sizeof(float) atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(n_parts, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:final];
[enc setBytes:args length:sizeof(args) atIndex:0];
[enc setBuffer:pbuf offset:p_val atIndex:1];
[enc setBuffer:pbuf offset:p_idx atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(tokens) offset:ds4_gpu_tensor_offset(tokens) atIndex:3];
[enc setBuffer:ds4_gpu_tensor_buffer(x) offset:ds4_gpu_tensor_offset(x) atIndex:4];
[enc setBuffer:ebuf offset:(NSUInteger)e_inner atIndex:5];
[enc setBuffer:ds4_gpu_tensor_buffer(conf_proj) offset:ds4_gpu_tensor_offset(conf_proj) atIndex:6];
[enc setBuffer:ds4_gpu_tensor_buffer(conf) offset:ds4_gpu_tensor_offset(conf) atIndex:7];
[enc dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
}
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 Markov chain");
}
}
int ds4_gpu_dsv41_gather_kv(ds4_gpu_tensor *out, const ds4_gpu_tensor *source,
const ds4_gpu_tensor *ids, uint32_t source_rows,
uint32_t selected_rows) {
if (!source_rows || !selected_rows || selected_rows > 512 ||
selected_rows > source_rows ||
!dsv41_tensor_has_floats(source, (uint64_t)source_rows * 512u) ||
!dsv41_tensor_has_floats(out, (uint64_t)selected_rows * 512u) ||
!dsv41_tensor_has_floats(ids, selected_rows)) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
const uint64_t row_bytes = 512u * sizeof(float);
const ds4_gpu_get_rows_args args = {
.ne00t = 512, .ne00 = 512, .nb01 = row_bytes,
.nb02 = (uint64_t)source_rows * row_bytes,
.nb03 = (uint64_t)source_rows * row_bytes,
.ne10 = (int32_t)selected_rows, .nb10 = sizeof(int32_t),
.nb11 = (uint64_t)selected_rows * sizeof(int32_t),
.nb12 = (uint64_t)selected_rows * sizeof(int32_t),
.nb1 = row_bytes, .nb2 = (uint64_t)selected_rows * row_bytes,
.nb3 = (uint64_t)selected_rows * row_bytes,
};
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
id<MTLComputeCommandEncoder> enc = cb ? ds4_gpu_compute_encoder(cb) : nil;
if (!enc) return 0;
[enc setComputePipelineState:g_get_rows_f32_pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
const ds4_gpu_tensor *inputs[] = {source, ids, out};
for (NSUInteger i = 0; i < 3; i++)
[enc setBuffer:ds4_gpu_tensor_buffer(inputs[i])
offset:ds4_gpu_tensor_offset(inputs[i]) atIndex:i + 1];
[enc dispatchThreadgroups:MTLSizeMake(2u * selected_rows, 1, 1)
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "V4.1 sparse KV gather");
}
}
void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) {
(void)enabled;
}
/* Qwen3.8-Flash-Next kernels (metal/qwen4.metal): one dispatch per wrapper,
* weights bound from the model map by offset, transients as ds4_gpu_tensors.
* Every wrapper joins the active command batch. */
typedef struct {
id<MTLBuffer> buf;
NSUInteger off;
} qwen4_bind;
static bool qwen4_bind_tensor(qwen4_bind *b, const ds4_gpu_tensor *t, uint64_t min_bytes, const char *what) {
if (!t || !ds4_gpu_tensor_buffer(t) || ds4_gpu_tensor_bytes(t) < min_bytes) {
fprintf(stderr, "ds4: Qwen3.8 %s buffer is missing or undersized (%" PRIu64 " < %" PRIu64 ")\n",
what, t ? ds4_gpu_tensor_bytes(t) : 0, min_bytes);
return false;
}
b->buf = ds4_gpu_tensor_buffer(t);
b->off = ds4_gpu_tensor_offset(t);
return true;
}
static bool qwen4_bind_weight(qwen4_bind *b, const void *map, uint64_t size,
uint64_t offset, uint64_t bytes, const char *what) {
uint64_t inner = 0;
if (!map || offset > size || bytes > size - offset) {
fprintf(stderr, "ds4: Qwen3.8 %s range is outside the mapped model\n", what);
return false;
}
b->buf = ds4_gpu_wrap_model_range(map, size, offset, bytes, &inner);
b->off = (NSUInteger)inner;
return b->buf != nil;
}
enum {
QWEN4_K_HC_NORM_F16 = 0,
QWEN4_K_HC_NORM_F32,
QWEN4_K_HC_NORM_Q8,
QWEN4_K_HC_NORM_REUSE_F16,
QWEN4_K_HC_NORM_REUSE_F32,
QWEN4_K_HC_NORM_REUSE_Q8,
QWEN4_K_HC_GATE_MIX_F16,
QWEN4_K_HC_GATE_MIX_F16_PF,
QWEN4_K_HC_GATE_MIX_F32,
QWEN4_K_HC_GATE_MIX_Q8,
QWEN4_K_HC_GATE_MIX_PAIR_F16,
QWEN4_K_HC_GATE_MIX_PAIR_F16_PF,
QWEN4_K_HC_GATE_MIX_PAIR_F32,
QWEN4_K_HC_GATE_MIX_PAIR_Q8,
QWEN4_K_MULTI_GEMV,
QWEN4_K_HC_COMBINE,
QWEN4_K_CONV_STREAM,
QWEN4_K_GDN_PREP,
QWEN4_K_GDN_SCAN,
QWEN4_K_GDN_SCAN_R4,
QWEN4_K_GDN_OUT,
QWEN4_K_PLE_GATE,
QWEN4_K_PLE_CONV,
QWEN4_K_ROUTER,
QWEN4_K_ATTN_PREP,
QWEN4_K_IDX_BLOCK_KEY,
QWEN4_K_IDX_SCORE,
QWEN4_K_IDX_SCORE_VEC,
QWEN4_K_IDX_SELECT,
QWEN4_K_IDX_SELECT_PRE,
QWEN4_K_IDX_SCORE_MM,
QWEN4_K_IDX_EXPAND,
QWEN4_K_ATTN_DECODE_NPT8,
QWEN4_K_ATTN_DECODE_NPT4,
QWEN4_K_ATTN_DECODE_NPT1,
QWEN4_K_ATTN_MERGE_NPT8,
QWEN4_K_ATTN_MERGE_NPT4,
QWEN4_K_ATTN_MERGE_NPT1,
QWEN4_K_ATTN_MERGE_WIDE_NPT8,
QWEN4_K_ATTN_MERGE_WIDE_NPT4,
QWEN4_K_ATTN_MM,
QWEN4_K_MOE_MID,
QWEN4_K_MOE_MID_Q4K,
QWEN4_K_MOE_MID_Q4K_NR1,
QWEN4_K_MOE_DOWN,
QWEN4_K_MOE_DOWN_MXFP4_PF,
QWEN4_K_MOE_MID_Q4K_GROUPED,
QWEN4_K_MOE_DOWN_MXFP4_GROUPED,
QWEN4_K_MOE_REDUCE,
QWEN4_K_HC_COMBINE_NORM,
QWEN4_K_ARGMAX,
QWEN4_K_MTP_STAGE,
QWEN4_K_MTP_COMBINE,
QWEN4_K_GDN_FRONT,
QWEN4_K_MOE_BUILD_LISTS,
QWEN4_K_MOE_MM_MID,
QWEN4_K_MOE_MM_MID_NT1,
QWEN4_K_MOE_MM_MID_NT2,
QWEN4_K_MOE_MM_DOWN,
QWEN4_K_MOE_MM_DOWN_NT1,
QWEN4_K_MOE_MM_DOWN_NT2,
QWEN4_K_MOE_MM_DOWN_NT8,
QWEN4_K_MOE_MM_MID_NT8,
QWEN4_K_MOE_MM_MID_NAX,
QWEN4_K_MOE_MM_DOWN_NAX,
QWEN4_K_MOE_MM_MID_NAX64,
QWEN4_K_MOE_MM_DOWN_NAX64,
QWEN4_K_MOE_MM_MID_NAXF,
QWEN4_K_MOE_MM_DOWN_NAXF,
QWEN4_K_MOE_MM_MID_NAXF64,
QWEN4_K_MOE_MM_DOWN_NAXF64,
QWEN4_K_MOE_MM_MID_NAXC,
QWEN4_K_MOE_MM_DOWN_NAXC,
QWEN4_K_MOE_MM_MID_NAXC64,
QWEN4_K_MOE_MM_DOWN_NAXC64,
QWEN4_K_ROWS_F32_TO_F16,
QWEN4_K_DENSE_MM,
QWEN4_K_DENSE_MM_REDUCE,
QWEN4_K_BATCH_MM_Q8_T1,
QWEN4_K_BATCH_MM_Q8_T2,
QWEN4_K_GDN_SCAN_ROWS,
QWEN4_K_CONV_STREAM_ROWS,
QWEN4_K_GDN_SCAN_ROWS2,
QWEN4_K_CONV_STREAM_ROWS2,
QWEN4_K_ATTN_PREP_ROWS,
QWEN4_K_IDX_BLOCK_KEY_ROWS,
QWEN4_K_IDX_SCORE_ROWS,
QWEN4_K_IDX_SELECT_ROWS,
QWEN4_K_IDX_EXPAND_ROWS,
QWEN4_K_ATTN_DECODE_ROWS_NPT8,
QWEN4_K_ATTN_DECODE_ROWS_NPT4,
QWEN4_K_ATTN_MERGE_ROWS_NPT8,
QWEN4_K_ATTN_MERGE_ROWS_NPT4,
QWEN4_K_HC_LO_ACT,
QWEN4_K_HC_MIX_ROWS,
QWEN4_K_VIS_PATCH_FINISH,
QWEN4_K_VIS_LAYERNORM,
QWEN4_K_VIS_QKV_ROPE,
QWEN4_K_VIS_ATTENTION,
QWEN4_K_VIS_BIAS_RESIDUAL,
QWEN4_K_VIS_BIAS_ACT,
QWEN4_K_COUNT,
};
static const char *const qwen4_kernel_names[QWEN4_K_COUNT] = {
"kernel_qwen4_hc_norm_f16",
"kernel_qwen4_hc_norm_f32",
"kernel_qwen4_hc_norm_q8",
"kernel_qwen4_hc_norm_reuse_f16",
"kernel_qwen4_hc_norm_reuse_f32",
"kernel_qwen4_hc_norm_reuse_q8",
"kernel_qwen4_hc_gate_mix_f16",
"kernel_qwen4_hc_gate_mix_f16_pf",
"kernel_qwen4_hc_gate_mix_f32",
"kernel_qwen4_hc_gate_mix_q8",
"kernel_qwen4_hc_gate_mix_pair_f16",
"kernel_qwen4_hc_gate_mix_pair_f16_pf",
"kernel_qwen4_hc_gate_mix_pair_f32",
"kernel_qwen4_hc_gate_mix_pair_q8",
"kernel_qwen4_multi_gemv",
"kernel_qwen4_hc_combine",
"kernel_qwen4_conv_stream",
"kernel_qwen4_gdn_prep",
"kernel_qwen4_gdn_scan",
"kernel_qwen4_gdn_scan_r4",
"kernel_qwen4_gdn_out",
"kernel_qwen4_ple_gate",
"kernel_qwen4_ple_conv",
"kernel_qwen4_router_topk",
"kernel_qwen4_attn_prep",
"kernel_qwen4_idx_block_key",
"kernel_qwen4_idx_score",
"kernel_qwen4_idx_score_vec",
"kernel_qwen4_idx_select",
"kernel_qwen4_idx_select_pre",
"kernel_qwen4_idx_score_mm",
"kernel_qwen4_idx_expand",
"kernel_qwen4_attn_decode_npt8",
"kernel_qwen4_attn_decode_npt4",
"kernel_qwen4_attn_decode_npt1",
"kernel_qwen4_attn_merge_npt8",
"kernel_qwen4_attn_merge_npt4",
"kernel_qwen4_attn_merge_npt1",
"kernel_qwen4_attn_merge_wide_npt8",
"kernel_qwen4_attn_merge_wide_npt4",
"kernel_qwen4_attn_mm",
"kernel_qwen4_moe_mid",
"kernel_qwen4_moe_mid_q4k",
"kernel_qwen4_moe_mid_q4k_nr1",
"kernel_qwen4_moe_down",
"kernel_qwen4_moe_down_mxfp4_pf",
"kernel_qwen4_moe_mid_q4k_grouped",
"kernel_qwen4_moe_down_mxfp4_grouped",
"kernel_qwen4_moe_reduce",
"kernel_qwen4_hc_combine_norm_f16",
"kernel_qwen4_argmax",
"kernel_qwen4_mtp_stage",
"kernel_qwen4_mtp_combine",
"kernel_qwen4_gdn_front",
"kernel_qwen4_moe_build_lists",
"kernel_qwen4_moe_mm_mid",
"kernel_qwen4_moe_mm_mid_nt1",
"kernel_qwen4_moe_mm_mid_nt2",
"kernel_qwen4_moe_mm_down",
"kernel_qwen4_moe_mm_down_nt1",
"kernel_qwen4_moe_mm_down_nt2",
"kernel_qwen4_moe_mm_down_nt8",
"kernel_qwen4_moe_mm_mid_nt8",
"kernel_qwen4_moe_mm_mid_nax",
"kernel_qwen4_moe_mm_down_nax",
"kernel_qwen4_moe_mm_mid_nax64",
"kernel_qwen4_moe_mm_down_nax64",
"kernel_qwen4_moe_mm_mid_naxf",
"kernel_qwen4_moe_mm_down_naxf",
"kernel_qwen4_moe_mm_mid_naxf64",
"kernel_qwen4_moe_mm_down_naxf64",
"kernel_qwen4_moe_mm_mid_naxc",
"kernel_qwen4_moe_mm_down_naxc",
"kernel_qwen4_moe_mm_mid_naxc64",
"kernel_qwen4_moe_mm_down_naxc64",
"kernel_qwen4_rows_f32_to_f16",
"kernel_qwen4_dense_mm",
"kernel_qwen4_dense_mm_reduce",
"kernel_qwen4_batch_mm_q8_t1",
"kernel_qwen4_batch_mm_q8_t2",
"kernel_qwen4_gdn_scan_rows",
"kernel_qwen4_conv_stream_rows",
"kernel_qwen4_gdn_scan_rows2",
"kernel_qwen4_conv_stream_rows2",
"kernel_qwen4_attn_prep_rows",
"kernel_qwen4_idx_block_key_rows",
"kernel_qwen4_idx_score_rows",
"kernel_qwen4_idx_select_rows",
"kernel_qwen4_idx_expand_rows",
"kernel_qwen4_attn_decode_rows_npt8",
"kernel_qwen4_attn_decode_rows_npt4",
"kernel_qwen4_attn_merge_rows_npt8",
"kernel_qwen4_attn_merge_rows_npt4",
"kernel_qwen4_hc_lo_act",
"kernel_qwen4_hc_mix_rows",
"kernel_qwen4_vis_patch_finish",
"kernel_qwen4_vis_layernorm",
"kernel_qwen4_vis_qkv_rope",
"kernel_qwen4_vis_attention",
"kernel_qwen4_vis_bias_residual",
"kernel_qwen4_vis_bias_act",
};
typedef struct {
uint32_t n_tokens, n_slots, n_out, in_dim, out_rows, weight_type, row_bytes, list_cap;
uint64_t expert_bytes;
uint32_t n_expert, tiles_per_launch, tail_base, expert_major;
} qwen4_moe_mm_args;
/* Expert-major tile order (DS4_QWEN4_MOE_MM_ORDER=0 restores the tile-major
* grid): consecutive threadgroups share a token tile and one expert's rows
* stay cache-resident across its tiles. */
static uint32_t qwen4_moe_mm_expert_major(void) {
const int override = ds4_gpu_env_bool("DS4_QWEN4_MOE_MM_ORDER");
return override >= 0 ? (uint32_t)override : ds4_gpu_device_is_m5_apple_silicon() ? 1u : 0u;
}
static MTLSize qwen4_moe_mm_grid(uint32_t row_blocks, uint32_t n_expert, uint32_t tiles, uint32_t expert_major) {
return expert_major ? MTLSizeMake((NSUInteger)row_blocks * tiles, n_expert, 1)
: MTLSizeMake(row_blocks, n_expert, tiles);
}
static id<MTLComputePipelineState> g_qwen4_pipelines[QWEN4_K_COUNT];
#define QWEN4_ATTN_NSG 4
#define QWEN4_ATTN_MAX_SPLITS 64
#define QWEN4_ATTN_ROWS_MAX 64 /* decode batch rows the attention rows kernels take */
typedef struct {
uint32_t n_tokens, n_slots, in_dim, out_rows, weight_type, row_bytes;
uint64_t expert_bytes;
uint32_t has_shared, shared_type, shared_row_bytes, n_total_expert;
uint32_t list_cap, pad0;
} qwen4_moe_args;
static bool qwen4_moe_mv_specialize(uint32_t type) {
/* Constant quantization and logical width remove the generic decode
* branches. Keep the original per-lane reduction order and padded stride.
* M3 Ultra uses low-bit and MXFP4 down rows; M5 uses MXFP4 down rows. */
const int override = ds4_gpu_env_bool("DS4_QWEN4_MOE_MV_SPECIALIZE");
return override >= 0 ? override != 0 :
((type == 16u || type == 10u || type == 39u) && ds4_gpu_device_name_contains("M3 Ultra")) ||
(type == 39u && ds4_gpu_device_is_m5_apple_silicon());
}
static uint32_t qwen4_moe_mv_rows(void) {
return (uint32_t)ds4_gpu_env_u64("DS4_QWEN4_MOE_MV_NR", 1u, 1u, 4u);
}
static uint32_t qwen4_moe_mv_groups(uint32_t type) {
/* Q2_K and MXFP4 use sixteen groups on M3 Ultra; MXFP4 also does on M5.
* Each group keeps its own rows and unchanged per-lane reduction order. */
const uint32_t default_nsg =
((type == 10u || type == 39u) && ds4_gpu_device_name_contains("M3 Ultra")) ||
(type == 39u && ds4_gpu_device_is_m5_apple_silicon()) ? 16u : 8u;
return (uint32_t)ds4_gpu_env_u64("DS4_QWEN4_MOE_MV_NSG", default_nsg, 1u, 16u);
}
/* resident: buffers the kernel reaches through GPU addresses rather than
* bindings (the decode batch's per-session caches); they are made resident
* for this dispatch. */
static int qwen4_dispatch_resident(int kernel, const void *args, size_t args_len,
const qwen4_bind *binds, int n_binds,
MTLSize grid, MTLSize tg, NSUInteger tg_mem,
const qwen4_bind *resident, uint32_t n_resident) {
if (!g_initialized && !ds4_gpu_init()) return 0;
@autoreleasepool {
id<MTLComputePipelineState> pipeline = nil;
if (kernel == QWEN4_K_MOE_MID || kernel == QWEN4_K_MOE_DOWN || kernel == QWEN4_K_MOE_DOWN_MXFP4_PF) {
const qwen4_moe_args *a = args;
const bool specialize = qwen4_moe_mv_specialize(a->weight_type);
const uint32_t values[] = {a->weight_type, a->shared_type, specialize ? a->in_dim : 0u, specialize ? qwen4_moe_mv_rows() : 0u};
NSString *key = [NSString stringWithFormat:@"%s_type=%u_shared=%u_dim=%u_rows=%u",
qwen4_kernel_names[kernel], values[0], values[1], values[2], values[3]];
pipeline = [g_pipeline_cache objectForKey:key];
if (!pipeline) {
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
for (uint32_t i = 0; specialize && i < 4u; i++)
[constants setConstantValue:&values[i] type:MTLDataTypeUInt atIndex:901u+i];
NSError *error = nil;
id<MTLFunction> fn = [g_library newFunctionWithName:
[NSString stringWithUTF8String:qwen4_kernel_names[kernel]]
constantValues:constants error:&error];
if (fn) pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Qwen MoE decode pipeline failed: %s\n", [[error localizedDescription] UTF8String]);
return 0;
}
[g_pipeline_cache setObject:pipeline forKey:key];
}
} else if (kernel >= QWEN4_K_MOE_MM_MID && kernel <= QWEN4_K_MOE_MM_DOWN_NAXC64) {
/* Keep each quantization's dequantizer constant through the K
* loop. M3 Ultra and M5 have balanced full-model measurements;
* other devices can opt in, and zero restores the generic kernel. */
const int override = ds4_gpu_env_bool("DS4_QWEN4_MOE_MM_SPECIALIZE");
const bool specialize = override >= 0 ? override != 0 :
ds4_gpu_device_name_contains("M3 Ultra") ||
ds4_gpu_device_is_m5_apple_silicon();
const uint32_t type = specialize ?
((const qwen4_moe_mm_args *)args)->weight_type : 0u;
if (type >= 40u) return 0;
const uint32_t tail_base = ((const qwen4_moe_mm_args *)args)->tail_base;
NSString *key = [NSString stringWithFormat:@"%s_type=%u_tail=%u",
qwen4_kernel_names[kernel], type, tail_base];
pipeline = [g_pipeline_cache objectForKey:key];
if (!pipeline) {
MTLFunctionConstantValues *constants = [[MTLFunctionConstantValues alloc] init];
[constants setConstantValue:&type type:MTLDataTypeUInt atIndex:900];
[constants setConstantValue:&tail_base type:MTLDataTypeUInt atIndex:905];
NSError *error = nil;
id<MTLFunction> fn = [g_library newFunctionWithName:
[NSString stringWithUTF8String:qwen4_kernel_names[kernel]]
constantValues:constants error:&error];
if (fn) pipeline = [g_device newComputePipelineStateWithFunction:fn error:&error];
if (!pipeline) {
fprintf(stderr, "ds4: Qwen MoE pipeline failed: %s\n",
[[error localizedDescription] UTF8String]);
return 0;
}
[g_pipeline_cache setObject:pipeline forKey:key];
}
} else {
if (!g_qwen4_pipelines[kernel]) {
g_qwen4_pipelines[kernel] = ds4_gpu_get_pipeline(qwen4_kernel_names[kernel]);
if (!g_qwen4_pipelines[kernel]) {
static int source_skew_note;
if (!source_skew_note) {
source_skew_note = 1;
fprintf(stderr,
"ds4: Qwen3.8 kernel '%s' is absent from the compiled Metal sources.\n"
"ds4: Metal sources are loaded from ./metal at startup, so a binary newer\n"
"ds4: than the working directory's metal/ files (e.g. run from an older\n"
"ds4: checkout) produces this error. Run from the matching checkout, update\n"
"ds4: metal/, or point DS4_METAL_QWEN4_SOURCE at a current qwen4.metal.\n",
qwen4_kernel_names[kernel]);
}
return 0;
}
}
pipeline = g_qwen4_pipelines[kernel];
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:args length:args_len atIndex:0];
for (int i = 0; i < n_binds; i++) {
[enc setBuffer:binds[i].buf offset:binds[i].off atIndex:(NSUInteger)(i + 1)];
}
if (tg_mem) [enc setThreadgroupMemoryLength:tg_mem atIndex:0];
if (n_resident) {
id<MTLResource> res[QWEN4_ATTN_ROWS_MAX * 5u];
if (n_resident > sizeof(res) / sizeof(res[0])) return 0;
for (uint32_t i = 0; i < n_resident; i++) res[i] = resident[i].buf;
[enc useResources:res count:n_resident usage:MTLResourceUsageRead | MTLResourceUsageWrite];
}
[enc dispatchThreadgroups:grid threadsPerThreadgroup:tg];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, qwen4_kernel_names[kernel]);
}
}
static int qwen4_dispatch(int kernel, const void *args, size_t args_len,
const qwen4_bind *binds, int n_binds,
MTLSize grid, MTLSize tg, NSUInteger tg_mem) {
return qwen4_dispatch_resident(kernel, args, args_len, binds, n_binds, grid, tg, tg_mem, NULL, 0u);
}
/* bytes of one hc mixer row of n elements: f16, f32 or q8_0 */
static uint64_t qwen4_hc_row_bytes(uint32_t weight_type, uint64_t n) {
return weight_type == 1u ? n * 2u : weight_type == 0u ? n * 4u : weight_type == 8u ? (n / 32u) * 34u : 0u;
}
static int qwen4_hc_kernel(uint32_t weight_type, int k16, int k32, int kq8) {
return weight_type == 1u ? k16 : weight_type == 0u ? k32 : kq8;
}
static uint32_t qwen4_expert_row_bytes(uint32_t weight_type, uint32_t in_dim) {
if (in_dim == 0 || (in_dim % 32u) != 0) return 0;
switch (weight_type) {
case 8u: return (in_dim / 32u) * 34u; /* q8_0 */
case 12u: return (in_dim % 256u) ? 0u : (in_dim / 256u) * 144u; /* q4_K */
case 10u: return (in_dim % 256u) ? 0u : (in_dim / 256u) * 84u; /* q2_K */
case 16u: return (in_dim % 256u) ? 0u : (in_dim / 256u) * 66u; /* iq2_xxs */
case 39u: return (in_dim / 32u) * 17u; /* mxfp4 */
case 2u: return (in_dim / 32u) * 18u; /* q4_0 */
case 1u: return in_dim * 2u; /* f16 */
case 30u: return in_dim * 2u; /* bf16 */
case 0u: return in_dim * 4u; /* f32 */
default: return 0;
}
}
int ds4_gpu_qwen4_decode_fusions_enabled(void) {
const int override = ds4_gpu_env_bool("DS4_QWEN4_DECODE_FUSIONS");
return override >= 0 ? override : ds4_gpu_device_name_contains("M3 Ultra");
}
int ds4_gpu_qwen4_hc_combine_norm_tensor(
ds4_gpu_tensor *next_R, const ds4_gpu_tensor *blk, const ds4_gpu_tensor *old_inj,
ds4_gpu_tensor *xn, ds4_gpu_tensor *inj_part, const ds4_gpu_tensor *R,
const void *model_map, uint64_t model_size, uint64_t gamma_offset, uint64_t inject_offset,
uint32_t weight_type, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc, uint32_t n_inject, float eps) {
const uint64_t dim = (uint64_t)n_embd * n_hc;
const uint64_t wrow = qwen4_hc_row_bytes(weight_type, dim);
struct { uint32_t n_tokens, n_embd, n_hc, n_inject; float eps; uint32_t pad0, pad1, pad2; } args =
{ n_tokens, n_embd, n_hc, n_inject, eps, 0, 0, 0 };
qwen4_bind b[8];
if (n_tokens != 1u || weight_type != 1u || n_inject != n_hc) return 0;
if (ds4_gpu_tensor_buffer(next_R) == ds4_gpu_tensor_buffer(R) ||
ds4_gpu_tensor_buffer(inj_part) == ds4_gpu_tensor_buffer(old_inj)) return 0;
if (n_tokens == 0 || n_embd == 0 || n_hc == 0 || n_hc > 8 || n_inject > 4 || wrow == 0 || (dim % 32) != 0 ||
!qwen4_bind_tensor(&b[0], R, n_tokens * dim * sizeof(float), "hc norm input") ||
!qwen4_bind_weight(&b[1], model_map, model_size, gamma_offset, dim * sizeof(float), "hc norm gamma") ||
!qwen4_bind_tensor(&b[3], xn, n_tokens * dim * sizeof(float), "hc norm output")) {
return 0;
}
if (n_inject) {
if (!qwen4_bind_weight(&b[2], model_map, model_size, inject_offset, (uint64_t)n_inject * wrow, "hc inject") ||
!qwen4_bind_tensor(&b[4], inj_part,
(uint64_t)n_tokens * n_hc * DS4_QWEN4_HC_CHUNKS * n_inject * sizeof(float),
"hc inject partials")) {
return 0;
}
} else {
b[2] = b[1];
b[4] = b[3];
}
if (!qwen4_bind_tensor(&b[5], next_R, dim * sizeof(float), "combine norm next R") ||
!qwen4_bind_tensor(&b[6], blk, n_embd * sizeof(float), "combine norm block") ||
!qwen4_bind_tensor(&b[7], old_inj, n_hc * DS4_QWEN4_HC_CHUNKS * n_hc * sizeof(float), "combine norm old inject")) return 0;
return qwen4_dispatch(QWEN4_K_HC_COMBINE_NORM, &args, sizeof(args), b, 8,
MTLSizeMake(n_hc * DS4_QWEN4_HC_CHUNKS, 1, 1), MTLSizeMake(128, 1, 1), 0);
}
int ds4_gpu_qwen4_hc_norm_tensor(
ds4_gpu_tensor *xn, ds4_gpu_tensor *inj_part, const ds4_gpu_tensor *R,
const void *model_map, uint64_t model_size, uint64_t gamma_offset, uint64_t inject_offset,
uint32_t weight_type, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc, uint32_t n_inject, float eps) {
const uint64_t dim = (uint64_t)n_embd * n_hc;
const uint64_t wrow = qwen4_hc_row_bytes(weight_type, dim);
struct { uint32_t n_tokens, n_embd, n_hc, n_inject; float eps; uint32_t pad0, pad1, pad2; } args =
{ n_tokens, n_embd, n_hc, n_inject, eps, 0, 0, 0 };
qwen4_bind b[5];
if (n_tokens == 0 || n_embd == 0 || n_hc == 0 || n_hc > 8 || n_inject > 4 || wrow == 0 || (dim % 32) != 0 ||
!qwen4_bind_tensor(&b[0], R, n_tokens * dim * sizeof(float), "hc norm input") ||
!qwen4_bind_weight(&b[1], model_map, model_size, gamma_offset, dim * sizeof(float), "hc norm gamma") ||
!qwen4_bind_tensor(&b[3], xn, n_tokens * dim * sizeof(float), "hc norm output")) {
return 0;
}
if (n_inject) {
if (!qwen4_bind_weight(&b[2], model_map, model_size, inject_offset, (uint64_t)n_inject * wrow, "hc inject") ||
!qwen4_bind_tensor(&b[4], inj_part,
(uint64_t)n_tokens * n_hc * DS4_QWEN4_HC_CHUNKS * n_inject * sizeof(float),
"hc inject partials")) {
return 0;
}
} else {
b[2] = b[1];
b[4] = b[3];
}
/* Reuse wins at the measured large prefill shape. Decode and MTP keep
* the original chunk parallelism, including with an explicit override. */
bool reuse = false;
if (n_tokens > 2u) {
const char *reuse_env = getenv("DS4_QWEN4_HC_NORM_REUSE");
if (reuse_env != NULL && strcmp(reuse_env, "1") == 0) {
reuse = true;
} else if (reuse_env == NULL || strcmp(reuse_env, "0") != 0) {
reuse = n_tokens >= 8192u && n_embd == 2560u && n_hc == 4u && n_inject == 4u &&
(ds4_gpu_device_name_contains("M3 Ultra") || ds4_gpu_device_is_m5_apple_silicon());
}
}
const int kernel = reuse
? qwen4_hc_kernel(weight_type, QWEN4_K_HC_NORM_REUSE_F16, QWEN4_K_HC_NORM_REUSE_F32, QWEN4_K_HC_NORM_REUSE_Q8)
: qwen4_hc_kernel(weight_type, QWEN4_K_HC_NORM_F16, QWEN4_K_HC_NORM_F32, QWEN4_K_HC_NORM_Q8);
return qwen4_dispatch(kernel, &args, sizeof(args), b, 5,
MTLSizeMake(reuse ? n_hc : n_hc * DS4_QWEN4_HC_CHUNKS, n_tokens, 1), MTLSizeMake(128, 1, 1), 0);
}
int ds4_gpu_qwen4_hc_gate_mix_tensor(
ds4_gpu_tensor *mixed, const ds4_gpu_tensor *xn, const ds4_gpu_tensor *lo,
const void *model_map, uint64_t model_size, uint64_t up_offset,
uint32_t weight_type, uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc, uint32_t n_rank) {
const uint64_t dim = (uint64_t)n_embd * n_hc;
const uint64_t wrow = qwen4_hc_row_bytes(weight_type, n_rank);
struct { uint32_t n_tokens, n_embd, n_hc, n_rank; } args = { n_tokens, n_embd, n_hc, n_rank };
qwen4_bind b[4];
if (n_tokens == 0 || n_hc != 4 || (n_rank % 4) != 0 || wrow == 0 || (weight_type == 8u && (n_rank % 32) != 0) ||
!qwen4_bind_tensor(&b[0], xn, n_tokens * dim * sizeof(float), "hc gate input") ||
!qwen4_bind_tensor(&b[1], lo, (uint64_t)n_tokens * n_rank * sizeof(float), "hc gate lowrank") ||
!qwen4_bind_weight(&b[2], model_map, model_size, up_offset, dim * wrow, "hc up") ||
!qwen4_bind_tensor(&b[3], mixed, (uint64_t)n_tokens * n_embd * sizeof(float), "hc mixed")) {
return 0;
}
const bool pair = n_tokens == 2u && getenv("DS4_QWEN4_NO_HC_PAIR") == NULL;
/* Register-prefetched F16 rows (same lane order and rounding, pinned
* against the plain kernel by tests/test_qwen4_kernels.c); M5 default. */
const int prefetch_override = ds4_gpu_env_bool("DS4_QWEN4_HC_MIX_PREFETCH");
const bool prefetch = weight_type == 1u &&
(prefetch_override >= 0 ? prefetch_override > 0 : ds4_gpu_device_is_m5_apple_silicon());
const int kernel = pair ? (prefetch ? QWEN4_K_HC_GATE_MIX_PAIR_F16_PF
: qwen4_hc_kernel(weight_type, QWEN4_K_HC_GATE_MIX_PAIR_F16,
QWEN4_K_HC_GATE_MIX_PAIR_F32, QWEN4_K_HC_GATE_MIX_PAIR_Q8))
: prefetch ? QWEN4_K_HC_GATE_MIX_F16_PF
: qwen4_hc_kernel(weight_type, QWEN4_K_HC_GATE_MIX_F16, QWEN4_K_HC_GATE_MIX_F32,
QWEN4_K_HC_GATE_MIX_Q8);
/* More independent output rows share the activated inputs in MTP.
* Keep the per-row lane mapping and reduction order unchanged. */
const uint32_t default_nsg = n_embd == 2560u && n_rank == 320u &&
ds4_gpu_device_name_contains("M3 Ultra") ? 16u : 4u;
const uint32_t nsg = pair ?
(uint32_t)ds4_gpu_env_u64("DS4_QWEN4_HC_PAIR_NSG", default_nsg, 1u, 16u) : 4u;
return qwen4_dispatch(kernel, &args, sizeof(args), b, 4,
MTLSizeMake((n_embd + nsg - 1u) / nsg, pair ? 1u : n_tokens, 1), MTLSizeMake(nsg * 32u, 1, 1),
pair ? (NSUInteger)n_rank * 2u * sizeof(float) : 0u);
}
int ds4_gpu_qwen4_hc_combine_tensor(
ds4_gpu_tensor *R, const ds4_gpu_tensor *out, const ds4_gpu_tensor *inj,
uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc) {
struct { uint32_t n_tokens, n_embd, n_hc, pad0; } args = { n_tokens, n_embd, n_hc, 0 };
qwen4_bind b[3];
if (n_tokens == 0 || n_hc == 0 || n_hc > 8 ||
!qwen4_bind_tensor(&b[0], R, (uint64_t)n_tokens * n_embd * n_hc * sizeof(float), "hc combine residual") ||
!qwen4_bind_tensor(&b[1], out, (uint64_t)n_tokens * n_embd * sizeof(float), "hc combine block output") ||
!qwen4_bind_tensor(&b[2], inj, (uint64_t)n_tokens * n_hc * DS4_QWEN4_HC_CHUNKS * n_hc * sizeof(float),
"hc combine inject")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_HC_COMBINE, &args, sizeof(args), b, 3,
MTLSizeMake((n_embd + 255) / 256, n_tokens, 1), MTLSizeMake(256, 1, 1), 0);
}
/* Diagnostic rollback is read per dispatch, including within one loaded engine. */
static bool qwen4_prefill_reuse(uint32_t n_tokens) {
if (n_tokens <= 8u) return false;
const int override = ds4_gpu_env_bool("DS4_QWEN4_PREFILL_REUSE");
return override >= 0 ? override > 0 :
n_tokens >= 1024u && ds4_gpu_device_name_contains("M3 Ultra");
}
int ds4_gpu_qwen4_conv_stream_tensor(
ds4_gpu_tensor *x, ds4_gpu_tensor *state,
const void *model_map, uint64_t model_size, uint64_t weight_offset,
uint32_t n_tokens, uint32_t n_channels, uint32_t conv_kernel, bool apply_silu) {
struct { uint32_t n_tokens, n_channels, conv_kernel, apply_silu; } args =
{ n_tokens, n_channels, conv_kernel, apply_silu ? 1u : 0u };
qwen4_bind b[3];
if (n_tokens == 0 || n_channels == 0 || conv_kernel < 2 || conv_kernel > 4 ||
!qwen4_bind_tensor(&b[0], x, (uint64_t)n_tokens * n_channels * sizeof(float), "conv input") ||
!qwen4_bind_tensor(&b[1], state, (uint64_t)(conv_kernel - 1) * n_channels * sizeof(float), "conv state") ||
!qwen4_bind_weight(&b[2], model_map, model_size, weight_offset,
(uint64_t)n_channels * conv_kernel * sizeof(float), "conv weight")) {
return 0;
}
const uint64_t blocks = ((uint64_t)n_tokens + 63u) / 64u;
const uint64_t total = blocks * (conv_kernel - 1u) * n_channels;
if (total <= (1u << 28) && qwen4_prefill_reuse(n_tokens)) {
if (!ds4_gpu_ensure_scratch_buffer(&g_qwen4_conv_halo_buffer, &g_qwen4_conv_halo_bytes,
(NSUInteger)(total * 4u), "convolution block history")) return 0;
const char *names[] = {"kernel_qwen4_conv_halo", "kernel_qwen4_conv_blocked"};
for (int stage = 0; stage < 2; stage++) {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(names[stage]);
if (!pipeline) return 0;
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&args length:sizeof(args) atIndex:0];
[enc setBuffer:b[0].buf offset:b[0].off atIndex:1];
[enc setBuffer:b[1].buf offset:b[1].off atIndex:2];
if (stage) [enc setBuffer:b[2].buf offset:b[2].off atIndex:3];
[enc setBuffer:g_qwen4_conv_halo_buffer offset:0 atIndex:stage ? 4 : 3];
const uint64_t work = stage ? blocks * n_channels : total;
[enc dispatchThreadgroups:MTLSizeMake((work + 255u) / 256u, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
if (!ds4_gpu_finish_command_buffer(cb, owned, names[stage])) return 0;
}
return 1;
}
return qwen4_dispatch(QWEN4_K_CONV_STREAM, &args, sizeof(args), b, 3,
MTLSizeMake((n_channels + 255) / 256, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_gdn_prep_tensor(
ds4_gpu_tensor *qkv, ds4_gpu_tensor *a, ds4_gpu_tensor *b,
const void *model_map, uint64_t model_size, uint64_t ssm_a_offset, uint64_t dt_bias_offset,
uint32_t n_tokens, uint32_t n_k_head, uint32_t n_v_head, uint32_t head_dim) {
const uint64_t conv_dim = (uint64_t)(2u * n_k_head + n_v_head) * head_dim;
struct { uint32_t n_tokens, n_k_head, n_v_head, head_dim; } args = { n_tokens, n_k_head, n_v_head, head_dim };
qwen4_bind bd[5];
if (n_tokens == 0 || head_dim < 32 || head_dim > 128 || (head_dim % 32) != 0 ||
!qwen4_bind_tensor(&bd[0], qkv, n_tokens * conv_dim * sizeof(float), "gdn qkv") ||
!qwen4_bind_tensor(&bd[1], a, (uint64_t)n_tokens * n_v_head * sizeof(float), "gdn alpha") ||
!qwen4_bind_tensor(&bd[2], b, (uint64_t)n_tokens * n_v_head * sizeof(float), "gdn beta") ||
!qwen4_bind_weight(&bd[3], model_map, model_size, ssm_a_offset, (uint64_t)n_v_head * sizeof(float), "ssm_a") ||
!qwen4_bind_weight(&bd[4], model_map, model_size, dt_bias_offset, (uint64_t)n_v_head * sizeof(float),
"dt_bias")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_GDN_PREP, &args, sizeof(args), bd, 5,
MTLSizeMake(n_k_head, n_tokens, 1), MTLSizeMake(32, 1, 1), 0);
}
int ds4_gpu_qwen4_gdn_scan_tensor(
ds4_gpu_tensor *out, ds4_gpu_tensor *state, const ds4_gpu_tensor *qkv,
const ds4_gpu_tensor *a, const ds4_gpu_tensor *b,
uint32_t n_tokens, uint32_t n_k_head, uint32_t n_v_head, uint32_t head_dim,
ds4_gpu_tensor *snap_state, uint32_t snap_tok,
ds4_gpu_tensor *snap2_state, uint32_t snap2_tok) {
const uint64_t conv_dim = (uint64_t)(2u * n_k_head + n_v_head) * head_dim;
struct { uint32_t n_tokens, n_k_head, n_v_head, head_dim, snap_tok, snap2_tok, pad1, pad2; } args =
{ n_tokens, n_k_head, n_v_head, head_dim,
snap_state ? snap_tok : UINT32_MAX, snap2_state ? snap2_tok : UINT32_MAX, 0, 0 };
const uint64_t state_bytes = (uint64_t)n_v_head * head_dim * head_dim * sizeof(float);
qwen4_bind bd[7];
if (n_tokens == 0 || head_dim < 32 || head_dim > 128 || (head_dim % 32) != 0 || n_k_head == 0 ||
!qwen4_bind_tensor(&bd[0], qkv, n_tokens * conv_dim * sizeof(float), "gdn scan qkv") ||
!qwen4_bind_tensor(&bd[1], a, (uint64_t)n_tokens * n_v_head * sizeof(float), "gdn scan g") ||
!qwen4_bind_tensor(&bd[2], b, (uint64_t)n_tokens * n_v_head * sizeof(float), "gdn scan beta") ||
!qwen4_bind_tensor(&bd[3], state, state_bytes, "gdn state") ||
!qwen4_bind_tensor(&bd[4], out, (uint64_t)n_tokens * n_v_head * head_dim * sizeof(float), "gdn scan output")) {
return 0;
}
if (snap_state) {
if (!qwen4_bind_tensor(&bd[5], snap_state, state_bytes, "gdn snapshot")) return 0;
} else {
bd[5] = bd[3];
}
if (snap2_state) {
if (!qwen4_bind_tensor(&bd[6], snap2_state, state_bytes, "gdn snapshot 2")) return 0;
} else {
bd[6] = bd[3];
}
const bool decode_r4 = (n_tokens <= 2u || (n_tokens == 3u && g_qwen4_verify_rows_exact)) &&
getenv("DS4_QWEN4_NO_GDN_R4") == NULL;
if (head_dim == 128u && (n_tokens > 8u || decode_r4)) {
/* SIMD groups own independent value rows, so changing their grouping
* preserves the recurrent arithmetic. M3 Ultra prefers eight groups
* for verification and four at the measured large-prefill shape. */
const uint32_t default_nsg = n_tokens == 2u && ds4_gpu_device_name_contains("M3 Ultra") ? 8u : 1u;
const uint32_t prefill_default_nsg =
n_tokens >= 8192u && n_k_head == 16u && n_v_head == 48u &&
ds4_gpu_device_name_contains("M3 Ultra") ? 4u : 1u;
const uint32_t nsg = n_tokens <= 2u ?
(uint32_t)ds4_gpu_env_u64("DS4_QWEN4_GDN_NSG", default_nsg, 1u, 8u) :
(uint32_t)ds4_gpu_env_u64("DS4_QWEN4_GDN_PREFILL_NSG", prefill_default_nsg, 1u, 8u);
return qwen4_dispatch(QWEN4_K_GDN_SCAN_R4, &args, sizeof(args), bd, 7,
MTLSizeMake((head_dim + 4u * nsg - 1u) / (4u * nsg), n_v_head, 1),
MTLSizeMake(32u * nsg, 1, 1), 0);
}
return qwen4_dispatch(QWEN4_K_GDN_SCAN, &args, sizeof(args), bd, 7,
MTLSizeMake(head_dim, n_v_head, 1), MTLSizeMake(32, 1, 1), 0);
}
static ds4_gpu_tensor *g_qwen4_gdn_slots;
/* One decode step for every row of a batch against its own slot of the state
* pool. The grid keeps the per-row kernel's shape and simd-group grouping, so
* each row's arithmetic is unchanged and the result is bit-identical; the rows
* merely share a dispatch long enough to reach steady state. */
int ds4_gpu_qwen4_gdn_scan_rows_tensor(
ds4_gpu_tensor *out, ds4_gpu_tensor *state_pool,
const ds4_gpu_tensor *qkv, const ds4_gpu_tensor *ga, const ds4_gpu_tensor *gb,
const uint32_t *slots, uint32_t n_rows,
uint32_t n_k_head, uint32_t n_v_head, uint32_t head_dim,
uint32_t state_stride, uint32_t qkv_stride, uint32_t out_stride) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (head_dim != 128u || n_rows == 0u || n_rows > 64u) return 0;
struct { uint32_t n_rows, n_k_head, n_v_head, head_dim, state_stride, qkv_stride, out_stride, pad0; } args =
{ n_rows, n_k_head, n_v_head, head_dim, state_stride, qkv_stride, out_stride, 0 };
const uint64_t slot_bytes = (uint64_t)n_rows * sizeof(uint32_t);
if (!g_qwen4_gdn_slots) {
g_qwen4_gdn_slots = ds4_gpu_tensor_alloc(64u * sizeof(uint32_t));
if (!g_qwen4_gdn_slots) return 0;
}
if (ds4_gpu_tensor_write(g_qwen4_gdn_slots, 0, slots, slot_bytes) == 0) return 0;
qwen4_bind b[6];
if (!qwen4_bind_tensor(&b[0], qkv, (uint64_t)n_rows * qkv_stride * sizeof(float), "gdn rows qkv") ||
!qwen4_bind_tensor(&b[1], ga, (uint64_t)n_rows * n_v_head * sizeof(float), "gdn rows decay") ||
!qwen4_bind_tensor(&b[2], gb, (uint64_t)n_rows * n_v_head * sizeof(float), "gdn rows beta") ||
!qwen4_bind_tensor(&b[3], state_pool, 0, "gdn rows state pool") ||
!qwen4_bind_tensor(&b[4], out, (uint64_t)n_rows * out_stride * sizeof(float), "gdn rows output") ||
!qwen4_bind_tensor(&b[5], g_qwen4_gdn_slots, slot_bytes, "gdn rows slots")) {
return 0;
}
/* Simd groups own independent value rows, so their grouping changes the
* dispatch shape without touching a row's recurrent arithmetic. */
const uint32_t nsg = 1u;
return qwen4_dispatch(QWEN4_K_GDN_SCAN_ROWS, &args, sizeof(args), b, 6,
MTLSizeMake((head_dim + 4u * nsg - 1u) / (4u * nsg), n_v_head, n_rows),
MTLSizeMake(32u * nsg, 1, 1), 0);
}
/* One decode token per batch row, each against its own slot of the history
* pool; see kernel_qwen4_conv_stream_rows. */
int ds4_gpu_qwen4_conv_stream_rows_tensor(
ds4_gpu_tensor *x, ds4_gpu_tensor *hist_pool,
const void *model_map, uint64_t model_size, uint64_t weight_offset,
const uint32_t *slots, uint32_t n_rows, uint32_t n_channels,
uint32_t conv_kernel, uint32_t state_stride, uint32_t x_stride, int apply_silu) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (n_rows == 0u || n_rows > 64u || conv_kernel < 2u || conv_kernel > 4u) return 0;
struct { uint32_t n_rows, n_channels, conv_kernel, apply_silu, state_stride, x_stride, pad0, pad1; } args =
{ n_rows, n_channels, conv_kernel, apply_silu ? 1u : 0u, state_stride, x_stride, 0, 0 };
const uint64_t slot_bytes = (uint64_t)n_rows * sizeof(uint32_t);
if (!g_qwen4_gdn_slots) {
g_qwen4_gdn_slots = ds4_gpu_tensor_alloc(64u * sizeof(uint32_t));
if (!g_qwen4_gdn_slots) return 0;
}
if (ds4_gpu_tensor_write(g_qwen4_gdn_slots, 0, slots, slot_bytes) == 0) return 0;
qwen4_bind b[4];
if (!qwen4_bind_tensor(&b[0], x, (uint64_t)n_rows * x_stride * sizeof(float), "conv rows x") ||
!qwen4_bind_tensor(&b[1], hist_pool, 0, "conv rows history pool") ||
!qwen4_bind_weight(&b[2], model_map, model_size, weight_offset,
(uint64_t)n_channels * conv_kernel * sizeof(float), "conv rows weight") ||
!qwen4_bind_tensor(&b[3], g_qwen4_gdn_slots, slot_bytes, "conv rows slots")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_CONV_STREAM_ROWS, &args, sizeof(args), b, 4,
MTLSizeMake((n_channels + 255u) / 256u, n_rows, 1), MTLSizeMake(256, 1, 1), 0);
}
/* ---- GDN rows by address: one or two tokens per session ---------------- */
static uint64_t qwen4_tensor_address(const ds4_gpu_tensor *t);
typedef struct {
uint64_t state, hist, snap_state, snap_hist;
uint32_t row0, n_tok, pad0, pad1;
} qwen4_gdn_row_entry;
int ds4_gpu_qwen4_gdn_rows_stage(ds4_gpu_tensor *table, uint64_t entry0,
const ds4_gpu_qwen4_gdn_row *rows, uint32_t n_rows) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!rows || n_rows == 0u || n_rows > QWEN4_ATTN_ROWS_MAX) return 0;
qwen4_gdn_row_entry e[QWEN4_ATTN_ROWS_MAX];
for (uint32_t i = 0; i < n_rows; i++) {
e[i].state = qwen4_tensor_address(rows[i].state);
e[i].hist = qwen4_tensor_address(rows[i].hist);
e[i].snap_state = qwen4_tensor_address(rows[i].snap_state);
e[i].snap_hist = qwen4_tensor_address(rows[i].snap_hist);
if (!e[i].state || !e[i].hist || (rows[i].snap_state && !e[i].snap_state) ||
(rows[i].snap_hist && !e[i].snap_hist) || rows[i].n_tok == 0u || rows[i].n_tok > 2u) {
return 0;
}
e[i].row0 = rows[i].row0;
e[i].n_tok = rows[i].n_tok;
e[i].pad0 = e[i].pad1 = 0u;
}
return ds4_gpu_tensor_write(table, entry0 * sizeof(e[0]), e, (uint64_t)n_rows * sizeof(e[0]));
}
static uint32_t qwen4_gdn_rows_resident(qwen4_bind *out, const ds4_gpu_qwen4_gdn_row *rows, uint32_t n_rows) {
uint32_t n = 0;
for (uint32_t i = 0; i < n_rows; i++) {
const ds4_gpu_tensor *ts[4] = { rows[i].state, rows[i].hist, rows[i].snap_state, rows[i].snap_hist };
for (int j = 0; j < 4; j++) {
if (!ts[j]) continue;
out[n].buf = ds4_gpu_tensor_buffer(ts[j]);
out[n].off = 0;
n++;
}
}
return n;
}
static bool qwen4_bind_gdn_rows(qwen4_bind *b, const ds4_gpu_tensor *table, uint64_t entry0, uint32_t n_rows) {
const uint64_t off = entry0 * sizeof(qwen4_gdn_row_entry);
if (!table || !ds4_gpu_tensor_buffer(table) ||
ds4_gpu_tensor_bytes(table) < off + (uint64_t)n_rows * sizeof(qwen4_gdn_row_entry)) {
fprintf(stderr, "ds4: Qwen3.8 GDN row table is missing or undersized\n");
return false;
}
b->buf = ds4_gpu_tensor_buffer(table);
b->off = ds4_gpu_tensor_offset(table) + (NSUInteger)off;
return true;
}
int ds4_gpu_qwen4_gdn_scan_rows2_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *qkv, const ds4_gpu_tensor *ga, const ds4_gpu_tensor *gb,
const ds4_gpu_tensor *table, uint64_t entry0, const ds4_gpu_qwen4_gdn_row *rows, uint32_t n_rows,
uint32_t n_batch_rows, uint32_t n_k_head, uint32_t n_v_head, uint32_t head_dim,
uint32_t qkv_stride, uint32_t out_stride) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (head_dim != 128u || n_rows == 0u || n_rows > QWEN4_ATTN_ROWS_MAX) return 0;
struct { uint32_t n_rows, n_k_head, n_v_head, head_dim, state_stride, qkv_stride, out_stride, pad0; } args =
{ n_rows, n_k_head, n_v_head, head_dim, 0u, qkv_stride, out_stride, 0 };
qwen4_bind b[5], res[QWEN4_ATTN_ROWS_MAX * 4u];
if (!qwen4_bind_tensor(&b[0], qkv, (uint64_t)n_batch_rows * qkv_stride * sizeof(float), "gdn rows qkv") ||
!qwen4_bind_tensor(&b[1], ga, (uint64_t)n_batch_rows * n_v_head * sizeof(float), "gdn rows decay") ||
!qwen4_bind_tensor(&b[2], gb, (uint64_t)n_batch_rows * n_v_head * sizeof(float), "gdn rows beta") ||
!qwen4_bind_tensor(&b[3], out, (uint64_t)n_batch_rows * out_stride * sizeof(float), "gdn rows output") ||
!qwen4_bind_gdn_rows(&b[4], table, entry0, n_rows)) {
return 0;
}
const uint32_t nsg = 1u;
return qwen4_dispatch_resident(QWEN4_K_GDN_SCAN_ROWS2, &args, sizeof(args), b, 5,
MTLSizeMake((head_dim + 4u * nsg - 1u) / (4u * nsg), n_v_head, n_rows),
MTLSizeMake(32u * nsg, 1, 1), 0, res, qwen4_gdn_rows_resident(res, rows, n_rows));
}
int ds4_gpu_qwen4_conv_stream_rows2_tensor(
ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset,
const ds4_gpu_tensor *table, uint64_t entry0, const ds4_gpu_qwen4_gdn_row *rows, uint32_t n_rows,
uint32_t n_batch_rows, uint32_t n_channels, uint32_t conv_kernel, uint32_t x_stride, int apply_silu) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (n_rows == 0u || n_rows > QWEN4_ATTN_ROWS_MAX || conv_kernel < 2u || conv_kernel > 4u) return 0;
struct { uint32_t n_rows, n_channels, conv_kernel, apply_silu, state_stride, x_stride, pad0, pad1; } args =
{ n_rows, n_channels, conv_kernel, apply_silu ? 1u : 0u, 0u, x_stride, 0, 0 };
qwen4_bind b[3], res[QWEN4_ATTN_ROWS_MAX * 4u];
if (!qwen4_bind_tensor(&b[0], x, (uint64_t)n_batch_rows * x_stride * sizeof(float), "conv rows x") ||
!qwen4_bind_weight(&b[1], model_map, model_size, weight_offset,
(uint64_t)n_channels * conv_kernel * sizeof(float), "conv rows weight") ||
!qwen4_bind_gdn_rows(&b[2], table, entry0, n_rows)) {
return 0;
}
return qwen4_dispatch_resident(QWEN4_K_CONV_STREAM_ROWS2, &args, sizeof(args), b, 3,
MTLSizeMake((n_channels + 255u) / 256u, n_rows, 1), MTLSizeMake(256, 1, 1), 0,
res, qwen4_gdn_rows_resident(res, rows, n_rows));
}
int ds4_gpu_qwen4_gdn_out_tensor(
ds4_gpu_tensor *o, const ds4_gpu_tensor *z,
const void *model_map, uint64_t model_size, uint64_t weight_offset,
uint32_t n_tokens, uint32_t n_head, uint32_t head_dim, float eps) {
struct { uint32_t n_tokens, n_head, head_dim; float eps; } args = { n_tokens, n_head, head_dim, eps };
qwen4_bind b[3];
const uint64_t bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
if (n_tokens == 0 || head_dim < 32 || (head_dim % 32) != 0 ||
!qwen4_bind_tensor(&b[0], o, bytes, "gdn out") ||
!qwen4_bind_tensor(&b[1], z, bytes, "gdn gate") ||
!qwen4_bind_weight(&b[2], model_map, model_size, weight_offset, (uint64_t)head_dim * sizeof(float),
"ssm_norm")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_GDN_OUT, &args, sizeof(args), b, 3,
MTLSizeMake(n_head, n_tokens, 1), MTLSizeMake(32, 1, 1), 0);
}
int ds4_gpu_qwen4_ple_gate_tensor(
ds4_gpu_tensor *gated, ds4_gpu_tensor *normed, const ds4_gpu_tensor *R,
const ds4_gpu_tensor *key, const ds4_gpu_tensor *value,
const void *model_map, uint64_t model_size,
uint64_t g_key_offset, uint64_t g_query_offset, uint64_t g_conv_offset,
uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc, float eps) {
const uint64_t dim = (uint64_t)n_embd * n_hc;
struct { uint32_t n_tokens, n_embd, n_hc; float eps; } args = { n_tokens, n_embd, n_hc, eps };
qwen4_bind b[8];
if (n_tokens == 0 ||
!qwen4_bind_tensor(&b[0], R, n_tokens * dim * sizeof(float), "ple residual") ||
!qwen4_bind_tensor(&b[1], key, n_tokens * dim * sizeof(float), "ple key") ||
!qwen4_bind_tensor(&b[2], value, (uint64_t)n_tokens * n_embd * sizeof(float), "ple value") ||
!qwen4_bind_weight(&b[3], model_map, model_size, g_key_offset, dim * sizeof(float), "ple norm_key") ||
!qwen4_bind_weight(&b[4], model_map, model_size, g_query_offset, dim * sizeof(float), "ple norm_query") ||
!qwen4_bind_weight(&b[5], model_map, model_size, g_conv_offset, dim * sizeof(float), "ple norm_conv") ||
!qwen4_bind_tensor(&b[6], gated, n_tokens * dim * sizeof(float), "ple gated") ||
!qwen4_bind_tensor(&b[7], normed, n_tokens * dim * sizeof(float), "ple normed")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_PLE_GATE, &args, sizeof(args), b, 8,
MTLSizeMake(n_tokens, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_ple_conv_tensor(
ds4_gpu_tensor *R, const ds4_gpu_tensor *gated, const ds4_gpu_tensor *normed,
ds4_gpu_tensor *history, const void *model_map, uint64_t model_size, uint64_t weight_offset,
uint32_t weight_type, uint32_t n_tokens, uint32_t n_channels, uint32_t conv_kernel, uint32_t dilation,
ds4_gpu_tensor *snap_history, uint32_t snap_tok,
ds4_gpu_tensor *snap2_history, uint32_t snap2_tok) {
const uint32_t wsize = weight_type == 1u ? 2u : weight_type == 0u ? 4u : 0u;
struct { uint32_t n_tokens, n_channels, conv_kernel, dilation, weight_f16, snap_tok, snap2_tok, pad2; } args =
{ n_tokens, n_channels, conv_kernel, dilation, weight_type == 1u ? 1u : 0u,
snap_history ? snap_tok : UINT32_MAX, snap2_history ? snap2_tok : UINT32_MAX, 0 };
if (wsize == 0) return 0;
const uint64_t rows = (uint64_t)n_tokens * n_channels * sizeof(float);
const uint32_t hist = (conv_kernel - 1u) * dilation;
qwen4_bind b[7];
if (n_tokens == 0 || conv_kernel < 2 || conv_kernel > 4 || dilation == 0 || hist > 9 ||
!qwen4_bind_tensor(&b[0], R, rows, "ple residual") ||
!qwen4_bind_tensor(&b[1], gated, rows, "ple gated") ||
!qwen4_bind_tensor(&b[2], normed, rows, "ple normed") ||
!qwen4_bind_tensor(&b[3], history, (uint64_t)hist * n_channels * sizeof(float), "ple history") ||
!qwen4_bind_weight(&b[4], model_map, model_size, weight_offset,
(uint64_t)n_channels * conv_kernel * wsize, "ple conv weight")) {
return 0;
}
if (snap_history) {
if (!qwen4_bind_tensor(&b[5], snap_history,
(uint64_t)(conv_kernel - 1) * dilation * n_channels * sizeof(float), "ple snapshot")) {
return 0;
}
} else {
b[5] = b[3];
}
if (snap2_history) {
if (!qwen4_bind_tensor(&b[6], snap2_history,
(uint64_t)(conv_kernel - 1) * dilation * n_channels * sizeof(float), "ple snapshot 2")) {
return 0;
}
} else {
b[6] = b[3];
}
return qwen4_dispatch(QWEN4_K_PLE_CONV, &args, sizeof(args), b, 7,
MTLSizeMake((n_channels + 255) / 256, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_router_topk_tensor(
ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, const ds4_gpu_tensor *logits,
const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t gate_offset,
uint32_t gate_type, uint32_t in_dim, ds4_gpu_tensor *shared_gate,
uint32_t n_tokens, uint32_t n_expert, uint32_t n_used) {
const uint32_t gate_bytes = in_dim ? qwen4_expert_row_bytes(gate_type, in_dim) : 0u;
struct { uint32_t n_tokens, n_expert, n_used, gate_type, in_dim, pad0, pad1, pad2; } args =
{ n_tokens, n_expert, n_used, gate_type, in_dim, 0, 0, 0 };
qwen4_bind b[6];
if (n_tokens == 0 || n_expert == 0 || n_expert > 512 || n_used == 0 || n_used > n_expert || n_used > 16 ||
(in_dim && gate_bytes == 0) ||
!qwen4_bind_tensor(&b[0], logits, (uint64_t)n_tokens * n_expert * sizeof(float), "router logits") ||
!qwen4_bind_tensor(&b[1], selected, (uint64_t)n_tokens * n_used * sizeof(int32_t), "router selected") ||
!qwen4_bind_tensor(&b[2], weights, (uint64_t)n_tokens * n_used * sizeof(float), "router weights")) {
return 0;
}
if (in_dim) {
if (!qwen4_bind_tensor(&b[3], x, (uint64_t)n_tokens * in_dim * sizeof(float), "router input") ||
!qwen4_bind_weight(&b[4], model_map, model_size, gate_offset, gate_bytes, "shared gate row") ||
!qwen4_bind_tensor(&b[5], shared_gate, (uint64_t)n_tokens * sizeof(float), "shared gate logit")) {
return 0;
}
} else {
b[3] = b[0];
b[4] = b[0];
b[5] = b[2];
}
return qwen4_dispatch(QWEN4_K_ROUTER, &args, sizeof(args), b, 6,
MTLSizeMake(n_tokens, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
static float g_qwen4_rope_freq[32];
static float g_qwen4_rope_mscale = 1.0f;
static bool g_qwen4_rope_set = false;
void ds4_gpu_qwen4_set_rope(const float *freq, uint32_t n_pairs, float mscale) {
memset(g_qwen4_rope_freq, 0, sizeof(g_qwen4_rope_freq));
for (uint32_t i = 0; freq && i < n_pairs && i < 32u; i++) g_qwen4_rope_freq[i] = freq[i];
g_qwen4_rope_mscale = freq ? mscale : 1.0f;
g_qwen4_rope_set = freq != NULL;
}
static void qwen4_rope_fill(float *freq, float *mscale, uint32_t n_rot, float base) {
for (uint32_t i = 0; i < 32u; i++) {
freq[i] = g_qwen4_rope_set ? g_qwen4_rope_freq[i] :
i < n_rot / 2u ? powf(base, -2.0f * (float)i / (float)n_rot) : 0.0f;
}
*mscale = g_qwen4_rope_set ? g_qwen4_rope_mscale : 1.0f;
}
int ds4_gpu_qwen4_attn_prep_tensor(
ds4_gpu_tensor *q_out, ds4_gpu_tensor *gate_out, ds4_gpu_tensor *k_cache, ds4_gpu_tensor *v_cache,
ds4_gpu_tensor *iq_out, ds4_gpu_tensor *ik_cache,
const ds4_gpu_tensor *qg, const ds4_gpu_tensor *kproj, const ds4_gpu_tensor *vproj,
const ds4_gpu_tensor *iq, const ds4_gpu_tensor *ik, const ds4_gpu_tensor *pos3,
const void *model_map, uint64_t model_size,
uint64_t g_q_offset, uint64_t g_k_offset, uint64_t g_iq_offset,
uint32_t n_tokens, uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim, uint32_t n_rot,
uint32_t n_idx_head, uint32_t idx_dim, uint32_t pos0, uint32_t cache_cap,
float rope_base, float eps) {
struct {
uint32_t n_tokens, n_head, n_head_kv, head_dim, n_rot, n_idx_head, idx_dim, pos0, cache_cap;
float rope_base, eps; uint32_t pad0; float rope_mscale; float rope_freq[32];
} args = { n_tokens, n_head, n_head_kv, head_dim, n_rot, n_idx_head, idx_dim, pos0, cache_cap,
rope_base, eps, 0, 1.0f, { 0 } };
qwen4_rope_fill(args.rope_freq, &args.rope_mscale, n_rot, rope_base);
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
const uint64_t kv_bytes = (uint64_t)n_tokens * n_head_kv * head_dim * sizeof(float);
const uint64_t iq_bytes = (uint64_t)n_tokens * n_idx_head * idx_dim * sizeof(float);
const uint64_t cache_bytes = (uint64_t)cache_cap * n_head_kv * head_dim * 2u;
qwen4_bind b[15];
if (n_tokens == 0 || head_dim < 32 || head_dim > 256 || (head_dim % 32) != 0 ||
idx_dim < 32 || idx_dim > 128 || (idx_dim % 32) != 0 || n_rot > 64 || (n_rot % 2) != 0 ||
n_rot > idx_dim || (uint64_t)pos0 + n_tokens > cache_cap || n_head_kv == 0 || (n_head % n_head_kv) != 0 ||
!qwen4_bind_tensor(&b[0], qg, 2u * q_bytes, "attn q/gate projection") ||
!qwen4_bind_tensor(&b[1], kproj, kv_bytes, "attn k projection") ||
!qwen4_bind_tensor(&b[2], vproj, kv_bytes, "attn v projection") ||
!qwen4_bind_tensor(&b[3], iq, iq_bytes, "indexer q projection") ||
!qwen4_bind_tensor(&b[4], ik, (uint64_t)n_tokens * idx_dim * sizeof(float), "indexer k projection") ||
!qwen4_bind_weight(&b[5], model_map, model_size, g_q_offset, (uint64_t)head_dim * sizeof(float),
"attn q_norm") ||
!qwen4_bind_weight(&b[6], model_map, model_size, g_k_offset, (uint64_t)head_dim * sizeof(float),
"attn k_norm") ||
!qwen4_bind_weight(&b[7], model_map, model_size, g_iq_offset, (uint64_t)idx_dim * sizeof(float),
"indexer q_norm") ||
!qwen4_bind_tensor(&b[8], q_out, q_bytes, "attn q") ||
!qwen4_bind_tensor(&b[9], gate_out, q_bytes, "attn gate") ||
!qwen4_bind_tensor(&b[10], k_cache, cache_bytes, "k cache") ||
!qwen4_bind_tensor(&b[11], v_cache, cache_bytes, "v cache") ||
!qwen4_bind_tensor(&b[12], iq_out, iq_bytes, "indexer q") ||
!qwen4_bind_tensor(&b[13], ik_cache, (uint64_t)cache_cap * idx_dim * sizeof(float), "indexer k cache") ||
!qwen4_bind_tensor(&b[14], pos3, (uint64_t)cache_cap * 16u, "rope positions")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_ATTN_PREP, &args, sizeof(args), b, 15,
MTLSizeMake(n_head + n_head_kv + n_idx_head + 1u, n_tokens, 1), MTLSizeMake(32, 1, 1), 0);
}
int ds4_gpu_qwen4_idx_block_key_tensor(
ds4_gpu_tensor *block_key, const ds4_gpu_tensor *ik_cache, const ds4_gpu_tensor *pos3,
const void *model_map, uint64_t model_size, uint64_t g_ik_offset,
uint32_t block0, uint32_t n_blocks, uint32_t ratio, uint32_t idx_dim, uint32_t n_rot,
float rope_base, float eps) {
struct { uint32_t block0, n_blocks, ratio, idx_dim, n_rot; float rope_base, eps; uint32_t pad0;
float rope_mscale; float rope_freq[32]; } args =
{ block0, n_blocks, ratio, idx_dim, n_rot, rope_base, eps, 0, 1.0f, { 0 } };
qwen4_rope_fill(args.rope_freq, &args.rope_mscale, n_rot, rope_base);
const uint64_t n_keys = ((uint64_t)block0 + n_blocks) * ratio;
qwen4_bind b[4];
if (n_blocks == 0 || ratio == 0 || idx_dim < 32 || idx_dim > 128 || (idx_dim % 32) != 0 || n_rot > idx_dim ||
!qwen4_bind_tensor(&b[0], ik_cache, n_keys * idx_dim * sizeof(float), "indexer k cache") ||
!qwen4_bind_weight(&b[1], model_map, model_size, g_ik_offset, (uint64_t)idx_dim * sizeof(float),
"indexer k_norm") ||
!qwen4_bind_tensor(&b[2], pos3, n_keys * 16u, "rope positions") ||
!qwen4_bind_tensor(&b[3], block_key, ((uint64_t)block0 + n_blocks) * idx_dim * 2u, "block keys")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_IDX_BLOCK_KEY, &args, sizeof(args), b, 4,
MTLSizeMake(n_blocks, 1, 1), MTLSizeMake(32, 1, 1), 0);
}
int ds4_gpu_qwen4_idx_score_tensor(
ds4_gpu_tensor *score, ds4_gpu_tensor *tile_max, const ds4_gpu_tensor *iq, const ds4_gpu_tensor *block_key,
uint32_t n_tokens, uint32_t n_blocks, uint32_t n_idx_head, uint32_t idx_dim,
uint32_t pos0, uint32_t ratio) {
struct { uint32_t n_tokens, n_blocks, n_idx_head, idx_dim, pos0, ratio, pad0, pad1; } args =
{ n_tokens, n_blocks, n_idx_head, idx_dim, pos0, ratio, 0, 0 };
qwen4_bind b[4];
if (n_tokens == 0 || n_blocks == 0 || ratio == 0 ||
!qwen4_bind_tensor(&b[0], iq, (uint64_t)n_tokens * n_idx_head * idx_dim * sizeof(float), "indexer q") ||
!qwen4_bind_tensor(&b[1], block_key, (uint64_t)n_blocks * idx_dim * 2u, "block keys") ||
!qwen4_bind_tensor(&b[2], score, (uint64_t)n_tokens * n_blocks * sizeof(float), "indexer scores")) {
return 0;
}
if (n_tokens > 2u && n_idx_head == 4u && idx_dim == 128u) {
return qwen4_dispatch(QWEN4_K_IDX_SCORE_MM, &args, sizeof(args), b, 3,
MTLSizeMake((n_blocks + 63) / 64, (n_tokens + 15) / 16, 1), MTLSizeMake(128, 1, 1), 0);
}
/* staged queries and vector key loads, plus the tile maxima the
* prefiltered selector needs; measured on M5, other devices keep the
* scalar scorer */
const int vec_override = ds4_gpu_env_bool("DS4_QWEN4_IDX_SCORE_VEC");
const bool vec = tile_max && n_idx_head * idx_dim <= 512u && (idx_dim & 3u) == 0u &&
(vec_override >= 0 ? vec_override != 0 : ds4_gpu_device_is_m5_apple_silicon());
if (vec) {
if (!qwen4_bind_tensor(&b[3], tile_max, (uint64_t)n_tokens * ((n_blocks + 7u) / 8u) * sizeof(uint32_t),
"indexer tile maxima")) return 0;
return qwen4_dispatch(QWEN4_K_IDX_SCORE_VEC, &args, sizeof(args), b, 4,
MTLSizeMake((n_blocks + 127) / 128, n_tokens, 1), MTLSizeMake(128, 1, 1), 0);
}
return qwen4_dispatch(QWEN4_K_IDX_SCORE, &args, sizeof(args), b, 3,
MTLSizeMake((n_blocks + 127) / 128, n_tokens, 1), MTLSizeMake(128, 1, 1), 0);
}
int ds4_gpu_qwen4_idx_select_tensor(
ds4_gpu_tensor *sel, const ds4_gpu_tensor *score, const ds4_gpu_tensor *tile_max,
uint32_t n_blocks, uint32_t n_tokens, uint32_t top_k) {
struct { uint32_t n_tokens, n_blocks, top_k, pad0; } args = { n_tokens, n_blocks, top_k, 0 };
qwen4_bind b[3];
if (n_tokens == 0 || top_k == 0 || top_k > n_blocks ||
!qwen4_bind_tensor(&b[0], score, (uint64_t)n_tokens * n_blocks * sizeof(float), "indexer scores") ||
!qwen4_bind_tensor(&b[2], sel, (uint64_t)n_tokens * top_k * sizeof(int32_t), "selected blocks")) {
return 0;
}
/* Long rows: bound the k-th score by the k-th tile maximum and select
* over the surviving tiles in threadgroup memory (exact; see the kernel).
* The scorer that emits tile maxima only runs for decode batches. */
const int pre_override = ds4_gpu_env_bool("DS4_QWEN4_IDX_PREFILTER");
/* Scalar/MM scorers do not populate tile maxima. An independent
* scorer override must also disable their consumer. */
const int vec_override = ds4_gpu_env_bool("DS4_QWEN4_IDX_SCORE_VEC");
const bool vec = vec_override >= 0 ? vec_override != 0 : ds4_gpu_device_is_m5_apple_silicon();
const bool pre = tile_max && vec && n_tokens <= 2u && n_blocks > 8u * top_k &&
(pre_override >= 0 ? pre_override != 0 : ds4_gpu_device_is_m5_apple_silicon());
if (pre) {
if (!qwen4_bind_tensor(&b[1], tile_max, (uint64_t)n_tokens * ((n_blocks + 7u) / 8u) * sizeof(uint32_t),
"indexer tile maxima")) return 0;
return qwen4_dispatch(QWEN4_K_IDX_SELECT_PRE, &args, sizeof(args), b, 3,
MTLSizeMake(n_tokens, 1, 1), MTLSizeMake(1024, 1, 1), 0);
}
b[1] = b[2];
return qwen4_dispatch(QWEN4_K_IDX_SELECT, &args, sizeof(args), b, 2,
MTLSizeMake(n_tokens, 1, 1), MTLSizeMake(1024, 1, 1), 0);
}
int ds4_gpu_qwen4_idx_expand_tensor(
ds4_gpu_tensor *sel_tokens, ds4_gpu_tensor *n_sel, const ds4_gpu_tensor *sel_blocks,
uint32_t n_tokens, uint32_t n_sel_blocks, uint32_t ratio, uint32_t pos0, uint32_t sel_stride) {
struct { uint32_t n_tokens, n_sel_blocks, ratio, pos0, sel_stride, pad0, pad1, pad2; } args =
{ n_tokens, n_sel_blocks, ratio, pos0, sel_stride, 0, 0, 0 };
qwen4_bind b[3];
if (n_tokens == 0 || ratio == 0 || sel_stride < n_sel_blocks * ratio + ratio - 1u ||
!qwen4_bind_tensor(&b[0], sel_blocks, (uint64_t)n_tokens * n_sel_blocks * sizeof(int32_t), "selected blocks") ||
!qwen4_bind_tensor(&b[1], sel_tokens, (uint64_t)n_tokens * sel_stride * sizeof(int32_t), "selected tokens") ||
!qwen4_bind_tensor(&b[2], n_sel, (uint64_t)n_tokens * sizeof(uint32_t), "selected counts")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_IDX_EXPAND, &args, sizeof(args), b, 3,
MTLSizeMake(n_tokens, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
/* keys per split of the partial-softmax decode path (DS4_QWEN4_ATTN_SPLIT_KEYS overrides) */
/* Read per call, not cached, so an A/B harness can switch it per step. */
static uint32_t qwen4_attn_split_keys(void) {
const char *env = getenv("DS4_QWEN4_ATTN_SPLIT_KEYS");
const int v = env ? atoi(env) : 0;
return v > 0 ? (uint32_t)v : 32u;
}
uint64_t ds4_gpu_qwen4_attn_part_floats(uint32_t n_tokens, uint32_t n_head, uint32_t head_dim) {
return (uint64_t)n_tokens * n_head * QWEN4_ATTN_MAX_SPLITS * (2u + head_dim);
}
int ds4_gpu_qwen4_attn_decode_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *q, const ds4_gpu_tensor *gate,
const ds4_gpu_tensor *k_cache, const ds4_gpu_tensor *v_cache,
const ds4_gpu_tensor *sel_tokens, const ds4_gpu_tensor *n_sel, ds4_gpu_tensor *part,
uint32_t n_tokens, uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim,
uint32_t pos0, bool use_sel, uint32_t sel_stride, float scale) {
const uint32_t n_keys = use_sel ? sel_stride : pos0 + n_tokens;
uint32_t n_splits = 1;
if (part) {
n_splits = (n_keys + qwen4_attn_split_keys() - 1) / qwen4_attn_split_keys();
if (n_splits < 1) n_splits = 1;
if (n_splits > QWEN4_ATTN_MAX_SPLITS) n_splits = QWEN4_ATTN_MAX_SPLITS;
}
const uint32_t keys_per_split = (n_keys + n_splits - 1) / n_splits;
struct { uint32_t n_tokens, n_head, n_head_kv, head_dim, pos0, use_sel, sel_stride; float scale;
uint32_t n_splits, keys_per_split, pad0, pad1; } args =
{ n_tokens, n_head, n_head_kv, head_dim, pos0, use_sel ? 1u : 0u, sel_stride, scale,
n_splits, keys_per_split, 0, 0 };
const uint64_t q_bytes = (uint64_t)n_tokens * n_head * head_dim * sizeof(float);
const uint64_t cache_bytes = (uint64_t)(pos0 + n_tokens) * n_head_kv * head_dim * 2u;
const uint64_t part_bytes = (uint64_t)n_tokens * n_head * n_splits * (2u + head_dim) * sizeof(float);
qwen4_bind b[8];
if (n_tokens == 0 || n_head_kv == 0 || (n_head % n_head_kv) != 0 || n_head / n_head_kv > 12 ||
(head_dim != 32 && head_dim != 128 && head_dim != 256) ||
!qwen4_bind_tensor(&b[0], q, q_bytes, "attn q") ||
!qwen4_bind_tensor(&b[1], gate, q_bytes, "attn gate") ||
!qwen4_bind_tensor(&b[2], k_cache, cache_bytes, "k cache") ||
!qwen4_bind_tensor(&b[3], v_cache, cache_bytes, "v cache") ||
!qwen4_bind_tensor(&b[6], out, q_bytes, "attn out")) {
return 0;
}
if (use_sel) {
if (!qwen4_bind_tensor(&b[4], sel_tokens, (uint64_t)n_tokens * sel_stride * sizeof(int32_t),
"selected tokens") ||
!qwen4_bind_tensor(&b[5], n_sel, (uint64_t)n_tokens * sizeof(uint32_t), "selected counts")) {
return 0;
}
} else {
b[4] = b[0];
b[5] = b[0];
}
if (n_splits > 1) {
if (!qwen4_bind_tensor(&b[7], part, part_bytes, "attn partials")) return 0;
} else {
b[7] = b[6];
}
/* Keep short prefill tails on the decode arithmetic. Rounding queries
* to half for just a few rows can flip nearly tied expert selections
* relative to one-token continuation; large prefills retain matrix tiles. */
if (n_splits == 1 && n_tokens > 8u && head_dim == 256u && n_head / n_head_kv <= 16u &&
getenv("DS4_QWEN4_NO_ATTN_MM") == NULL) {
return qwen4_dispatch(QWEN4_K_ATTN_MM, &args, sizeof(args), b, 7,
MTLSizeMake(n_head_kv, n_tokens, 1), MTLSizeMake(128, 1, 1), 0);
}
const int kd = head_dim == 256u ? QWEN4_K_ATTN_DECODE_NPT8 :
head_dim == 128u ? QWEN4_K_ATTN_DECODE_NPT4 : QWEN4_K_ATTN_DECODE_NPT1;
const int km = head_dim == 256u ? QWEN4_K_ATTN_MERGE_NPT8 :
head_dim == 128u ? QWEN4_K_ATTN_MERGE_NPT4 : QWEN4_K_ATTN_MERGE_NPT1;
if (!qwen4_dispatch(kd, &args, sizeof(args), b, 8,
MTLSizeMake(n_splits, n_head_kv, n_tokens), MTLSizeMake(32 * QWEN4_ATTN_NSG, 1, 1), 0)) {
return 0;
}
if (n_splits == 1) return 1;
qwen4_bind mb[3] = { b[7], b[1], b[6] };
/* One thread per dim merges the same split chain with eight times the
* threads; measured on M5, other devices keep the simdgroup merge. */
const int wide_override = ds4_gpu_env_bool("DS4_QWEN4_ATTN_MERGE_WIDE");
const bool wide = head_dim >= 128u &&
(wide_override >= 0 ? wide_override != 0 : ds4_gpu_device_is_m5_apple_silicon());
if (wide) {
return qwen4_dispatch(head_dim == 256u ? QWEN4_K_ATTN_MERGE_WIDE_NPT8 : QWEN4_K_ATTN_MERGE_WIDE_NPT4,
&args, sizeof(args), mb, 3,
MTLSizeMake(n_head, n_tokens, 1), MTLSizeMake(head_dim, 1, 1), 0);
}
return qwen4_dispatch(km, &args, sizeof(args), mb, 3,
MTLSizeMake(n_head, n_tokens, 1), MTLSizeMake(32, 1, 1), 0);
}
/* ---- decode batch: attention over rows ----
*
* The caches stay per session, so the rows kernels find each row's through
* a table of GPU addresses (one entry per row, staged per layer) and the
* dispatch makes those buffers resident. Every kernel keeps the single-row
* arithmetic, so a row's outputs match the per-row dispatches bit for bit. */
typedef struct {
uint64_t k_cache, v_cache, ik_cache, block_key, pos3;
uint32_t pos, n_blocks, use_sel, pad0;
} qwen4_attn_row_entry;
static uint64_t qwen4_tensor_address(const ds4_gpu_tensor *t) {
return t ? ds4_gpu_buffer_address(ds4_gpu_tensor_buffer(t), ds4_gpu_tensor_offset(t)) : 0u;
}
int ds4_gpu_qwen4_attn_rows_stage(ds4_gpu_tensor *table, uint64_t entry0,
const ds4_gpu_qwen4_attn_row *rows, uint32_t n_rows, uint32_t ratio) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if (!rows || n_rows == 0u || n_rows > QWEN4_ATTN_ROWS_MAX || ratio == 0u) return 0;
qwen4_attn_row_entry e[QWEN4_ATTN_ROWS_MAX];
for (uint32_t i = 0; i < n_rows; i++) {
e[i].k_cache = qwen4_tensor_address(rows[i].k_cache);
e[i].v_cache = qwen4_tensor_address(rows[i].v_cache);
e[i].ik_cache = qwen4_tensor_address(rows[i].ik_cache);
e[i].block_key = qwen4_tensor_address(rows[i].block_key);
e[i].pos3 = qwen4_tensor_address(rows[i].pos3);
if (!e[i].k_cache || !e[i].v_cache || !e[i].ik_cache || !e[i].block_key || !e[i].pos3) {
fprintf(stderr, "ds4: Qwen3.8 attention rows need GPU buffer addresses (macOS 13+)\n");
return 0;
}
e[i].pos = rows[i].pos;
e[i].n_blocks = (rows[i].pos + 1u) / ratio;
e[i].use_sel = rows[i].use_sel ? 1u : 0u;
e[i].pad0 = 0u;
}
return ds4_gpu_tensor_write(table, entry0 * sizeof(e[0]), e, (uint64_t)n_rows * sizeof(e[0]));
}
static bool qwen4_bind_rows(qwen4_bind *b, const ds4_gpu_tensor *table, uint64_t entry0, uint32_t n_rows) {
const uint64_t off = entry0 * sizeof(qwen4_attn_row_entry);
if (!table || !ds4_gpu_tensor_buffer(table) ||
ds4_gpu_tensor_bytes(table) < off + (uint64_t)n_rows * sizeof(qwen4_attn_row_entry)) {
fprintf(stderr, "ds4: Qwen3.8 attention row table is missing or undersized\n");
return false;
}
b->buf = ds4_gpu_tensor_buffer(table);
b->off = ds4_gpu_tensor_offset(table) + (NSUInteger)off;
return true;
}
static uint32_t qwen4_rows_resident(qwen4_bind *out, const ds4_gpu_qwen4_attn_row *rows, uint32_t n_rows) {
uint32_t n = 0;
for (uint32_t i = 0; i < n_rows; i++) {
const ds4_gpu_tensor *ts[5] = { rows[i].k_cache, rows[i].v_cache, rows[i].ik_cache,
rows[i].block_key, rows[i].pos3 };
for (int j = 0; j < 5; j++) { out[n].buf = ds4_gpu_tensor_buffer(ts[j]); out[n].off = 0; n++; }
}
return n;
}
int ds4_gpu_qwen4_attn_prep_rows_tensor(
ds4_gpu_tensor *q_out, ds4_gpu_tensor *gate_out, ds4_gpu_tensor *iq_out,
const ds4_gpu_tensor *qg, const ds4_gpu_tensor *kproj, const ds4_gpu_tensor *vproj,
const ds4_gpu_tensor *iq, const ds4_gpu_tensor *ik,
const ds4_gpu_tensor *table, uint64_t entry0, const ds4_gpu_qwen4_attn_row *rows, uint32_t n_rows,
const void *model_map, uint64_t model_size,
uint64_t g_q_offset, uint64_t g_k_offset, uint64_t g_iq_offset,
uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim, uint32_t n_rot,
uint32_t n_idx_head, uint32_t idx_dim, float rope_base, float eps) {
/* The caller checks each row's position against its own cache; the
* kernel does not read cache_cap. */
struct {
uint32_t n_tokens, n_head, n_head_kv, head_dim, n_rot, n_idx_head, idx_dim, pos0, cache_cap;
float rope_base, eps; uint32_t pad0; float rope_mscale; float rope_freq[32];
} args = { n_rows, n_head, n_head_kv, head_dim, n_rot, n_idx_head, idx_dim, 0u, 0u,
rope_base, eps, 0, 1.0f, { 0 } };
qwen4_rope_fill(args.rope_freq, &args.rope_mscale, n_rot, rope_base);
const uint64_t q_bytes = (uint64_t)n_rows * n_head * head_dim * sizeof(float);
const uint64_t kv_bytes = (uint64_t)n_rows * n_head_kv * head_dim * sizeof(float);
const uint64_t iq_bytes = (uint64_t)n_rows * n_idx_head * idx_dim * sizeof(float);
qwen4_bind b[12], res[QWEN4_ATTN_ROWS_MAX * 5u];
if (n_rows == 0 || n_rows > QWEN4_ATTN_ROWS_MAX || head_dim < 32 || head_dim > 256 || (head_dim % 32) != 0 ||
idx_dim < 32 || idx_dim > 128 || (idx_dim % 32) != 0 || n_rot > 64 || (n_rot % 2) != 0 ||
n_rot > idx_dim || n_head_kv == 0 || (n_head % n_head_kv) != 0) {
return 0;
}
if (!qwen4_bind_tensor(&b[0], qg, 2u * q_bytes, "attn rows q/gate projection") ||
!qwen4_bind_tensor(&b[1], kproj, kv_bytes, "attn rows k projection") ||
!qwen4_bind_tensor(&b[2], vproj, kv_bytes, "attn rows v projection") ||
!qwen4_bind_tensor(&b[3], iq, iq_bytes, "indexer rows q projection") ||
!qwen4_bind_tensor(&b[4], ik, (uint64_t)n_rows * idx_dim * sizeof(float), "indexer rows k projection") ||
!qwen4_bind_weight(&b[5], model_map, model_size, g_q_offset, (uint64_t)head_dim * sizeof(float), "attn q_norm") ||
!qwen4_bind_weight(&b[6], model_map, model_size, g_k_offset, (uint64_t)head_dim * sizeof(float), "attn k_norm") ||
!qwen4_bind_weight(&b[7], model_map, model_size, g_iq_offset, (uint64_t)idx_dim * sizeof(float), "indexer q_norm") ||
!qwen4_bind_tensor(&b[8], q_out, q_bytes, "attn rows q") ||
!qwen4_bind_tensor(&b[9], gate_out, q_bytes, "attn rows gate") ||
!qwen4_bind_tensor(&b[10], iq_out, iq_bytes, "indexer rows q") ||
!qwen4_bind_rows(&b[11], table, entry0, n_rows)) {
return 0;
}
return qwen4_dispatch_resident(QWEN4_K_ATTN_PREP_ROWS, &args, sizeof(args), b, 12,
MTLSizeMake(n_head + n_head_kv + n_idx_head + 1u, n_rows, 1), MTLSizeMake(32, 1, 1), 0,
res, qwen4_rows_resident(res, rows, n_rows));
}
int ds4_gpu_qwen4_idx_block_key_rows_tensor(
const ds4_gpu_tensor *table, uint64_t entry0, const ds4_gpu_qwen4_attn_row *rows, uint32_t n_rows,
const void *model_map, uint64_t model_size, uint64_t g_ik_offset,
uint32_t ratio, uint32_t idx_dim, uint32_t n_rot, float rope_base, float eps) {
struct { uint32_t block0, n_blocks, ratio, idx_dim, n_rot; float rope_base, eps; uint32_t pad0;
float rope_mscale; float rope_freq[32]; } args =
{ 0u, n_rows, ratio, idx_dim, n_rot, rope_base, eps, 0, 1.0f, { 0 } };
qwen4_rope_fill(args.rope_freq, &args.rope_mscale, n_rot, rope_base);
qwen4_bind b[2], res[QWEN4_ATTN_ROWS_MAX * 5u];
if (n_rows == 0 || n_rows > QWEN4_ATTN_ROWS_MAX || ratio == 0 || idx_dim < 32 || idx_dim > 128 ||
(idx_dim % 32) != 0 || n_rot > idx_dim ||
!qwen4_bind_weight(&b[0], model_map, model_size, g_ik_offset, (uint64_t)idx_dim * sizeof(float), "indexer k_norm") ||
!qwen4_bind_rows(&b[1], table, entry0, n_rows)) {
return 0;
}
return qwen4_dispatch_resident(QWEN4_K_IDX_BLOCK_KEY_ROWS, &args, sizeof(args), b, 2,
MTLSizeMake(n_rows, 1, 1), MTLSizeMake(32, 1, 1), 0,
res, qwen4_rows_resident(res, rows, n_rows));
}
int ds4_gpu_qwen4_idx_score_rows_tensor(
ds4_gpu_tensor *score, ds4_gpu_tensor *tile_max, const ds4_gpu_tensor *iq,
const ds4_gpu_tensor *table, uint64_t entry0, const ds4_gpu_qwen4_attn_row *rows, uint32_t n_rows,
uint32_t n_block_stride, uint32_t n_idx_head, uint32_t idx_dim, uint32_t ratio) {
struct { uint32_t n_tokens, n_blocks, n_idx_head, idx_dim, pos0, ratio, pad0, pad1; } args =
{ n_rows, n_block_stride, n_idx_head, idx_dim, 0u, ratio, 0, 0 };
const uint32_t n_tiles = (n_block_stride + 7u) / 8u;
qwen4_bind b[4], res[QWEN4_ATTN_ROWS_MAX * 5u];
if (n_rows == 0 || n_rows > QWEN4_ATTN_ROWS_MAX || n_block_stride == 0 || ratio == 0 ||
n_idx_head * idx_dim > 512u || (idx_dim & 3u) != 0u ||
!qwen4_bind_tensor(&b[0], iq, (uint64_t)n_rows * n_idx_head * idx_dim * sizeof(float), "indexer rows q") ||
!qwen4_bind_tensor(&b[1], score, (uint64_t)n_rows * n_block_stride * sizeof(float), "indexer rows scores") ||
!qwen4_bind_tensor(&b[2], tile_max, (uint64_t)n_rows * n_tiles * sizeof(uint32_t), "indexer rows tile maxima") ||
!qwen4_bind_rows(&b[3], table, entry0, n_rows)) {
return 0;
}
return qwen4_dispatch_resident(QWEN4_K_IDX_SCORE_ROWS, &args, sizeof(args), b, 4,
MTLSizeMake((n_block_stride + 127) / 128, n_rows, 1), MTLSizeMake(128, 1, 1), 0,
res, qwen4_rows_resident(res, rows, n_rows));
}
int ds4_gpu_qwen4_idx_select_rows_tensor(
ds4_gpu_tensor *sel, const ds4_gpu_tensor *score, const ds4_gpu_tensor *tile_max,
const ds4_gpu_tensor *table, uint64_t entry0, const ds4_gpu_qwen4_attn_row *rows, uint32_t n_rows,
uint32_t n_block_stride, uint32_t top_k) {
struct { uint32_t n_tokens, n_blocks, top_k, pad0; } args = { n_rows, n_block_stride, top_k, 0 };
const uint32_t n_tiles = (n_block_stride + 7u) / 8u;
qwen4_bind b[4];
if (n_rows == 0 || n_rows > QWEN4_ATTN_ROWS_MAX || top_k == 0 || top_k > n_block_stride ||
!qwen4_bind_tensor(&b[0], score, (uint64_t)n_rows * n_block_stride * sizeof(float), "indexer rows scores") ||
!qwen4_bind_tensor(&b[1], tile_max, (uint64_t)n_rows * n_tiles * sizeof(uint32_t), "indexer rows tile maxima") ||
!qwen4_bind_tensor(&b[2], sel, (uint64_t)n_rows * top_k * sizeof(int32_t), "rows selected blocks") ||
!qwen4_bind_rows(&b[3], table, entry0, n_rows)) {
return 0;
}
(void)rows;
return qwen4_dispatch(QWEN4_K_IDX_SELECT_ROWS, &args, sizeof(args), b, 4,
MTLSizeMake(n_rows, 1, 1), MTLSizeMake(1024, 1, 1), 0);
}
int ds4_gpu_qwen4_idx_expand_rows_tensor(
ds4_gpu_tensor *sel_tokens, ds4_gpu_tensor *n_sel, const ds4_gpu_tensor *sel_blocks,
const ds4_gpu_tensor *table, uint64_t entry0, uint32_t n_rows,
uint32_t n_sel_blocks, uint32_t ratio, uint32_t sel_stride) {
struct { uint32_t n_tokens, n_sel_blocks, ratio, pos0, sel_stride, pad0, pad1, pad2; } args =
{ n_rows, n_sel_blocks, ratio, 0u, sel_stride, 0, 0, 0 };
qwen4_bind b[4];
if (n_rows == 0 || n_rows > QWEN4_ATTN_ROWS_MAX || ratio == 0 || sel_stride < n_sel_blocks * ratio + ratio - 1u ||
!qwen4_bind_tensor(&b[0], sel_blocks, (uint64_t)n_rows * n_sel_blocks * sizeof(int32_t), "rows selected blocks") ||
!qwen4_bind_tensor(&b[1], sel_tokens, (uint64_t)n_rows * sel_stride * sizeof(int32_t), "rows selected tokens") ||
!qwen4_bind_tensor(&b[2], n_sel, (uint64_t)n_rows * sizeof(uint32_t), "rows selected counts") ||
!qwen4_bind_rows(&b[3], table, entry0, n_rows)) {
return 0;
}
return qwen4_dispatch(QWEN4_K_IDX_EXPAND_ROWS, &args, sizeof(args), b, 4,
MTLSizeMake(n_rows, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_attn_decode_rows_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *q, const ds4_gpu_tensor *gate,
const ds4_gpu_tensor *sel_tokens, const ds4_gpu_tensor *n_sel, ds4_gpu_tensor *part,
const ds4_gpu_tensor *table, uint64_t entry0, const ds4_gpu_qwen4_attn_row *rows, uint32_t n_rows,
uint32_t n_head, uint32_t n_head_kv, uint32_t head_dim, uint32_t sel_stride, float scale) {
/* Each row takes the splits the single-row dispatch gives its key count
* (see the kernel); the partials use the split cap as their stride. */
struct { uint32_t n_tokens, n_head, n_head_kv, head_dim, pos0, use_sel, sel_stride; float scale;
uint32_t n_splits, keys_per_split, pad0, pad1; } args =
{ n_rows, n_head, n_head_kv, head_dim, 0u, 0u, sel_stride, scale,
QWEN4_ATTN_MAX_SPLITS, qwen4_attn_split_keys(), 0, 0 };
const uint64_t q_bytes = (uint64_t)n_rows * n_head * head_dim * sizeof(float);
const uint64_t part_bytes = ds4_gpu_qwen4_attn_part_floats(n_rows, n_head, head_dim) * sizeof(float);
qwen4_bind b[7], res[QWEN4_ATTN_ROWS_MAX * 5u];
if (n_rows == 0 || n_rows > QWEN4_ATTN_ROWS_MAX || n_head_kv == 0 || (n_head % n_head_kv) != 0 ||
n_head / n_head_kv > 12 || (head_dim != 128 && head_dim != 256) ||
!qwen4_bind_tensor(&b[0], q, q_bytes, "attn rows q") ||
!qwen4_bind_tensor(&b[1], gate, q_bytes, "attn rows gate") ||
!qwen4_bind_tensor(&b[2], sel_tokens, (uint64_t)n_rows * sel_stride * sizeof(int32_t), "rows selected tokens") ||
!qwen4_bind_tensor(&b[3], n_sel, (uint64_t)n_rows * sizeof(uint32_t), "rows selected counts") ||
!qwen4_bind_tensor(&b[4], out, q_bytes, "attn rows out") ||
!qwen4_bind_tensor(&b[5], part, part_bytes, "attn rows partials") ||
!qwen4_bind_rows(&b[6], table, entry0, n_rows)) {
return 0;
}
const bool npt8 = head_dim == 256u;
if (!qwen4_dispatch_resident(npt8 ? QWEN4_K_ATTN_DECODE_ROWS_NPT8 : QWEN4_K_ATTN_DECODE_ROWS_NPT4,
&args, sizeof(args), b, 7,
MTLSizeMake(QWEN4_ATTN_MAX_SPLITS, n_head_kv, n_rows),
MTLSizeMake(32 * QWEN4_ATTN_NSG, 1, 1), 0,
res, qwen4_rows_resident(res, rows, n_rows))) {
return 0;
}
qwen4_bind mb[4] = { b[5], b[1], b[4], b[6] };
return qwen4_dispatch(npt8 ? QWEN4_K_ATTN_MERGE_ROWS_NPT8 : QWEN4_K_ATTN_MERGE_ROWS_NPT4,
&args, sizeof(args), mb, 4,
MTLSizeMake(n_head, n_rows, 1), MTLSizeMake(head_dim, 1, 1), 0);
}
int ds4_gpu_qwen4_moe_mid_tensor(
ds4_gpu_tensor *mid, const ds4_gpu_tensor *x, const ds4_gpu_tensor *selected,
const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset,
uint32_t weight_type, uint32_t n_total_expert, uint32_t n_tokens, uint32_t n_slots,
uint32_t in_dim, uint32_t ff_dim,
uint64_t shared_gate_offset, uint64_t shared_up_offset, uint32_t shared_type) {
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, in_dim);
const uint64_t expert_bytes = (uint64_t)row_bytes * ff_dim;
const bool has_shared = shared_type != UINT32_MAX;
const uint32_t sh_row_bytes = has_shared ? qwen4_expert_row_bytes(shared_type, in_dim) : 0u;
const uint32_t n_out = n_slots + (has_shared ? 1u : 0u);
const uint64_t experts_bytes = expert_bytes * n_total_expert;
const uint64_t shared_bytes = (uint64_t)sh_row_bytes * ff_dim;
qwen4_moe_args args = { n_tokens, n_slots, in_dim, ff_dim, weight_type, row_bytes, expert_bytes,
has_shared ? 1u : 0u, has_shared ? shared_type : 0u, sh_row_bytes, n_total_expert, 0u, 0u };
qwen4_bind b[7];
if (n_tokens == 0 || n_slots == 0 || row_bytes == 0 || ff_dim == 0 || (has_shared && sh_row_bytes == 0) ||
!qwen4_bind_weight(&b[0], model_map, model_size, gate_offset, experts_bytes, "moe gate experts") ||
!qwen4_bind_weight(&b[1], model_map, model_size, up_offset, experts_bytes, "moe up experts") ||
!qwen4_bind_tensor(&b[2], selected, (uint64_t)n_tokens * n_slots * sizeof(int32_t), "moe selected") ||
!qwen4_bind_tensor(&b[3], x, (uint64_t)n_tokens * in_dim * sizeof(float), "moe input") ||
!qwen4_bind_tensor(&b[4], mid, (uint64_t)n_tokens * n_out * ff_dim * sizeof(float), "moe mid")) {
return 0;
}
if (has_shared) {
if (!qwen4_bind_weight(&b[5], model_map, model_size, shared_gate_offset, shared_bytes, "shared gate") ||
!qwen4_bind_weight(&b[6], model_map, model_size, shared_up_offset, shared_bytes, "shared up")) {
return 0;
}
} else {
b[5] = b[0];
b[6] = b[1];
}
const bool q4k = weight_type == 12u && getenv("DS4_QWEN4_NO_Q4K_MID") == NULL;
const bool m3_ultra = q4k && ds4_gpu_device_name_contains("M3 Ultra");
/* One row per SIMD group improves both single-token decode and the
* two-token MTP verifier on M3 Ultra without changing dot-product order.
* M5 measured the single-token case with four groups per threadgroup and
* the two-row MTP passes with four. */
const bool m5_single = q4k && (n_tokens <= 2u || (n_tokens == 3u && g_qwen4_verify_rows_exact)) &&
ds4_gpu_device_is_m5_apple_silicon();
const uint32_t default_nr = (m3_ultra && n_tokens <= 2u) || m5_single ? 1u : 2u;
const bool specialize = !q4k && qwen4_moe_mv_specialize(weight_type);
const uint64_t nr_env = q4k ?
ds4_gpu_env_u64("DS4_QWEN4_Q4K_MID_NR", default_nr, 1u, UINT64_MAX) :
(specialize ? qwen4_moe_mv_rows() : 2u);
const uint32_t nr = nr_env >= 1u && nr_env <= (q4k ? 2u : 4u) ? (uint32_t)nr_env : default_nr;
/* NR2 without an NSG override restores the former ordered dispatch. */
const uint32_t default_nsg = nr != 1u ? 2u : m3_ultra ? 8u : m5_single ? 4u : 2u;
const uint32_t nsg = q4k ?
(uint32_t)ds4_gpu_env_u64("DS4_QWEN4_Q4K_MID_NSG", default_nsg, 1u, 8u) :
(specialize ? qwen4_moe_mv_groups(weight_type) : 4u);
const uint32_t rows_per_tg = nr * nsg;
const int kernel = !q4k ? QWEN4_K_MOE_MID : nr == 1u ? QWEN4_K_MOE_MID_Q4K_NR1 : QWEN4_K_MOE_MID_Q4K;
return qwen4_dispatch(kernel, &args, sizeof(args), b, 7,
MTLSizeMake((ff_dim + rows_per_tg - 1) / rows_per_tg, n_out, n_tokens),
MTLSizeMake(32u * nsg, 1, 1), 0);
}
int ds4_gpu_qwen4_moe_down_tensor(
ds4_gpu_tensor *part, const ds4_gpu_tensor *mid, const ds4_gpu_tensor *selected,
const void *model_map, uint64_t model_size, uint64_t down_offset,
uint32_t weight_type, uint32_t n_total_expert, uint32_t n_tokens, uint32_t n_slots,
uint32_t ff_dim, uint32_t out_dim,
uint64_t shared_down_offset, uint32_t shared_type) {
const uint32_t weight_dim = weight_type == 10u ? (ff_dim + 255u) / 256u * 256u : ff_dim;
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, weight_dim);
const uint64_t expert_bytes = (uint64_t)row_bytes * out_dim;
const bool has_shared = shared_type != UINT32_MAX;
const uint32_t sh_row_bytes = has_shared ? qwen4_expert_row_bytes(shared_type, ff_dim) : 0u;
const uint32_t n_out = n_slots + (has_shared ? 1u : 0u);
const uint64_t experts_bytes = expert_bytes * n_total_expert;
const uint64_t shared_bytes = (uint64_t)sh_row_bytes * out_dim;
qwen4_moe_args args = { n_tokens, n_slots, ff_dim, out_dim, weight_type, row_bytes, expert_bytes,
has_shared ? 1u : 0u, has_shared ? shared_type : 0u, sh_row_bytes, 0u, 0u, 0u };
qwen4_bind b[5];
if (n_tokens == 0 || n_slots == 0 || row_bytes == 0 || (ff_dim % 32u) != 0 ||
out_dim == 0 || (has_shared && sh_row_bytes == 0) ||
!qwen4_bind_weight(&b[0], model_map, model_size, down_offset, experts_bytes, "moe down experts") ||
!qwen4_bind_tensor(&b[1], selected, (uint64_t)n_tokens * n_slots * sizeof(int32_t), "moe selected") ||
!qwen4_bind_tensor(&b[2], mid, (uint64_t)n_tokens * n_out * ff_dim * sizeof(float), "moe mid") ||
!qwen4_bind_tensor(&b[3], part, (uint64_t)n_tokens * n_out * out_dim * sizeof(float), "moe partial")) {
return 0;
}
if (has_shared) {
if (!qwen4_bind_weight(&b[4], model_map, model_size, shared_down_offset, shared_bytes, "shared down")) return 0;
} else {
b[4] = b[0];
}
const uint32_t nsg = qwen4_moe_mv_specialize(weight_type) ? qwen4_moe_mv_groups(weight_type) : 4u;
const uint32_t rows_per_tg = nsg * (qwen4_moe_mv_specialize(weight_type) ? qwen4_moe_mv_rows() : 2u);
/* MXFP4 rows with four blocks per lane requested ahead (same lane map and
* chain order, byte-identical); M5 default, DS4_QWEN4_MOE_DOWN_PREFETCH=0/1
* overrides on any device. */
const int prefetch_override = ds4_gpu_env_bool("DS4_QWEN4_MOE_DOWN_PREFETCH");
const bool prefetch = weight_type == 39u && (ff_dim % 32u) == 0 &&
(prefetch_override >= 0 ? prefetch_override > 0 : ds4_gpu_device_is_m5_apple_silicon());
return qwen4_dispatch(prefetch ? QWEN4_K_MOE_DOWN_MXFP4_PF : QWEN4_K_MOE_DOWN, &args, sizeof(args), b, 5,
MTLSizeMake((out_dim + rows_per_tg - 1) / rows_per_tg, n_out, n_tokens),
MTLSizeMake(32u * nsg, 1, 1), 0);
}
/* Decode batch over its distinct experts (kernel_qwen4_moe_mid_q4k_grouped):
* Q4_K gate/up, no shared slot, per-expert pair lists from
* ds4_gpu_qwen4_moe_build_lists_tensor. Same grid as the per-token kernel;
* the pairs that do not own their expert return. */
int ds4_gpu_qwen4_moe_mid_grouped_tensor(
ds4_gpu_tensor *mid, const ds4_gpu_tensor *x, const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *lists, const ds4_gpu_tensor *counts, uint32_t list_cap,
const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset,
uint32_t weight_type, uint32_t n_total_expert, uint32_t n_tokens, uint32_t n_slots,
uint32_t in_dim, uint32_t ff_dim) {
if (weight_type != 12u || (in_dim % 256u) != 0) return 0;
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, in_dim);
const uint64_t expert_bytes = (uint64_t)row_bytes * ff_dim;
qwen4_moe_args args = { n_tokens, n_slots, in_dim, ff_dim, weight_type, row_bytes, expert_bytes,
0u, 0u, 0u, n_total_expert, list_cap, 0u };
qwen4_bind b[7];
if (n_tokens == 0 || n_slots == 0 || row_bytes == 0 || ff_dim == 0 || list_cap == 0 ||
!qwen4_bind_weight(&b[0], model_map, model_size, gate_offset, expert_bytes * n_total_expert, "moe gate experts") ||
!qwen4_bind_weight(&b[1], model_map, model_size, up_offset, expert_bytes * n_total_expert, "moe up experts") ||
!qwen4_bind_tensor(&b[2], selected, (uint64_t)n_tokens * n_slots * sizeof(int32_t), "moe selected") ||
!qwen4_bind_tensor(&b[3], x, (uint64_t)n_tokens * in_dim * sizeof(float), "moe input") ||
!qwen4_bind_tensor(&b[4], mid, (uint64_t)n_tokens * n_slots * ff_dim * sizeof(float), "moe mid") ||
!qwen4_bind_tensor(&b[5], lists, (uint64_t)n_total_expert * list_cap * sizeof(int32_t), "moe lists") ||
!qwen4_bind_tensor(&b[6], counts, (uint64_t)n_total_expert * sizeof(int32_t), "moe counts")) {
return 0;
}
const uint32_t nsg = 2u, rows_per_tg = 2u * nsg;
return qwen4_dispatch(QWEN4_K_MOE_MID_Q4K_GROUPED, &args, sizeof(args), b, 7,
MTLSizeMake((ff_dim + rows_per_tg - 1) / rows_per_tg, n_slots, n_tokens),
MTLSizeMake(32u * nsg, 1, 1), 0);
}
/* The MXFP4 down rows of the same batch, on the same lists. */
int ds4_gpu_qwen4_moe_down_grouped_tensor(
ds4_gpu_tensor *part, const ds4_gpu_tensor *mid, const ds4_gpu_tensor *selected,
const ds4_gpu_tensor *lists, const ds4_gpu_tensor *counts, uint32_t list_cap,
const void *model_map, uint64_t model_size, uint64_t down_offset,
uint32_t weight_type, uint32_t n_total_expert, uint32_t n_tokens, uint32_t n_slots,
uint32_t ff_dim, uint32_t out_dim) {
if (weight_type != 39u || (ff_dim % 32u) != 0) return 0;
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, ff_dim);
const uint64_t expert_bytes = (uint64_t)row_bytes * out_dim;
qwen4_moe_args args = { n_tokens, n_slots, ff_dim, out_dim, weight_type, row_bytes, expert_bytes,
0u, 0u, 0u, n_total_expert, list_cap, 0u };
qwen4_bind b[6];
if (n_tokens == 0 || n_slots == 0 || row_bytes == 0 || out_dim == 0 || list_cap == 0 ||
!qwen4_bind_weight(&b[0], model_map, model_size, down_offset, expert_bytes * n_total_expert, "moe down experts") ||
!qwen4_bind_tensor(&b[1], selected, (uint64_t)n_tokens * n_slots * sizeof(int32_t), "moe selected") ||
!qwen4_bind_tensor(&b[2], mid, (uint64_t)n_tokens * n_slots * ff_dim * sizeof(float), "moe mid") ||
!qwen4_bind_tensor(&b[3], part, (uint64_t)n_tokens * n_slots * out_dim * sizeof(float), "moe partial") ||
!qwen4_bind_tensor(&b[4], lists, (uint64_t)n_total_expert * list_cap * sizeof(int32_t), "moe lists") ||
!qwen4_bind_tensor(&b[5], counts, (uint64_t)n_total_expert * sizeof(int32_t), "moe counts")) {
return 0;
}
const uint32_t nsg = 4u, rows_per_tg = 2u * nsg;
return qwen4_dispatch(QWEN4_K_MOE_DOWN_MXFP4_GROUPED, &args, sizeof(args), b, 6,
MTLSizeMake((out_dim + rows_per_tg - 1) / rows_per_tg, n_slots, n_tokens),
MTLSizeMake(32u * nsg, 1, 1), 0);
}
int ds4_gpu_qwen4_moe_reduce_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *part, const ds4_gpu_tensor *weights,
const ds4_gpu_tensor *shared_gate, const ds4_gpu_tensor *shared, ds4_gpu_tensor *R, const ds4_gpu_tensor *inj,
uint32_t n_tokens, uint32_t n_slots, uint32_t part_stride, uint32_t dim, uint32_t n_hc) {
const uint32_t shared_src = !shared_gate ? 0u : shared ? 2u : 1u;
struct { uint32_t n_tokens, n_slots, dim, shared_src, n_hc, part_stride, pad1, pad2; } args =
{ n_tokens, n_slots, dim, shared_src, R ? n_hc : 0u, part_stride, 0, 0 };
qwen4_bind b[7];
if (n_tokens == 0 || n_slots == 0 || dim == 0 || (R && n_hc == 0) || n_hc > 8 || part_stride < n_slots ||
(shared_src == 1u && part_stride < n_slots + 1u) ||
!qwen4_bind_tensor(&b[0], part, (uint64_t)n_tokens * part_stride * dim * sizeof(float), "moe partial") ||
!qwen4_bind_tensor(&b[1], weights, (uint64_t)n_tokens * n_slots * sizeof(float), "moe weights") ||
!qwen4_bind_tensor(&b[3], out, (uint64_t)n_tokens * dim * sizeof(float), "moe out")) {
return 0;
}
if (shared_gate) {
if (!qwen4_bind_tensor(&b[2], shared_gate, (uint64_t)n_tokens * sizeof(float), "shared gate")) return 0;
} else {
b[2] = b[1];
}
if (R) {
if (!qwen4_bind_tensor(&b[4], R, (uint64_t)n_tokens * dim * n_hc * sizeof(float), "moe residual") ||
!qwen4_bind_tensor(&b[5], inj, (uint64_t)n_tokens * n_hc * DS4_QWEN4_HC_CHUNKS * n_hc * sizeof(float),
"moe inject")) {
return 0;
}
} else {
b[4] = b[3];
b[5] = b[1];
}
if (shared_src == 2u) {
if (!qwen4_bind_tensor(&b[6], shared, (uint64_t)n_tokens * dim * sizeof(float), "shared expert output")) {
return 0;
}
} else {
b[6] = b[0];
}
return qwen4_dispatch(QWEN4_K_MOE_REDUCE, &args, sizeof(args), b, 7,
MTLSizeMake((dim + 255) / 256, n_tokens, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_moe_build_lists_tensor(
ds4_gpu_tensor *lists, ds4_gpu_tensor *counts, const ds4_gpu_tensor *selected,
uint32_t n_tokens, uint32_t n_slots, uint32_t n_expert, uint32_t list_cap) {
qwen4_moe_mm_args args = { n_tokens, n_slots, 0, 0, 0, 0, 0, list_cap, 0, n_expert, 0, 0, 0 };
qwen4_bind b[3];
if (n_tokens == 0 || n_slots == 0 || n_expert == 0 || n_expert > 512 || list_cap == 0 ||
!qwen4_bind_tensor(&b[0], selected, (uint64_t)n_tokens * n_slots * sizeof(int32_t), "moe selected") ||
!qwen4_bind_tensor(&b[1], lists, (uint64_t)n_expert * list_cap * sizeof(int32_t), "moe lists") ||
!qwen4_bind_tensor(&b[2], counts, (uint64_t)n_expert * sizeof(int32_t), "moe counts")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_MOE_BUILD_LISTS, &args, sizeof(args), b, 3,
MTLSizeMake(1, 1, 1), MTLSizeMake(512, 1, 1), 0);
}
static uint32_t qwen4_moe_mm_tiles(uint32_t n_tokens, bool mid) {
uint32_t tiles = (n_tokens + 31u) / 32u;
/* Spread large routed batches over more independent tiles on M3 Ultra
* (gate/up and down) and on M5 (gate/up; the down tiles measured flat).
* Each tile keeps the same K loop and accumulation order. */
const bool spread = ds4_gpu_device_name_contains("M3 Ultra") ||
(mid && ds4_gpu_device_is_m5_apple_silicon());
const uint32_t default_cap = spread ?
(n_tokens >= 8192u ? 32u : n_tokens >= 4096u ? 16u : 8u) : 8u;
const uint32_t cap = (uint32_t)ds4_gpu_env_u64(
mid ? "DS4_QWEN4_MOE_MID_TILES" : "DS4_QWEN4_MOE_DOWN_TILES", default_cap, 1u, 32u);
return tiles > cap ? cap : tiles;
}
static uint32_t qwen4_moe_mm_nt(uint32_t n_tokens, uint32_t type, const char *env_name) {
/* Low-bit experts receive few tokens in short prefills. Smaller token
* tiles avoid unused matrix products while retaining the same K order.
* Keep 32-token tiles for larger batches and unmeasured devices/types. */
uint32_t default_nt = 4u;
if ((type == 16u || type == 10u) && ds4_gpu_device_name_contains("M3 Ultra")) {
if (n_tokens <= 512u) default_nt = 1u;
else if (n_tokens <= 1024u) default_nt = 2u;
}
const bool down = strstr(env_name, "DOWN") != NULL;
/* Reuse each decoded down-weight tile across 64 tokens. The K loop and
* partial-tail accumulation stay unchanged; short batches keep small tiles. */
if (down && type == 10u && n_tokens >= 8192u && qwen4_prefill_reuse(n_tokens)) default_nt = 8u;
/* M5: 64-token gate/up and down tiles for the Q4_K / MXFP4 pack at
* chunk-sized batches (same tile arithmetic; remainders take the
* 8/16/32-token kernels). Measured on M5 Max, see
* speed-bench/qwen38-m5-round4.md. */
if ((type == 12u || type == 39u) && n_tokens >= 4096u && ds4_gpu_device_is_m5_apple_silicon()) default_nt = 8u;
const uint32_t nt = (uint32_t)ds4_gpu_env_u64(env_name, default_nt, 1u, 8u);
return nt == 1u || nt == 2u || nt == 4u || nt == 8u ? nt : default_nt;
}
/* Routed tiles on the Metal 4 tensor ops (drift class: the cooperative
* matmul's accumulation order differs from the simdgroup tiles; needs the
* tensor API). DS4_QWEN4_MOE_MM_NAX: 0 simdgroup tiles; 1 tensor-op
* tiles of 32 tokens; 2 tensor-op tiles of 64 tokens; 3/4 the same widths
* with the activation operand kept at full float precision (x gathered from
* the float tensor, no half rounding of x or of the mid hand-off, no
* conversion pass); 5/6 compensated tiles (64/32 tokens), which stage the
* half rounding residual of the operand beside it and run the tensor op on
* both, reaching the 3/4 accuracy while keeping the fp16 tensor-op rate.
* Unset selects 2 (the 64-token half tiles) for every tensor tier on
* devices with the tensor API: its operand rounding matches the simdgroup
* tiles and its fixture NLL sits within noise of them, so the compensated
* level stays an env-selectable accuracy option. 0 forces the simdgroup
* tiles. Returns the token tile width, 0 for the simdgroup path. */
static long qwen4_moe_mm_nax_level(uint32_t type) {
const char *v = getenv("DS4_QWEN4_MOE_MM_NAX");
/* default 2 for every tensor tier: level 2 matches the simdgroup tiles'
* operand rounding and differs only by accumulation order, and its
* fixture NLL is within noise of them on both packs (Q4 0.20503 vs
* 0.20505, Q2 0.30418 vs 0.30346), while it carries the full prefill
* gain (+48/+51% Q4, +29/+36% Q2). The compensated level 5 keeps the
* best absolute NLL and stays one env variable away. */
if (!v || !v[0]) return (type == 12u || type == 39u || type == 16u || type == 10u) ? 2 : 0;
return strtol(v, NULL, 10);
}
static uint32_t qwen4_moe_mm_nax(uint32_t type) {
if (!ds4_gpu_mpp_available()) return 0;
if (!(type == 12u || type == 39u || type == 16u || type == 10u)) return 0;
const long n = qwen4_moe_mm_nax_level(type);
if (n <= 0) return 0;
return (n == 1 || n == 4 || n == 6) ? 32u : 64u;
}
/* float-activation tensor tiles (levels 3 and 4) */
static bool qwen4_moe_mm_nax_fx(uint32_t type) {
const long n = qwen4_moe_mm_nax_level(type);
return n == 3 || n == 4;
}
/* compensated tensor tiles (levels 5 and 6): the fast fp16 tensor kernels run
* twice per K step, on the half-rounded operand and on its half residual, so
* the activation operand enters at ~2^-22 relative instead of 2^-11 */
static bool qwen4_moe_mm_nax_comp(uint32_t type) {
const long n = qwen4_moe_mm_nax_level(type);
return n == 5 || n == 6;
}
static ds4_gpu_tensor *g_qwen4_nax_half_x, *g_qwen4_nax_half_mid, *g_qwen4_nax_half_midr;
static uint64_t g_qwen4_nax_half_x_bytes, g_qwen4_nax_half_mid_bytes, g_qwen4_nax_half_midr_bytes;
static const ds4_gpu_tensor *g_qwen4_nax_half_mid_for;
static uint64_t g_qwen4_nax_half_mid_count;
static void qwen4_nax_release_scratch(void) {
if (g_qwen4_nax_half_x) ds4_gpu_tensor_free(g_qwen4_nax_half_x);
if (g_qwen4_nax_half_mid) ds4_gpu_tensor_free(g_qwen4_nax_half_mid);
if (g_qwen4_nax_half_midr) ds4_gpu_tensor_free(g_qwen4_nax_half_midr);
g_qwen4_nax_half_x = g_qwen4_nax_half_mid = g_qwen4_nax_half_midr = NULL;
g_qwen4_nax_half_x_bytes = g_qwen4_nax_half_mid_bytes = g_qwen4_nax_half_midr_bytes = 0;
g_qwen4_nax_half_mid_for = NULL;
g_qwen4_nax_half_mid_count = 0;
}
static ds4_gpu_tensor *qwen4_nax_scratch(ds4_gpu_tensor **slot, uint64_t *slot_bytes, uint64_t bytes) {
if (!*slot || *slot_bytes < bytes) {
if (*slot) ds4_gpu_tensor_free(*slot);
*slot = ds4_gpu_tensor_alloc(bytes);
*slot_bytes = *slot ? bytes : 0;
}
return *slot;
}
static ds4_gpu_tensor *qwen4_nax_half_operand(ds4_gpu_tensor **slot, uint64_t *slot_bytes,
const ds4_gpu_tensor *src, uint64_t count) {
const uint64_t bytes = count * sizeof(uint16_t);
if (count == 0 || (count % 4u) != 0) return NULL;
ds4_gpu_tensor *dst = qwen4_nax_scratch(slot, slot_bytes, bytes);
if (!dst) return NULL;
struct { uint32_t n4; } args = { (uint32_t)(count / 4u) };
qwen4_bind b[2];
if (!qwen4_bind_tensor(&b[0], src, count * sizeof(float), "nax half operand source") ||
!qwen4_bind_tensor(&b[1], dst, bytes, "nax half operand")) return NULL;
const uint32_t n_threads = (args.n4 + 3u) / 4u; /* 16 values per thread */
if (!qwen4_dispatch(QWEN4_K_ROWS_F32_TO_F16, &args, sizeof(args), b, 2,
MTLSizeMake((n_threads + 255u) / 256u, 1, 1), MTLSizeMake(256, 1, 1), 0)) return NULL;
return dst;
}
static bool qwen4_moe_mm_tails(uint32_t type, uint32_t nt) {
/* Remainder tiles measured on M3 Ultra for the low-bit experts and on M5
* for Q4_K gate/up with MXFP4 down. */
const int override = ds4_gpu_env_bool("DS4_QWEN4_MOE_TAILS");
return nt > 1u && (override >= 0 ? override != 0 :
((type == 16u || type == 10u) && ds4_gpu_device_name_contains("M3 Ultra")) ||
((type == 12u || type == 39u) && ds4_gpu_device_is_m5_apple_silicon()));
}
int ds4_gpu_qwen4_moe_mm_mid_tensor(
ds4_gpu_tensor *mid, const ds4_gpu_tensor *x, const ds4_gpu_tensor *lists, const ds4_gpu_tensor *counts,
const void *model_map, uint64_t model_size, uint64_t gate_offset, uint64_t up_offset,
uint32_t weight_type, uint32_t n_expert, uint32_t n_tokens, uint32_t n_slots, uint32_t n_out,
uint32_t in_dim, uint32_t ff_dim, uint32_t list_cap) {
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, in_dim);
const uint64_t expert_bytes = (uint64_t)row_bytes * ff_dim;
const uint32_t tiles = qwen4_moe_mm_tiles(n_tokens, true);
const uint32_t nt = qwen4_moe_mm_nt(n_tokens, weight_type, "DS4_QWEN4_MOE_MID_NT");
const int kernel = nt == 1u ? QWEN4_K_MOE_MM_MID_NT1 :
nt == 2u ? QWEN4_K_MOE_MM_MID_NT2 :
nt == 8u ? QWEN4_K_MOE_MM_MID_NT8 : QWEN4_K_MOE_MM_MID;
qwen4_moe_mm_args args = { n_tokens, n_slots, n_out, in_dim, ff_dim, weight_type, row_bytes, list_cap,
expert_bytes, n_expert, tiles, 0, qwen4_moe_mm_expert_major() };
const bool tails = qwen4_moe_mm_tails(weight_type, nt);
if (tails) args.tail_base = nt * 8u;
qwen4_bind b[8];
if (n_tokens == 0 || n_slots == 0 || n_out < n_slots || row_bytes == 0 ||
(weight_type != 8u && weight_type != 39u && weight_type != 12u && weight_type != 10u && weight_type != 16u && weight_type != 2u) ||
(in_dim % 64) != 0 || ff_dim == 0 || n_expert == 0 || n_expert > 512 ||
!qwen4_bind_weight(&b[0], model_map, model_size, gate_offset, expert_bytes * n_expert, "moe gate experts") ||
!qwen4_bind_weight(&b[1], model_map, model_size, up_offset, expert_bytes * n_expert, "moe up experts") ||
!qwen4_bind_tensor(&b[2], lists, (uint64_t)n_expert * list_cap * sizeof(int32_t), "moe lists") ||
!qwen4_bind_tensor(&b[3], counts, (uint64_t)n_expert * sizeof(int32_t), "moe counts") ||
!qwen4_bind_tensor(&b[4], x, (uint64_t)n_tokens * in_dim * sizeof(float), "moe input") ||
!qwen4_bind_tensor(&b[5], mid, (uint64_t)n_tokens * n_out * ff_dim * sizeof(float), "moe mid")) {
return 0;
}
const uint32_t nax = qwen4_moe_mm_nax(weight_type);
const bool nax_fx = qwen4_moe_mm_nax_fx(weight_type);
const bool nax_comp = qwen4_moe_mm_nax_comp(weight_type);
if (nax) {
args.tail_base = 0;
const uint64_t mid_count = (uint64_t)n_tokens * n_out * ff_dim;
if (nax_comp) {
/* compensated tiles: x stays bound at b[4] (float) and the tiles
* write both the half copy of mid and its half residual */
ds4_gpu_tensor *mh = qwen4_nax_scratch(&g_qwen4_nax_half_mid, &g_qwen4_nax_half_mid_bytes, mid_count * sizeof(uint16_t));
ds4_gpu_tensor *mr = qwen4_nax_scratch(&g_qwen4_nax_half_midr, &g_qwen4_nax_half_midr_bytes, mid_count * sizeof(uint16_t));
if (!mh || !mr || !qwen4_bind_tensor(&b[6], mh, mid_count * sizeof(uint16_t), "moe mid half") ||
!qwen4_bind_tensor(&b[7], mr, mid_count * sizeof(uint16_t), "moe mid residual")) return 0;
g_qwen4_nax_half_mid_for = mid;
g_qwen4_nax_half_mid_count = mid_count;
} else if (nax_fx) {
/* float-activation tiles: x stays bound at b[4]; they write no
* half copy of mid, so b[6] is a never-read placeholder */
if (!qwen4_bind_tensor(&b[6], mid, mid_count * sizeof(uint16_t), "moe mid")) return 0;
} else {
ds4_gpu_tensor *xh = qwen4_nax_half_operand(&g_qwen4_nax_half_x, &g_qwen4_nax_half_x_bytes, x, (uint64_t)n_tokens * in_dim);
ds4_gpu_tensor *mh = qwen4_nax_scratch(&g_qwen4_nax_half_mid, &g_qwen4_nax_half_mid_bytes, mid_count * sizeof(uint16_t));
if (!xh || !mh || !qwen4_bind_tensor(&b[4], xh, (uint64_t)n_tokens * in_dim * sizeof(uint16_t), "moe input half") ||
!qwen4_bind_tensor(&b[6], mh, mid_count * sizeof(uint16_t), "moe mid half")) return 0;
g_qwen4_nax_half_mid_for = mid;
g_qwen4_nax_half_mid_count = mid_count;
}
const bool nax_tails = nax == 64u && qwen4_moe_mm_tails(weight_type, 8u);
if (nax_tails) args.tail_base = 64u;
const int nax_mid = nax_comp ? (nax == 64u ? QWEN4_K_MOE_MM_MID_NAXC64 : QWEN4_K_MOE_MM_MID_NAXC)
: nax_fx ? (nax == 64u ? QWEN4_K_MOE_MM_MID_NAXF64 : QWEN4_K_MOE_MM_MID_NAXF)
: (nax == 64u ? QWEN4_K_MOE_MM_MID_NAX64 : QWEN4_K_MOE_MM_MID_NAX);
const int nax_mid_tail = nax_comp ? QWEN4_K_MOE_MM_MID_NAXC : nax_fx ? QWEN4_K_MOE_MM_MID_NAXF : QWEN4_K_MOE_MM_MID_NAX;
if (!qwen4_dispatch(nax_mid, &args, sizeof(args), b, nax_comp ? 8 : 7,
qwen4_moe_mm_grid((ff_dim + 63u) / 64u, n_expert, tiles, args.expert_major),
MTLSizeMake(128, 1, 1), nax == 64u ? 16384u : nax_fx || nax_comp ? 12288u : 10240u)) return 0;
if (!nax_tails) return 1;
args.tiles_per_launch = 1;
return qwen4_dispatch(nax_mid_tail, &args, sizeof(args), b, nax_comp ? 8 : 7,
qwen4_moe_mm_grid((ff_dim + 63u) / 64u, n_expert, 1, args.expert_major),
MTLSizeMake(128, 1, 1), nax_fx || nax_comp ? 12288u : 10240u);
}
if (!qwen4_dispatch(kernel, &args, sizeof(args), b, 6,
qwen4_moe_mm_grid((ff_dim + 31u) / 32u, n_expert, tiles, args.expert_major),
MTLSizeMake(128, 1, 1), 0)) return 0;
if (tails) {
const MTLSize tail_grid = MTLSizeMake((ff_dim + 31u) / 32u, n_expert, 1);
if (!qwen4_dispatch(QWEN4_K_MOE_MM_MID_NT1, &args, sizeof(args), b, 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 2u && !qwen4_dispatch(QWEN4_K_MOE_MM_MID_NT2, &args, sizeof(args), b, 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 4u && !qwen4_dispatch(QWEN4_K_MOE_MM_MID, &args, sizeof(args), b, 6, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
}
return 1;
}
int ds4_gpu_qwen4_moe_mm_down_tensor(
ds4_gpu_tensor *part, const ds4_gpu_tensor *mid, const ds4_gpu_tensor *lists, const ds4_gpu_tensor *counts,
const void *model_map, uint64_t model_size, uint64_t down_offset,
uint32_t weight_type, uint32_t n_expert, uint32_t n_tokens, uint32_t n_slots, uint32_t n_out,
uint32_t ff_dim, uint32_t out_dim, uint32_t list_cap) {
const uint32_t weight_dim = weight_type == 10u ? (ff_dim + 255u) / 256u * 256u : ff_dim;
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, weight_dim);
const uint64_t expert_bytes = (uint64_t)row_bytes * out_dim;
const uint32_t tiles = qwen4_moe_mm_tiles(n_tokens, false);
const uint32_t nt = qwen4_moe_mm_nt(n_tokens, weight_type, "DS4_QWEN4_MOE_DOWN_NT");
const int kernel = nt == 1u ? QWEN4_K_MOE_MM_DOWN_NT1 :
nt == 2u ? QWEN4_K_MOE_MM_DOWN_NT2 :
nt == 8u ? QWEN4_K_MOE_MM_DOWN_NT8 : QWEN4_K_MOE_MM_DOWN;
qwen4_moe_mm_args args = { n_tokens, n_slots, n_out, ff_dim, out_dim, weight_type, row_bytes, list_cap,
expert_bytes, n_expert, tiles, 0, qwen4_moe_mm_expert_major() };
const bool tails = qwen4_moe_mm_tails(weight_type, nt);
if (tails) args.tail_base = nt * 8u;
qwen4_bind b[6];
if (n_tokens == 0 || n_slots == 0 || n_out < n_slots || row_bytes == 0 ||
(weight_type != 8u && weight_type != 39u && weight_type != 12u && weight_type != 10u && weight_type != 16u && weight_type != 2u) ||
(ff_dim % 64) != 0 || out_dim == 0 || n_expert == 0 || n_expert > 512 ||
!qwen4_bind_weight(&b[0], model_map, model_size, down_offset, expert_bytes * n_expert, "moe down experts") ||
!qwen4_bind_tensor(&b[1], lists, (uint64_t)n_expert * list_cap * sizeof(int32_t), "moe lists") ||
!qwen4_bind_tensor(&b[2], counts, (uint64_t)n_expert * sizeof(int32_t), "moe counts") ||
!qwen4_bind_tensor(&b[3], mid, (uint64_t)n_tokens * n_out * ff_dim * sizeof(float), "moe mid") ||
!qwen4_bind_tensor(&b[4], part, (uint64_t)n_tokens * n_out * out_dim * sizeof(float), "moe partial")) {
return 0;
}
const uint32_t nax = qwen4_moe_mm_nax(weight_type);
const bool nax_fx = qwen4_moe_mm_nax_fx(weight_type);
bool nax_comp = qwen4_moe_mm_nax_comp(weight_type);
if (nax) {
args.tail_base = 0;
const uint64_t mid_count = (uint64_t)n_tokens * n_out * ff_dim;
if (nax_comp) {
/* the compensated tiles read the half copy of mid and its
* residual; both must come from the mid tiles' matching launch,
* otherwise fall back to the uncompensated tiles */
if (g_qwen4_nax_half_mid_for == mid && g_qwen4_nax_half_mid_count == mid_count &&
g_qwen4_nax_half_mid && g_qwen4_nax_half_midr) {
if (!qwen4_bind_tensor(&b[3], g_qwen4_nax_half_mid, mid_count * sizeof(uint16_t), "moe mid half") ||
!qwen4_bind_tensor(&b[5], g_qwen4_nax_half_midr, mid_count * sizeof(uint16_t), "moe mid residual")) return 0;
g_qwen4_nax_half_mid_for = NULL;
g_qwen4_nax_half_mid_count = 0;
} else {
nax_comp = false;
}
}
if (!nax_comp && !nax_fx) {
/* the half tiles read the mid tiles' half copy of mid */
ds4_gpu_tensor *mh = g_qwen4_nax_half_mid_for == mid && g_qwen4_nax_half_mid_count == mid_count && g_qwen4_nax_half_mid
? g_qwen4_nax_half_mid
: qwen4_nax_half_operand(&g_qwen4_nax_half_mid, &g_qwen4_nax_half_mid_bytes, mid, mid_count);
g_qwen4_nax_half_mid_for = NULL;
g_qwen4_nax_half_mid_count = 0;
if (!mh || !qwen4_bind_tensor(&b[3], mh, mid_count * sizeof(uint16_t), "moe mid half")) return 0;
}
const bool nax_tails = nax == 64u && qwen4_moe_mm_tails(weight_type, 8u);
if (nax_tails) args.tail_base = 64u;
const int nax_down = nax_comp ? (nax == 64u ? QWEN4_K_MOE_MM_DOWN_NAXC64 : QWEN4_K_MOE_MM_DOWN_NAXC)
: nax_fx ? (nax == 64u ? QWEN4_K_MOE_MM_DOWN_NAXF64 : QWEN4_K_MOE_MM_DOWN_NAXF)
: (nax == 64u ? QWEN4_K_MOE_MM_DOWN_NAX64 : QWEN4_K_MOE_MM_DOWN_NAX);
const int nax_down_tail = nax_comp ? QWEN4_K_MOE_MM_DOWN_NAXC : nax_fx ? QWEN4_K_MOE_MM_DOWN_NAXF : QWEN4_K_MOE_MM_DOWN_NAX;
if (!qwen4_dispatch(nax_down, &args, sizeof(args), b, nax_comp ? 6 : 5,
qwen4_moe_mm_grid((out_dim + 63u) / 64u, n_expert, tiles, args.expert_major),
MTLSizeMake(128, 1, 1), nax == 64u ? 16384u : 8192u)) return 0;
if (!nax_tails) return 1;
args.tiles_per_launch = 1;
return qwen4_dispatch(nax_down_tail, &args, sizeof(args), b, nax_comp ? 6 : 5,
qwen4_moe_mm_grid((out_dim + 63u) / 64u, n_expert, 1, args.expert_major),
MTLSizeMake(128, 1, 1), 8192u);
}
if (!qwen4_dispatch(kernel, &args, sizeof(args), b, 5,
qwen4_moe_mm_grid((out_dim + 31u) / 32u, n_expert, tiles, args.expert_major),
MTLSizeMake(128, 1, 1), 0)) return 0;
if (tails) {
const MTLSize tail_grid = MTLSizeMake((out_dim + 31u) / 32u, n_expert, 1);
if (!qwen4_dispatch(QWEN4_K_MOE_MM_DOWN_NT1, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 2u && !qwen4_dispatch(QWEN4_K_MOE_MM_DOWN_NT2, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
if (nt > 4u && !qwen4_dispatch(QWEN4_K_MOE_MM_DOWN, &args, sizeof(args), b, 5, tail_grid, MTLSizeMake(128, 1, 1), 0)) return 0;
}
return 1;
}
int ds4_gpu_qwen4_matmul_q8_0_weights_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *w,
uint32_t in_dim, uint32_t out_dim, const ds4_gpu_tensor *x) {
if (!g_initialized && !ds4_gpu_init()) return 0;
if ((in_dim & 31u) != 0 || !out_dim) return 0;
const uint64_t row_bytes = (uint64_t)(in_dim / 32u) * 34u;
@autoreleasepool {
id<MTLBuffer> xbuf = ds4_gpu_tensor_buffer(x);
id<MTLBuffer> outbuf = ds4_gpu_tensor_buffer(out);
id<MTLBuffer> wbuf = ds4_gpu_tensor_buffer(w);
if (!xbuf || !outbuf || !wbuf ||
ds4_gpu_tensor_bytes(x) < (uint64_t)in_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(out) < (uint64_t)out_dim * sizeof(float) ||
ds4_gpu_tensor_bytes(w) < (uint64_t)out_dim * row_bytes) {
fprintf(stderr, "ds4: Metal Q8_0 tensor-weight matvec received undersized buffers\n");
return 0;
}
int owned = 0;
id<MTLCommandBuffer> cb = ds4_gpu_command_buffer(&owned);
if (!cb) return 0;
ds4_gpu_q8_0_matvec_args mv_args = ds4_gpu_make_q8_0_mv_args(in_dim, out_dim);
ds4_gpu_mv_dispatch mv_dispatch = ds4_gpu_make_q8_0_mv_dispatch();
if (out_dim > 65536u) mv_dispatch.nsg = 8;
mv_args.nr0 = mv_dispatch.nr0;
id<MTLComputePipelineState> pipeline =
ds4_gpu_get_mul_mv_pipeline(mv_dispatch.function_name, mv_dispatch.nsg);
if (!pipeline) return 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState:pipeline];
[enc setBytes:&mv_args length:sizeof(mv_args) atIndex:0];
[enc setBuffer:wbuf offset:ds4_gpu_tensor_offset(w) atIndex:1];
[enc setBuffer:xbuf offset:ds4_gpu_tensor_offset(x) atIndex:2];
[enc setBuffer:outbuf offset:ds4_gpu_tensor_offset(out) atIndex:3];
[enc setThreadgroupMemoryLength:mv_dispatch.smem atIndex:0];
[enc dispatchThreadgroups:MTLSizeMake(((NSUInteger)out_dim + (NSUInteger)mv_dispatch.nr0 - 1u) / (NSUInteger)mv_dispatch.nr0, 1, 1)
threadsPerThreadgroup:MTLSizeMake(32, (NSUInteger)mv_dispatch.nsg, 1)];
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "Q8_0 tensor-weight matvec") ? 1 : 0;
}
}
int ds4_gpu_qwen4_argmax_tensor(ds4_gpu_tensor *out_idx, ds4_gpu_tensor *scratch,
const ds4_gpu_tensor *logits, uint32_t n_vocab) {
if (!n_vocab || n_vocab > INT32_MAX) return 0;
const uint32_t chunks = (n_vocab + 4095u) / 4096u;
struct { uint32_t n, finish; } args = {n_vocab, 0};
qwen4_bind b[3];
if (!qwen4_bind_tensor(&b[0], logits, (uint64_t)n_vocab * sizeof(float), "argmax logits") ||
!qwen4_bind_tensor(&b[1], scratch, (uint64_t)chunks * 8u, "argmax partials") ||
!qwen4_bind_tensor(&b[2], out_idx, sizeof(int32_t), "argmax index")) return 0;
if (!qwen4_dispatch(QWEN4_K_ARGMAX, &args, sizeof(args), b, 3,
MTLSizeMake(chunks, 1, 1), MTLSizeMake(256, 1, 1), 0)) return 0;
args.n = chunks;
args.finish = 1;
return qwen4_dispatch(QWEN4_K_ARGMAX, &args, sizeof(args), b, 3,
MTLSizeMake(1, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_mtp_stage_tensor(
ds4_gpu_tensor *cat, const ds4_gpu_tensor *e, const ds4_gpu_tensor *R,
const void *model_map, uint64_t model_size, uint64_t g_e_offset, uint64_t g_h_offset,
uint32_t n_embd, uint32_t n_hc, float eps) {
struct { uint32_t n_embd, n_hc, pad0; float eps; } args = { n_embd, n_hc, 0u, eps };
qwen4_bind b[5];
if (n_embd == 0 || n_hc == 0 ||
!qwen4_bind_tensor(&b[0], e, (uint64_t)n_embd * sizeof(float), "mtp embedding") ||
!qwen4_bind_tensor(&b[1], R, (uint64_t)n_embd * n_hc * sizeof(float), "mtp streams") ||
!qwen4_bind_weight(&b[2], model_map, model_size, g_e_offset, (uint64_t)n_embd * sizeof(float), "mtp enorm") ||
!qwen4_bind_weight(&b[3], model_map, model_size, g_h_offset, (uint64_t)n_embd * n_hc * sizeof(float),
"mtp hnorm") ||
!qwen4_bind_tensor(&b[4], cat, (uint64_t)(n_hc + 1u) * 2u * n_embd * sizeof(float), "mtp concat")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_MTP_STAGE, &args, sizeof(args), b, 5,
MTLSizeMake(n_hc + 1u, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_mtp_combine_tensor(
ds4_gpu_tensor *R_out, const ds4_gpu_tensor *proj, uint32_t n_embd, uint32_t n_hc) {
struct { uint32_t n_embd, n_hc; } args = { n_embd, n_hc };
qwen4_bind b[2];
if (n_embd == 0 || n_hc == 0 ||
!qwen4_bind_tensor(&b[0], proj, (uint64_t)(n_hc + 1u) * n_embd * sizeof(float), "mtp projection") ||
!qwen4_bind_tensor(&b[1], R_out, (uint64_t)n_embd * n_hc * sizeof(float), "mtp residual")) {
return 0;
}
const uint32_t n = n_embd * n_hc;
return qwen4_dispatch(QWEN4_K_MTP_COMBINE, &args, sizeof(args), b, 2,
MTLSizeMake((n + 255u) / 256u, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_gdn_front_tensor(
ds4_gpu_tensor *qkv, ds4_gpu_tensor *state, const ds4_gpu_tensor *mixed,
ds4_gpu_tensor *ga, ds4_gpu_tensor *gb,
const void *model_map, uint64_t model_size, uint64_t conv_offset,
uint64_t alpha_offset, uint64_t beta_offset, uint64_t ssm_a_offset, uint64_t dt_bias_offset,
uint32_t weight_type, uint32_t n_tokens, uint32_t n_k_head, uint32_t n_v_head, uint32_t head_dim,
uint32_t conv_kernel, uint32_t in_dim, ds4_gpu_tensor *snap_state, uint32_t snap_tok,
ds4_gpu_tensor *snap2_state, uint32_t snap2_tok) {
const uint32_t row_bytes = qwen4_expert_row_bytes(weight_type, in_dim);
struct { uint32_t n_tokens, n_k_head, n_v_head, head_dim, conv_kernel, weight_type, in_dim, row_bytes,
snap_tok, snap2_tok, pad1, pad2; } args =
{ n_tokens, n_k_head, n_v_head, head_dim, conv_kernel, weight_type, in_dim, row_bytes,
snap_state ? snap_tok : UINT32_MAX, snap2_state ? snap2_tok : UINT32_MAX, 0, 0 };
const uint64_t conv_dim = 2ull * n_k_head * head_dim + (uint64_t)n_v_head * head_dim;
const uint64_t state_bytes = (uint64_t)(conv_kernel - 1) * conv_dim * sizeof(float);
const uint64_t proj_bytes = (uint64_t)n_v_head * row_bytes;
const uint64_t vhead_bytes = (uint64_t)n_v_head * sizeof(float);
qwen4_bind b[12];
if (n_tokens == 0 || n_k_head == 0 || head_dim % 32 != 0 || n_v_head % n_k_head != 0 ||
conv_kernel < 2 || conv_kernel > 4 || row_bytes == 0 ||
!qwen4_bind_tensor(&b[0], qkv, (uint64_t)n_tokens * conv_dim * sizeof(float), "gdn front qkv") ||
!qwen4_bind_tensor(&b[1], state, state_bytes, "gdn front conv state") ||
!qwen4_bind_weight(&b[2], model_map, model_size, conv_offset, conv_dim * conv_kernel * sizeof(float),
"gdn front conv weight") ||
!qwen4_bind_tensor(&b[3], mixed, (uint64_t)n_tokens * in_dim * sizeof(float), "gdn front input") ||
!qwen4_bind_weight(&b[4], model_map, model_size, alpha_offset, proj_bytes, "gdn front alpha") ||
!qwen4_bind_weight(&b[5], model_map, model_size, beta_offset, proj_bytes, "gdn front beta") ||
!qwen4_bind_weight(&b[6], model_map, model_size, ssm_a_offset, vhead_bytes, "gdn front ssm_a") ||
!qwen4_bind_weight(&b[7], model_map, model_size, dt_bias_offset, vhead_bytes, "gdn front dt_bias") ||
!qwen4_bind_tensor(&b[8], ga, (uint64_t)n_tokens * n_v_head * sizeof(float), "gdn front decay") ||
!qwen4_bind_tensor(&b[9], gb, (uint64_t)n_tokens * n_v_head * sizeof(float), "gdn front beta out")) {
return 0;
}
if (snap_state) {
if (!qwen4_bind_tensor(&b[10], snap_state, state_bytes, "gdn front snapshot")) return 0;
} else {
b[10] = b[1];
}
if (snap2_state) {
if (!qwen4_bind_tensor(&b[11], snap2_state, state_bytes, "gdn front snapshot 2")) return 0;
} else {
b[11] = b[1];
}
/* One threadgroup per key head: simdgroups 0/1 normalize, the rest take
* the alpha/beta rows, and every channel's conv stays on one thread, so
* the thread count only spreads the conv wider. M5 measured 1024. */
const uint64_t nth = ds4_gpu_env_u64("DS4_QWEN4_GDN_FRONT_THREADS",
ds4_gpu_device_is_m5_apple_silicon() ? 1024u : 256u, 96u, 1024u);
return qwen4_dispatch(QWEN4_K_GDN_FRONT, &args, sizeof(args), b, 12,
MTLSizeMake(n_k_head, 1, 1), MTLSizeMake((NSUInteger)(nth / 32u * 32u), 1, 1), 0);
}
int ds4_gpu_qwen4_multi_gemv_tensor(
const ds4_gpu_tensor *x, uint32_t n_tokens, uint32_t in_dim, uint32_t n_out,
ds4_gpu_tensor *const *outs, const void *model_map, uint64_t model_size,
const uint64_t *offsets, const uint32_t *types, const uint32_t *out_rows) {
struct { uint32_t n_tokens, in_dim, n_out, pad0, rows[4], types[4], row_bytes[4]; } args = {0};
qwen4_bind b[9];
if (!x || n_tokens == 0 || in_dim == 0 || (in_dim % 32) != 0 || n_out == 0 || n_out > 4 ||
!qwen4_bind_tensor(&b[0], x, (uint64_t)n_tokens * in_dim * sizeof(float), "multi gemv input")) {
return 0;
}
args.n_tokens = n_tokens;
args.in_dim = in_dim;
args.n_out = n_out;
uint32_t total = 0;
for (uint32_t i = 0; i < 4; i++) {
if (i < n_out) {
const uint32_t rb = qwen4_expert_row_bytes(types[i], in_dim);
if (rb == 0 || out_rows[i] == 0 ||
!qwen4_bind_weight(&b[1 + i], model_map, model_size, offsets[i], (uint64_t)rb * out_rows[i],
"multi gemv weight") ||
!qwen4_bind_tensor(&b[5 + i], outs[i], (uint64_t)n_tokens * out_rows[i] * sizeof(float),
"multi gemv output")) {
return 0;
}
args.rows[i] = out_rows[i];
args.types[i] = types[i];
args.row_bytes[i] = rb;
total += out_rows[i];
} else {
b[1 + i] = b[1];
b[5 + i] = b[5];
}
}
return qwen4_dispatch(QWEN4_K_MULTI_GEMV, &args, sizeof(args), b, 9,
MTLSizeMake((total + 7) / 8, n_tokens, 1), MTLSizeMake(128, 1, 1), 0);
}
static ds4_gpu_tensor *g_qwen4_dense_mm_partials;
static uint64_t g_qwen4_dense_mm_partials_bytes;
static void qwen4_batch_release_scratch(void) {
ds4_gpu_tensor_free(g_qwen4_dense_mm_partials);
ds4_gpu_tensor_free(g_qwen4_gdn_slots);
g_qwen4_dense_mm_partials = g_qwen4_gdn_slots = NULL;
g_qwen4_dense_mm_partials_bytes = 0;
for (unsigned i = 0; i < QWEN4_K_COUNT; i++) g_qwen4_pipelines[i] = nil;
}
/* Scratch for the k-split planes, allocated once and generously: the
* command buffer holds unretained references, so a scratch replaced while
* earlier dispatches of the same buffer still point at it is read after its
* release. Growing it per request was what made wide splits look wrong. */
static bool qwen4_dense_mm_partials_ensure(uint64_t bytes) {
if (g_qwen4_dense_mm_partials && g_qwen4_dense_mm_partials_bytes >= bytes) return true;
if (g_qwen4_dense_mm_partials) {
fprintf(stderr, "ds4: Qwen3.8 k-split scratch of %.1f MiB cannot grow to %.1f MiB mid-batch\n",
ds4_gpu_mib(g_qwen4_dense_mm_partials_bytes), ds4_gpu_mib(bytes));
return false;
}
const uint64_t grow = bytes > (64u << 20) ? bytes : (64u << 20);
g_qwen4_dense_mm_partials = ds4_gpu_tensor_alloc(grow);
if (!g_qwen4_dense_mm_partials) return false;
g_qwen4_dense_mm_partials_bytes = grow;
return true;
}
int ds4_gpu_qwen4_dense_mm_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t weight_type,
uint32_t n_tokens, uint32_t in_dim, uint32_t out_rows) {
const uint32_t row_bytes = weight_type == 0u ? in_dim * 4u : weight_type == 1u ? in_dim * 2u :
qwen4_expert_row_bytes(weight_type, in_dim);
struct { uint32_t n_tokens, in_dim, out_rows, weight_type, row_bytes, n_split, pad1, pad2; } args =
{ n_tokens, in_dim, out_rows, weight_type, row_bytes, 1, 0, 0 };
qwen4_bind b[3];
if (n_tokens == 0 || row_bytes == 0 || (weight_type != 0u && weight_type != 1u && weight_type != 8u) ||
(weight_type == 8u && (in_dim % 32) != 0) || (in_dim % 8) != 0 || out_rows == 0) {
fprintf(stderr, "ds4: Qwen3.8 dense mm rejected type %u in %u out %u tokens %u\n",
weight_type, in_dim, out_rows, n_tokens);
return 0;
}
if (!qwen4_bind_weight(&b[0], model_map, model_size, weight_offset, (uint64_t)row_bytes * out_rows,
"dense mm weight") ||
!qwen4_bind_tensor(&b[1], x, (uint64_t)n_tokens * in_dim * sizeof(float), "dense mm input") ||
!qwen4_bind_tensor(&b[2], out, (uint64_t)n_tokens * out_rows * sizeof(float), "dense mm output")) {
return 0;
}
/* A projection whose output is narrow gives this grid only a few
* threadgroups when the batch is a decode step; split k until the
* machine has work, then sum the planes. A prefill chunk's token tiles
* already fill it, and its planes would not fit the scratch anyway.
* The split is chosen so the partials stay small and every threadgroup
* still walks enough k to amortize its staging. */
const uint32_t tiles = (out_rows + 31u) / 32u;
const uint32_t threadgroups = tiles * ((n_tokens + 31u) / 32u);
uint32_t n_split = 1u;
if (!ds4_gpu_env_u64("DS4_QWEN4_NO_DENSE_MM_KSPLIT", 0u, 0u, 1u)) {
const uint32_t nk = (in_dim + 31u) / 32u;
const uint32_t target = 128u;
const uint32_t want = threadgroups >= target ? 1u : (target + threadgroups - 1u) / threadgroups;
n_split = want > nk / 4u ? (nk / 4u ? nk / 4u : 1u) : want;
if (n_split > 64u) n_split = 64u;
}
args.n_split = n_split;
if (n_split <= 1u) {
return qwen4_dispatch(QWEN4_K_DENSE_MM, &args, sizeof(args), b, 3,
MTLSizeMake(tiles, (n_tokens + 31) / 32, 1), MTLSizeMake(128, 1, 1), 0);
}
const uint64_t plane = (uint64_t)n_tokens * out_rows * sizeof(float);
if (!qwen4_dense_mm_partials_ensure(plane * n_split)) return 0;
qwen4_bind bp[3];
bp[0] = b[0];
bp[1] = b[1];
if (!qwen4_bind_tensor(&bp[2], g_qwen4_dense_mm_partials, plane * n_split, "dense mm partials")) {
return 0;
}
if (!qwen4_dispatch(QWEN4_K_DENSE_MM, &args, sizeof(args), bp, 3,
MTLSizeMake(tiles, (n_tokens + 31) / 32, n_split), MTLSizeMake(128, 1, 1), 0)) {
return 0;
}
qwen4_bind br[2];
br[0] = bp[2];
br[1] = b[2];
const uint64_t n = (uint64_t)n_tokens * out_rows;
return qwen4_dispatch(QWEN4_K_DENSE_MM_REDUCE, &args, sizeof(args), br, 2,
MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1), MTLSizeMake(256, 1, 1), 0);
}
/* Decode batch of 8 or 16 rows against a Q8 matrix (kernel_qwen4_batch_mm_q8):
* the k-split gives narrow projections enough simdgroups, the partial planes
* are summed by kernel_qwen4_dense_mm_reduce. */
int ds4_gpu_qwen4_batch_mm_q8_tensor(
ds4_gpu_tensor *out, const ds4_gpu_tensor *x,
const void *model_map, uint64_t model_size, uint64_t weight_offset,
uint32_t n_tokens, uint32_t in_dim, uint32_t out_rows) {
if ((n_tokens != 8u && n_tokens != 16u) || (in_dim % 32u) != 0 || (out_rows % 32u) != 0 || in_dim == 0) return 0;
const uint32_t row_bytes = (in_dim / 32u) * 34u;
const uint32_t nk = in_dim / 32u;
const uint32_t simdgroups = out_rows / 32u;
/* enough simdgroups to cover the machine several times over, no split
* shorter than four blocks */
const uint32_t target = 640u;
uint32_t n_split = simdgroups >= target ? 1u : (target + simdgroups - 1u) / simdgroups;
if (n_split > nk / 4u) n_split = nk / 4u ? nk / 4u : 1u;
if (n_split > 16u) n_split = 16u;
struct { uint32_t n_tokens, in_dim, out_rows, weight_type, row_bytes, n_split, pad1, pad2; } args =
{ n_tokens, in_dim, out_rows, 8u, row_bytes, n_split, 0, 0 };
qwen4_bind b[3];
if (!qwen4_bind_weight(&b[0], model_map, model_size, weight_offset, (uint64_t)row_bytes * out_rows, "batch mm weight") ||
!qwen4_bind_tensor(&b[1], x, (uint64_t)n_tokens * in_dim * sizeof(float), "batch mm input") ||
!qwen4_bind_tensor(&b[2], out, (uint64_t)n_tokens * out_rows * sizeof(float), "batch mm output")) {
return 0;
}
const MTLSize grid = MTLSizeMake((out_rows + 127u) / 128u, 1, n_split);
const int kernel = n_tokens == 16u ? QWEN4_K_BATCH_MM_Q8_T2 : QWEN4_K_BATCH_MM_Q8_T1;
if (n_split <= 1u) {
return qwen4_dispatch(kernel, &args, sizeof(args), b, 3, grid, MTLSizeMake(128, 1, 1), 0);
}
const uint64_t plane = (uint64_t)n_tokens * out_rows * sizeof(float);
if (!qwen4_dense_mm_partials_ensure(plane * n_split)) return 0;
qwen4_bind bp[3] = { b[0], b[1] };
if (!qwen4_bind_tensor(&bp[2], g_qwen4_dense_mm_partials, plane * n_split, "batch mm partials")) return 0;
if (!qwen4_dispatch(kernel, &args, sizeof(args), bp, 3, grid, MTLSizeMake(128, 1, 1), 0)) return 0;
qwen4_bind br[2] = { bp[2], b[2] };
const uint64_t n = (uint64_t)n_tokens * out_rows;
return qwen4_dispatch(QWEN4_K_DENSE_MM_REDUCE, &args, sizeof(args), br, 2,
MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1), MTLSizeMake(256, 1, 1), 0);
}
/* Vision tower (metal/qwen4_vision.metal). */
typedef struct { uint32_t rows, width, grid_w, n_head, head_dim, mode; float eps, scale; } qwen4_vis_args;
static int qwen4_vis_rows(int kernel, const qwen4_vis_args *a, qwen4_bind *b, int n_bind, uint32_t rows, uint32_t y) {
const int ok = qwen4_dispatch(kernel, a, sizeof(*a), b, n_bind, MTLSizeMake(rows, y, 1),
MTLSizeMake(kernel == QWEN4_K_VIS_ATTENTION ? 32 : 256, 1, 1), 0);
if (!ok) fprintf(stderr, "ds4: Qwen3.8 vision kernel %s failed\n", qwen4_kernel_names[kernel]);
return ok;
}
static int qwen4_vis_elems(int kernel, const qwen4_vis_args *a, qwen4_bind *b, int n_bind) {
const uint64_t n = (uint64_t)a->rows * a->width;
const int ok = qwen4_dispatch(kernel, a, sizeof(*a), b, n_bind, MTLSizeMake((NSUInteger)((n + 255) / 256), 1, 1),
MTLSizeMake(256, 1, 1), 0);
if (!ok) fprintf(stderr, "ds4: Qwen3.8 vision kernel %s failed\n", qwen4_kernel_names[kernel]);
return ok;
}
static bool qwen4_vis_w(qwen4_bind *b, const void *map, uint64_t size, uint64_t off, uint64_t n, const char *what) {
return qwen4_bind_weight(b, map, size, off, n * sizeof(float), what);
}
static bool qwen4_vis_t(qwen4_bind *b, const ds4_gpu_tensor *t, uint64_t n, const char *what) {
return qwen4_bind_tensor(b, t, n * sizeof(float), what);
}
int ds4_gpu_qwen4_vision_encode(float *out, const float *patches, const float *pos_embed, uint32_t n_patches,
uint32_t grid_w, const void *map, uint64_t size, const ds4_qwen4_vision_weights *w) {
if (!out || !patches || !pos_embed || !map || !w) return 0;
const uint32_t E = w->n_embd, FF = w->n_ff, H = w->n_head, D = H ? E / H : 0, P = w->n_patch, M = w->n_merge;
const uint32_t in_patch = 3u * P * P, mm = M * M, ME = E * mm, O = w->n_out;
if (n_patches == 0 || mm == 0 || (n_patches % mm) != 0 || grid_w == 0 || (grid_w % M) != 0 ||
(n_patches % grid_w) != 0 || ((n_patches / grid_w) % M) != 0 ||
(D != 64u && D != 72u) || H * D != E || FF == 0 || O == 0 || (in_patch % 32) != 0 || (E % 32) != 0 ||
n_patches > 65535u || ds4_gpu_commands_active()) {
fprintf(stderr, "ds4: invalid Qwen3.8 vision encoder input\n");
return 0;
}
if (!g_initialized && !ds4_gpu_init()) return 0;
const uint32_t merged = n_patches / mm;
const uint64_t N = n_patches;
ds4_gpu_tensor *patch = NULL, *a0 = NULL, *a1 = NULL, *pos = NULL, *x = NULL, *tmp = NULL, *qkv = NULL;
ds4_gpu_tensor *q = NULL, *k = NULL, *v = NULL, *attn = NULL, *ffn = NULL, *m0 = NULL, *res = NULL;
qwen4_vis_args a = { n_patches, E, grid_w, H, D, 0, w->eps, 1.0f / sqrtf((float)D) };
qwen4_bind b[5];
int ok = 0;
#define QWEN4_VISION_ALLOC(name_, count_) do { \
name_ = ds4_gpu_tensor_alloc((uint64_t)(count_) * sizeof(float)); \
if (!(name_)) goto cleanup; \
} while (0)
QWEN4_VISION_ALLOC(patch, N * in_patch);
QWEN4_VISION_ALLOC(a0, N * E);
QWEN4_VISION_ALLOC(a1, N * E);
QWEN4_VISION_ALLOC(pos, N * E);
QWEN4_VISION_ALLOC(x, N * E);
QWEN4_VISION_ALLOC(tmp, N * E);
QWEN4_VISION_ALLOC(qkv, N * 3u * E);
QWEN4_VISION_ALLOC(q, N * E);
QWEN4_VISION_ALLOC(k, N * E);
QWEN4_VISION_ALLOC(v, N * E);
QWEN4_VISION_ALLOC(attn, N * E);
QWEN4_VISION_ALLOC(ffn, N * FF);
QWEN4_VISION_ALLOC(m0, (uint64_t)merged * ME);
QWEN4_VISION_ALLOC(res, (uint64_t)merged * O);
#undef QWEN4_VISION_ALLOC
if (!ds4_gpu_tensor_write(patch, 0, patches, N * in_patch * sizeof(float)) ||
!ds4_gpu_tensor_write(pos, 0, pos_embed, N * E * sizeof(float))) goto cleanup;
if (!ds4_gpu_begin_commands()) goto cleanup;
ok = ds4_gpu_qwen4_dense_mm_tensor(a0, patch, map, size, w->patch_w0, w->patch_type, n_patches, in_patch, E) &&
ds4_gpu_qwen4_dense_mm_tensor(a1, patch, map, size, w->patch_w1, w->patch_type, n_patches, in_patch, E) &&
qwen4_vis_t(&b[0], a0, N * E, "vision patch a0") &&
qwen4_vis_t(&b[1], a1, N * E, "vision patch a1") &&
qwen4_vis_w(&b[2], map, size, w->patch_b, E, "vision patch bias") &&
qwen4_vis_t(&b[3], pos, N * E, "vision pos") &&
qwen4_vis_t(&b[4], x, N * E, "vision x") &&
qwen4_vis_elems(QWEN4_K_VIS_PATCH_FINISH, &a, b, 5);
for (uint32_t il = 0; ok && il < DS4_QWEN4_VISION_LAYERS; il++) {
const ds4_qwen4_vision_layer_weights *l = &w->layer[il];
a.rows = n_patches; a.width = E;
ok = qwen4_vis_t(&b[0], x, N * E, "vision x") &&
qwen4_vis_w(&b[1], map, size, l->ln1_w, E, "vision ln1") &&
qwen4_vis_w(&b[2], map, size, l->ln1_b, E, "vision ln1 bias") &&
qwen4_vis_t(&b[3], tmp, N * E, "vision tmp") &&
qwen4_vis_rows(QWEN4_K_VIS_LAYERNORM, &a, b, 4, n_patches, 1) &&
ds4_gpu_qwen4_dense_mm_tensor(qkv, tmp, map, size, l->qkv_w, l->qkv_type, n_patches, E, 3u * E) &&
qwen4_vis_t(&b[0], qkv, N * 3u * E, "vision qkv") &&
qwen4_vis_w(&b[1], map, size, l->qkv_b, 3u * E, "vision qkv bias") &&
qwen4_vis_t(&b[2], q, N * E, "vision q") &&
qwen4_vis_t(&b[3], k, N * E, "vision k") &&
qwen4_vis_t(&b[4], v, N * E, "vision v") &&
qwen4_vis_rows(QWEN4_K_VIS_QKV_ROPE, &a, b, 5, n_patches, 1) &&
qwen4_vis_t(&b[0], q, N * E, "vision q") &&
qwen4_vis_t(&b[1], k, N * E, "vision k") &&
qwen4_vis_t(&b[2], v, N * E, "vision v") &&
qwen4_vis_t(&b[3], attn, N * E, "vision attn") &&
qwen4_vis_rows(QWEN4_K_VIS_ATTENTION, &a, b, 4, n_patches, H) &&
ds4_gpu_qwen4_dense_mm_tensor(tmp, attn, map, size, l->out_w, l->out_type, n_patches, E, E) &&
qwen4_vis_t(&b[0], x, N * E, "vision x") &&
qwen4_vis_t(&b[1], tmp, N * E, "vision tmp") &&
qwen4_vis_w(&b[2], map, size, l->out_b, E, "vision out bias") &&
qwen4_vis_elems(QWEN4_K_VIS_BIAS_RESIDUAL, &a, b, 3) &&
qwen4_vis_t(&b[0], x, N * E, "vision x") &&
qwen4_vis_w(&b[1], map, size, l->ln2_w, E, "vision ln2") &&
qwen4_vis_w(&b[2], map, size, l->ln2_b, E, "vision ln2 bias") &&
qwen4_vis_t(&b[3], tmp, N * E, "vision tmp") &&
qwen4_vis_rows(QWEN4_K_VIS_LAYERNORM, &a, b, 4, n_patches, 1) &&
ds4_gpu_qwen4_dense_mm_tensor(ffn, tmp, map, size, l->up_w, l->up_type, n_patches, E, FF);
if (!ok) break;
a.width = FF; a.mode = 0;
ok = qwen4_vis_t(&b[0], ffn, N * FF, "vision ffn") &&
qwen4_vis_w(&b[1], map, size, l->up_b, FF, "vision up bias") &&
qwen4_vis_elems(QWEN4_K_VIS_BIAS_ACT, &a, b, 2) &&
ds4_gpu_qwen4_dense_mm_tensor(tmp, ffn, map, size, l->down_w, l->down_type, n_patches, FF, E);
a.width = E;
ok = ok && qwen4_vis_t(&b[0], x, N * E, "vision x") &&
qwen4_vis_t(&b[1], tmp, N * E, "vision tmp") &&
qwen4_vis_w(&b[2], map, size, l->down_b, E, "vision down bias") &&
qwen4_vis_elems(QWEN4_K_VIS_BIAS_RESIDUAL, &a, b, 3);
}
if (ok) {
a.rows = n_patches; a.width = E;
ok = qwen4_vis_t(&b[0], x, N * E, "vision x") &&
qwen4_vis_w(&b[1], map, size, w->post_ln_w, E, "vision post ln") &&
qwen4_vis_w(&b[2], map, size, w->post_ln_b, E, "vision post ln bias") &&
qwen4_vis_t(&b[3], tmp, N * E, "vision tmp") &&
qwen4_vis_rows(QWEN4_K_VIS_LAYERNORM, &a, b, 4, n_patches, 1) &&
ds4_gpu_qwen4_dense_mm_tensor(m0, tmp, map, size, w->mm0_w, w->mm0_type, merged, ME, ME);
a.rows = merged; a.width = ME; a.mode = 1;
ok = ok && qwen4_vis_t(&b[0], m0, (uint64_t)merged * ME, "vision merger") &&
qwen4_vis_w(&b[1], map, size, w->mm0_b, ME, "vision merger bias") &&
qwen4_vis_elems(QWEN4_K_VIS_BIAS_ACT, &a, b, 2) &&
ds4_gpu_qwen4_dense_mm_tensor(res, m0, map, size, w->mm2_w, w->mm2_type, merged, ME, O);
a.width = O; a.mode = 2;
ok = ok && qwen4_vis_t(&b[0], res, (uint64_t)merged * O, "vision out") &&
qwen4_vis_w(&b[1], map, size, w->mm2_b, O, "vision out bias") &&
qwen4_vis_elems(QWEN4_K_VIS_BIAS_ACT, &a, b, 2);
}
if (!ds4_gpu_end_commands()) ok = 0;
if (ok) ok = ds4_gpu_tensor_read(res, 0, out, (uint64_t)merged * O * sizeof(float));
cleanup:
ds4_gpu_tensor_free(res);
ds4_gpu_tensor_free(m0);
ds4_gpu_tensor_free(ffn);
ds4_gpu_tensor_free(attn);
ds4_gpu_tensor_free(v);
ds4_gpu_tensor_free(k);
ds4_gpu_tensor_free(q);
ds4_gpu_tensor_free(qkv);
ds4_gpu_tensor_free(tmp);
ds4_gpu_tensor_free(x);
ds4_gpu_tensor_free(pos);
ds4_gpu_tensor_free(a1);
ds4_gpu_tensor_free(a0);
ds4_gpu_tensor_free(patch);
return ok;
}
int ds4_gpu_qwen4_hc_lo_act_tensor(ds4_gpu_tensor *lo_act, const ds4_gpu_tensor *lo,
uint32_t n_tokens, uint32_t n_hc, uint32_t n_rank) {
struct { uint32_t n_tokens, n_embd, n_hc, n_rank; } args = { n_tokens, 0, n_hc, n_rank };
qwen4_bind b[2];
const uint64_t n = (uint64_t)n_tokens * n_rank;
if (n == 0 || n_hc == 0 ||
!qwen4_bind_tensor(&b[0], lo, n * sizeof(float), "hc lo") ||
!qwen4_bind_tensor(&b[1], lo_act, n * sizeof(float), "hc lo act")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_HC_LO_ACT, &args, sizeof(args), b, 2,
MTLSizeMake((n + 255) / 256, 1, 1), MTLSizeMake(256, 1, 1), 0);
}
int ds4_gpu_qwen4_hc_mix_rows_tensor(ds4_gpu_tensor *mixed, const ds4_gpu_tensor *u, const ds4_gpu_tensor *xn,
uint32_t n_tokens, uint32_t n_embd, uint32_t n_hc) {
struct { uint32_t n_tokens, n_embd, n_hc, n_rank; } args = { n_tokens, n_embd, n_hc, 0 };
qwen4_bind b[3];
const uint64_t dim = (uint64_t)n_embd * n_hc;
if (n_tokens == 0 || n_embd == 0 || n_hc == 0 ||
!qwen4_bind_tensor(&b[0], u, (uint64_t)n_tokens * dim * sizeof(float), "hc u") ||
!qwen4_bind_tensor(&b[1], xn, (uint64_t)n_tokens * dim * sizeof(float), "hc xn") ||
!qwen4_bind_tensor(&b[2], mixed, (uint64_t)n_tokens * n_embd * sizeof(float), "hc mixed")) {
return 0;
}
return qwen4_dispatch(QWEN4_K_HC_MIX_ROWS, &args, sizeof(args), b, 3,
MTLSizeMake((n_embd + 255) / 256, n_tokens, 1), MTLSizeMake(256, 1, 1), 0);
}
Reference in new issue View Git Blame Copy Permalink