Support DeepSeek V4 Flash Vision Experimental

This commit is contained in:
Georg Bauer committed 2026-10-04 20:32:19 +02:00
1 parent 7c18e738cd
commit 9bbe12e6b3
52 files changed
+4988 -100

No files matched your search

+1
View File
@@ -2,6 +2,7 @@ MIT License
Copyright (c) 2026 The ds4.c authors
Copyright (c) 2023-2026 The ggml authors
Copyright (c) 2023 DeepSeek
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
+7
View File
@@ -12,3 +12,10 @@ interfaces and are not copies of DS4 application code.
DS4Server builds and bundles this local snapshot; it does not read a sibling
DS4 checkout. Rust owns the model, graph, session, sampling, and lifecycle.
Objective-C remains only at the platform Metal boundary.
The Vision Experimental primitives added for issue #100 were reviewed against
DS4 `0aaea5a238fb41a35106a551e73c8409dfb751ac`. The new Metal file and visual
routing kernel follow that reference; image preprocessing, encoder graph,
visual prompt masks and session identities remain in Rust. See
`docs/issue100-vision/README.md` for the source and validation receipts.
+210 -14
View File
@@ -4551,6 +4551,7 @@ static NSString *ds4_gpu_full_source(void) {
@[@"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"],
@@ -28972,7 +28973,8 @@ static int ds4_gpu_encode_flash_attention_decode_raw_batch_heads(
uint32_t window,
uint32_t n_head,
uint32_t head_dim,
bool noncausal) {
bool noncausal,
const uint16_t *visual_mask) {
if (head_dim != 512 || n_head == 0 || n_tokens == 0 ||
n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap) {
return 0;
@@ -29037,7 +29039,9 @@ static int ds4_gpu_encode_flash_attention_decode_raw_batch_heads(
return 0;
}
if (noncausal) {
if (visual_mask) {
memcpy([mask_buffer contents], visual_mask, mask_bytes);
} else if (noncausal) {
memset([mask_buffer contents], 0, mask_bytes);
} else {
ds4_gpu_fill_raw_decode_batch_mask((uint16_t *)[mask_buffer contents],
@@ -29192,7 +29196,8 @@ static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(
uint32_t window,
uint32_t ratio,
uint32_t n_head,
uint32_t head_dim) {
uint32_t head_dim,
const uint16_t *visual_mask) {
if (n_comp == 0) {
return ds4_gpu_encode_flash_attention_decode_raw_batch_heads(cb,
heads,
@@ -29208,7 +29213,7 @@ static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(
window,
n_head,
head_dim,
false);
false, visual_mask);
}
if (head_dim != 512 || n_head == 0 || n_tokens == 0 ||
n_raw == 0 || raw_cap < n_raw || raw_start >= raw_cap ||
@@ -29290,6 +29295,9 @@ static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(
return 0;
}
if (visual_mask) {
memcpy([mask_buffer contents], visual_mask, mask_bytes);
} else {
ds4_gpu_fill_mixed_decode_batch_mask((uint16_t *)[mask_buffer contents],
n_tokens,
n_raw,
@@ -29297,6 +29305,7 @@ static int ds4_gpu_encode_flash_attention_decode_mixed_batch_heads(
pos0,
window,
ratio);
}
if (use_comp_mask) {
if (!ds4_gpu_encode_cpy_f32_f16_2d(cb,
maskbuf,
@@ -29618,7 +29627,7 @@ int ds4_gpu_attention_decode_raw_batch_heads_tensor(
window,
n_head,
head_dim,
false)) {
false, NULL)) {
return 0;
}
@@ -29680,7 +29689,7 @@ int ds4_gpu_attention_noncausal_raw_batch_heads_tensor(
0,
n_head,
head_dim,
true)) {
true, NULL)) {
return 0;
}
@@ -29755,7 +29764,7 @@ int ds4_gpu_attention_decode_mixed_batch_heads_tensor(
window,
ratio,
n_head,
head_dim)) {
head_dim, NULL)) {
return 0;
}
@@ -33034,7 +33043,8 @@ static int ds4_gpu_encode_router_select(
uint32_t n_expert_used,
float expert_weight_scale,
bool has_bias,
bool hash_mode) {
bool hash_mode,
const ds4_gpu_tensor *visual_bias, uint32_t vocab_size) {
id<MTLBuffer> selectedbuf = ds4_gpu_tensor_buffer(selected);
id<MTLBuffer> weightsbuf = ds4_gpu_tensor_buffer(weights);
id<MTLBuffer> probsbuf = ds4_gpu_tensor_buffer(probs);
@@ -33052,7 +33062,7 @@ static int ds4_gpu_encode_router_select(
fabsf(expert_weight_scale - 1.5f) <= 1.0e-6f;
int ok = 0;
if (flash_router_fast_path &&
if (!visual_bias && flash_router_fast_path &&
!g_quality_mode && n_tokens == 1 &&
getenv("DS4_METAL_DISABLE_ROUTER_SELECT_FUSION") == NULL) {
const bool pre_m5_device =
@@ -33211,7 +33221,7 @@ static int ds4_gpu_encode_router_select(
return 1;
}
if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) {
if (!visual_bias && flash_router_fast_path && !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");
@@ -33253,7 +33263,28 @@ static int ds4_gpu_encode_router_select(
}
if (!ok) return 0;
if (hash_mode) {
if (visual_bias) {
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline("kernel_dsv4_router_select_visual_batch");
if (!pipeline || !tokensbuf || n_expert != 256 || n_expert_used != 6 ||
ds4_gpu_tensor_bytes(visual_bias) < 256u * sizeof(float)) return 0;
const uint32_t args[5] = { hash_rows, vocab_size, n_tokens, has_bias, hash_mode };
const uint32_t zero = 0;
id<MTLComputeCommandEncoder> enc = ds4_gpu_compute_encoder(cb);
[enc setComputePipelineState: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 length:sizeof(zero) atIndex:2];
[enc setBuffer:ds4_gpu_tensor_buffer(visual_bias) offset:ds4_gpu_tensor_offset(visual_bias) atIndex:3];
if (hash_mode) [enc setBuffer:hashbuf offset:hash_off atIndex:4];
else [enc setBytes:&zero length:sizeof(zero) atIndex:4];
[enc setBuffer:tokensbuf offset:tokens_off atIndex:5];
[enc setBuffer:selectedbuf offset:selected_off atIndex:6];
[enc setThreadgroupMemoryLength:2048 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,
@@ -33323,7 +33354,7 @@ static int ds4_gpu_encode_router_select(
return 1;
}
if (flash_router_fast_path && !g_quality_mode && n_tokens == 1) {
if (!visual_bias && 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");
@@ -39168,7 +39199,7 @@ int ds4_gpu_router_select_tensor(
n_expert_used,
expert_weight_scale,
has_bias && !hash_mode,
hash_mode);
hash_mode, NULL, 0);
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
@@ -39264,7 +39295,104 @@ int ds4_gpu_router_select_batch_tensor(
n_expert_used,
expert_weight_scale,
has_bias && !hash_mode,
hash_mode);
hash_mode, NULL, 0);
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
if (!ok) return 0;
}
return 1;
}
int ds4_gpu_router_select_visual_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,
const ds4_gpu_tensor *visual_bias, uint32_t vocab_size) {
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),
NULL,
hash_rows,
n_tokens,
n_expert,
n_expert_used,
expert_weight_scale,
has_bias && !hash_mode,
hash_mode, visual_bias, vocab_size);
if (!had_batch) {
ok = ds4_gpu_end_commands() != 0 && ok;
}
@@ -45809,3 +45937,71 @@ int ds4_gpu_glm53_kda_prefill(
void ds4_gpu_set_glm_mtp_verify_mode(bool enabled) {
(void)enabled;
}
// Primitive dispatch only. The DeepSeek image graph and all layouts live in Rust.
int ds4_gpu_deepseek_vision_dispatch(
const char *kernel, const ds4_gpu_metal_binding *bindings, uint32_t count,
const uint32_t grid[3], const uint32_t threads[3],
uint32_t individual_threads, uint32_t shared_bytes) {
if (!kernel || !bindings || count == 0 || count > 31 || !grid || !threads ||
!grid[0] || !grid[1] || !grid[2] || !threads[0] || !threads[1] || !threads[2]) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
id<MTLComputePipelineState> pipeline = ds4_gpu_get_pipeline(kernel);
if (!pipeline || (uint64_t)threads[0] * threads[1] * threads[2] > pipeline.maxTotalThreadsPerThreadgroup) return 0;
for (uint32_t i = 0; i < count; i++) {
const ds4_gpu_metal_binding *b = &bindings[i];
if (b->index > 30 || (b->tensor ? b->offset >= ds4_gpu_tensor_bytes(b->tensor)
: !b->bytes || b->size == 0 || b->size > 4096)) 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];
for (uint32_t i = 0; i < count; i++) {
const ds4_gpu_metal_binding *b = &bindings[i];
if (b->tensor) {
[enc setBuffer:ds4_gpu_tensor_buffer(b->tensor)
offset:ds4_gpu_tensor_offset(b->tensor) + b->offset atIndex:b->index];
} else {
[enc setBytes:b->bytes length:b->size atIndex:b->index];
}
}
if (shared_bytes) [enc setThreadgroupMemoryLength:shared_bytes atIndex:0];
if (individual_threads) {
[enc dispatchThreads:MTLSizeMake(grid[0],grid[1],grid[2])
threadsPerThreadgroup:MTLSizeMake(threads[0],threads[1],threads[2])];
} else {
[enc dispatchThreadgroups:MTLSizeMake(grid[0],grid[1],grid[2])
threadsPerThreadgroup:MTLSizeMake(threads[0],threads[1],threads[2])];
}
ds4_gpu_end_compute_encoder(cb, enc);
return ds4_gpu_finish_command_buffer(cb, owned, "DeepSeek vision dispatch");
}
// Rust constructs the model-specific visual mask; this bridge only binds it
// to the same FlashAttention primitive used by ordinary DeepSeek prefill.
int ds4_gpu_attention_visual_batch_tensor(
ds4_gpu_tensor *heads, const void *map, uint64_t size, uint64_t sinks,
const ds4_gpu_tensor *q, const ds4_gpu_tensor *raw,
const ds4_gpu_tensor *comp, const ds4_gpu_tensor *comp_mask,
uint32_t use_mask, uint32_t rows, uint32_t pos,
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 uint16_t *mask, uint64_t mask_elements) {
if (!map || !heads || !q || !raw || !mask || !rows || !n_raw ||
raw_cap < n_raw || raw_start >= raw_cap || (n_comp && (!comp || !ratio)) ||
(use_mask && !comp_mask) || mask_elements != (uint64_t)rows * ((uint64_t)n_raw+n_comp) ||
sinks > size || (uint64_t)n_head * sizeof(float) > size-sinks) return 0;
if (!g_initialized && !ds4_gpu_init()) return 0;
uint64_t offset = 0;
id<MTLBuffer> weight = ds4_gpu_wrap_model_range(map,size,sinks,(uint64_t)n_head*sizeof(float),&offset);
if (!weight) 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,weight,offset,q,raw,comp,1,comp_mask,use_mask,
rows,pos,n_raw,raw_cap,raw_start,n_comp,window,ratio,n_head,head_dim,mask)) return 0;
return ds4_gpu_finish_command_buffer(cb,owned,"DeepSeek visual attention");
}