Support DeepSeek V4 Flash Vision Experimental
This commit is contained in:
1 parent
7c18e738cd
commit
9bbe12e6b3
52 files changed
+4988
-100
No files matched your search
@@ -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
|
||||
|
||||
@@ -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
@@ -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");
|
||||
}
|
||||
Reference in new issue
Block a user