From a68112a74bb2a96e3a25a62bc145f81c273ef13c Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 11 Jul 2026 09:12:26 +0200 Subject: [PATCH 1/4] fix(deepseek4): hand off full HC state between local shards --- server/src/deepseek4/deepseek4_graph.cpp | 27 ++++++++++++++++--- .../deepseek4_layer_split_adapter.cpp | 22 +++++++-------- .../deepseek4/deepseek4_layer_split_adapter.h | 4 +++ server/tests/test_deepseek4_unit.cpp | 16 +++++++++++ 4 files changed, 52 insertions(+), 17 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 3f81ba9aa..ce4ba65f2 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -4255,9 +4255,21 @@ bool deepseek4_step_layer_range( // Initialize HC state. // First shard (layer_begin=0): embed is token embeddings [n_embd × n_tokens], // replicate into n_hc streams. - // Later shards: embed is full HC state [hc_dim × n_tokens] from previous shard. - if (hc_state.size() != (size_t)hc_dim * (size_t)n_tokens) { - hc_state.resize((size_t)hc_dim * (size_t)n_tokens); + // Later shards: embed is full HC state [hc_dim × n_tokens] from previous + // shard — either hc_state's own buffer (local in-process handoff) or a + // separate buffer (IPC daemon). Detect the alias before any resize, which + // would invalidate embed. + const size_t hc_state_elems = (size_t)hc_dim * (size_t)n_tokens; + const bool boundary_in_place = embed != nullptr && embed == hc_state.data(); + if (hc_state.size() != hc_state_elems) { + if (boundary_in_place) { + std::fprintf(stderr, + "[deepseek4] HC boundary state size mismatch for layer range [%d,%d): " + "have %zu want %zu\n", + layer_begin, layer_end, hc_state.size(), hc_state_elems); + return false; + } + hc_state.resize(hc_state_elems); } if (layer_begin == 0) { // First shard: replicate embedding into all HC streams @@ -4269,7 +4281,14 @@ bool deepseek4_step_layer_range( } } else { // Later shard: embed contains full HC state from previous shard - memcpy(hc_state.data(), embed, sizeof(float) * (size_t)hc_dim * (size_t)n_tokens); + if (!embed) { + std::fprintf(stderr, "[deepseek4] missing HC boundary state for layer range [%d,%d)\n", + layer_begin, layer_end); + return false; + } + if (!boundary_in_place) { + memcpy(hc_state.data(), embed, sizeof(float) * hc_state_elems); + } } // Lazy-load per-layer HC weights on CPU (static to avoid reloading) diff --git a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp index a6054aa8b..28a4d2157 100644 --- a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp +++ b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp @@ -124,6 +124,13 @@ size_t DeepSeek4LayerSplitAdapter::hc_state_elements( return (size_t) weights.n_hc * (size_t) weights.n_embd; } +const float * DeepSeek4LayerSplitAdapter::local_shard_input( + size_t shard_index, + const std::vector & token_embeddings, + const std::vector & hc_state) { + return shard_index == 0 ? token_embeddings.data() : hc_state.data(); +} + int DeepSeek4LayerSplitAdapter::compute_auto_split_layers() const { // Check env override first int override_layers = env_int("DFLASH_DS4_CUDA_LAYERS", -1); @@ -440,7 +447,6 @@ bool DeepSeek4LayerSplitAdapter::run_forward( const auto forward_t0 = SplitClock::now(); const int n_tokens = (int)tokens.size(); const int n_embd = shards_[0].weights.n_embd; - const int n_hc = shards_[0].weights.n_hc; // Embed tokens on first shard auto & first_shard = shards_[0]; @@ -453,17 +459,6 @@ bool DeepSeek4LayerSplitAdapter::run_forward( DeepSeek4StepTelemetry tel_acc; if (timing) tel_acc.embed_us = split_elapsed_us(embed_t0, SplitClock::now()); - // Initialize HC state from embedding for new sequences - if (base_pos == 0 && cur_pos_ == 0) { - for (int t = 0; t < n_tokens; ++t) { - for (int h = 0; h < n_hc; ++h) { - std::memcpy(hc_state_.data() + (size_t)h * n_embd, - embed.data() + (size_t)t * n_embd, - (size_t)n_embd * sizeof(float)); - } - } - } - // If using mixed target split (remote Halo shard), delegate to that path if (use_mixed_target_split()) { return run_mixed_forward(tokens, base_pos, last_tok, logits_out); @@ -476,10 +471,11 @@ bool DeepSeek4LayerSplitAdapter::run_forward( const bool is_last = (si == shards_.size() - 1); std::vector * shard_logits = is_last ? logits_out : nullptr; DeepSeek4StepTelemetry step_tel; + const float * shard_input = local_shard_input(si, embed, hc_state_); if (!deepseek4_step_layer_range( shard.backend, shard.weights, shard.cache, - hc_state_, embed.data(), n_tokens, base_pos, + hc_state_, shard_input, n_tokens, base_pos, shard.layer_begin, shard.layer_end, shard_logits, tokens.data(), timing ? &step_tel : nullptr)) { std::fprintf(stderr, "[deepseek4-split] forward failed on shard %zu\n", si); diff --git a/server/src/deepseek4/deepseek4_layer_split_adapter.h b/server/src/deepseek4/deepseek4_layer_split_adapter.h index 67072d36a..82930d344 100644 --- a/server/src/deepseek4/deepseek4_layer_split_adapter.h +++ b/server/src/deepseek4/deepseek4_layer_split_adapter.h @@ -83,6 +83,10 @@ class DeepSeek4LayerSplitAdapter : public LayerSplitAdapter { int compute_auto_split_layers() const; static int estimate_cuda_layers_from_free_bytes(size_t free_bytes); static size_t hc_state_elements(const DeepSeek4Weights & weights); + static const float * local_shard_input( + size_t shard_index, + const std::vector & token_embeddings, + const std::vector & hc_state); DeepSeek4LayerSplitAdapterConfig cfg_; std::vector shards_; diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index bb32aeed2..f98652d6c 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -621,6 +621,21 @@ static void test_hc_state_dimensions() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_local_shards_receive_hc_boundary_state() { + std::fprintf(stderr, " test_local_shards_receive_hc_boundary_state ..."); + + const std::vector embeddings = {1.0f, 2.0f, 3.0f, 4.0f}; + const std::vector hc_state = {5.0f, 6.0f, 7.0f, 8.0f}; + TEST_ASSERT(DeepSeek4LayerSplitAdapter::local_shard_input( + 0, embeddings, hc_state) == embeddings.data()); + TEST_ASSERT(DeepSeek4LayerSplitAdapter::local_shard_input( + 1, embeddings, hc_state) == hc_state.data()); + TEST_ASSERT(DeepSeek4LayerSplitAdapter::local_shard_input( + 2, embeddings, hc_state) == hc_state.data()); + + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_loader_rejects_missing_required_metadata(ggml_backend_t backend) { std::fprintf(stderr, " test_loader_rejects_missing_required_metadata ..."); @@ -1558,6 +1573,7 @@ int main() { test_auto_split_computation(); test_layer_range_validation(); test_hc_state_dimensions(); + test_local_shards_receive_hc_boundary_state(); test_loader_rejects_missing_required_metadata(backend); test_loader_rejects_invalid_compress_ratio_type(backend); test_loader_rejects_zero_vocab_size(backend); From b6e50aa378339085b60e3fd5aef4ac8739591d85 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Sat, 11 Jul 2026 09:13:57 +0200 Subject: [PATCH 2/4] fix(deepseek4): scope HC runtime caches per shard --- server/src/deepseek4/deepseek4_graph.cpp | 266 ++++++++++++++++------ server/src/deepseek4/deepseek4_internal.h | 8 + 2 files changed, 207 insertions(+), 67 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index ce4ba65f2..c80365823 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2829,18 +2829,20 @@ static bool load_tensor_to_f16_cpu(std::vector & dst, ggml_tensor * t) return true; } -static void load_hc_weights_cpu(HcWeightsCpu & dst, ggml_tensor * fn, - ggml_tensor * scale, ggml_tensor * base) { - if (!fn || !scale || !base || dst.loaded) return; +static bool load_hc_weights_cpu(HcWeightsCpu & dst, ggml_tensor * fn, + ggml_tensor * scale, ggml_tensor * base) { + if (dst.loaded) return true; + if (!fn || !scale || !base) return false; if (!load_tensor_to_f16_cpu(dst.fn_data, fn) || !load_tensor_to_f32_cpu(dst.scale_data, scale) || !load_tensor_to_f32_cpu(dst.base_data, base)) { dst.fn_data.clear(); dst.scale_data.clear(); dst.base_data.clear(); - return; + return false; } dst.loaded = true; + return true; } static void release_hc_fn_device(HcWeightsCpu & w) { @@ -4226,6 +4228,171 @@ static int ds4_try_fused_decode_step( return 1; } +struct DeepSeek4LayerRangeRuntimeCache { + ~DeepSeek4LayerRangeRuntimeCache() { reset(); } + + const DeepSeek4Weights * owner_weights = nullptr; + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + int n_layer = 0; + int layer_begin = -1; + int layer_end = -1; + bool owns_output = false; + + std::vector hc_layer_weights; + HcWeightsCpu hc_output_weights; + std::vector hash_routing_tables; + std::vector cached_attn_allocs; + std::vector cached_decode_attn_hc_pre_graphs; + std::vector cached_decode_ffn_hc_pre_graphs; + DeepSeek4CachedDecodeHcPostGraph cached_decode_hc_post_graph; + std::vector> cached_decode_attn_graphs; + std::vector cached_decode_ffn_graphs; + DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; + DeepSeek4CachedLayerAlloc cached_dynamic_output_alloc; + DeepSeek4FusedDecodeCache fused_decode_graph_cache; + Ds4DecodeSharedInputs decode_shared_inputs; + + bool matches(const DeepSeek4Weights & w, + ggml_backend_t candidate_backend, + int candidate_begin, + int candidate_end, + bool candidate_owns_output) const { + return owner_weights == &w && + owner_ctx == w.ctx && + backend == candidate_backend && + n_layer == w.n_layer && + layer_begin == candidate_begin && + layer_end == candidate_end && + owns_output == candidate_owns_output; + } + + void reset() { + reset_hc_layer_weights_cpu(hc_layer_weights); + reset_hc_weights_cpu(hc_output_weights); + hash_routing_tables.clear(); + for (auto & alloc : cached_attn_allocs) alloc.free(); + cached_attn_allocs.clear(); + for (auto & graph : cached_decode_attn_hc_pre_graphs) graph.free(); + cached_decode_attn_hc_pre_graphs.clear(); + for (auto & graph : cached_decode_ffn_hc_pre_graphs) graph.free(); + cached_decode_ffn_hc_pre_graphs.clear(); + cached_decode_hc_post_graph.free(); + for (auto & per_layer : cached_decode_attn_graphs) { + for (auto & graph : per_layer) graph.free(); + } + cached_decode_attn_graphs.clear(); + for (auto & graph : cached_decode_ffn_graphs) graph.free(); + cached_decode_ffn_graphs.clear(); + cached_decode_output_graph.free(); + cached_dynamic_output_alloc.free(); + fused_decode_graph_cache.destroy(); + decode_shared_inputs.free(); + owner_weights = nullptr; + owner_ctx = nullptr; + backend = nullptr; + n_layer = 0; + layer_begin = -1; + layer_end = -1; + owns_output = false; + } +}; + +static bool ds4_hc_layer_weights_ready(const HcWeightsCpu & weights, + int n_embd, + int n_hc) { + const size_t hc_dim = (size_t)n_embd * (size_t)n_hc; + const size_t mix_dim = (size_t)(2 * n_hc + n_hc * n_hc); + return weights.loaded && + weights.fn_data.size() >= hc_dim * mix_dim && + weights.scale_data.size() >= 3 && + weights.base_data.size() >= mix_dim; +} + +static bool ds4_hc_output_weights_ready(const HcWeightsCpu & weights, + int n_embd, + int n_hc) { + const size_t hc_dim = (size_t)n_embd * (size_t)n_hc; + return weights.loaded && + weights.fn_data.size() >= hc_dim * (size_t)n_hc && + !weights.scale_data.empty() && + weights.base_data.size() >= (size_t)n_hc; +} + +static bool initialize_layer_range_runtime_cache( + DeepSeek4LayerRangeRuntimeCache & runtime, + ggml_backend_t backend, + const DeepSeek4Weights & w, + int layer_begin, + int layer_end, + bool owns_output) { + runtime.reset(); + if (layer_begin < 0 || layer_end < layer_begin || layer_end > w.n_layer) { + std::fprintf(stderr, + "[deepseek4] invalid HC cache layer range [%d,%d) for %d layers\n", + layer_begin, layer_end, w.n_layer); + return false; + } + + runtime.hc_layer_weights.resize((size_t)w.n_layer); + runtime.hash_routing_tables.assign((size_t)w.n_layer, {}); + runtime.cached_attn_allocs.assign((size_t)w.n_layer, {}); + runtime.cached_decode_attn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + runtime.cached_decode_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + runtime.cached_decode_attn_graphs.assign((size_t)w.n_layer, {}); + runtime.cached_decode_ffn_graphs.assign((size_t)w.n_layer, {}); + + for (int il = layer_begin; il < layer_end; ++il) { + const DeepSeek4Layer & layer = w.layers[(size_t)il]; + HcLayerWeightsCpu & cached = runtime.hc_layer_weights[(size_t)il]; + const bool attn_loaded = load_hc_weights_cpu( + cached.attn, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base); + const bool ffn_loaded = load_hc_weights_cpu( + cached.ffn, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base); + if (!attn_loaded || !ds4_hc_layer_weights_ready(cached.attn, w.n_embd, w.n_hc)) { + std::fprintf(stderr, + "[deepseek4] missing or invalid HC attention weights for layer %d\n", il); + runtime.reset(); + return false; + } + if (!ffn_loaded || !ds4_hc_layer_weights_ready(cached.ffn, w.n_embd, w.n_hc)) { + std::fprintf(stderr, + "[deepseek4] missing or invalid HC FFN weights for layer %d\n", il); + runtime.reset(); + return false; + } + if (ds4_backend_is_gpu(backend) && !ds4_env_flag("DFLASH_DS4_HC_CPU")) { + ensure_hc_fn_device(cached.attn, layer.hc_attn_fn); + ensure_hc_fn_device(cached.ffn, layer.hc_ffn_fn); + } + if (il < w.n_hash_layer && layer.ffn_gate_tid2eid) { + load_hash_routing_cpu(runtime.hash_routing_tables[(size_t)il], + layer.ffn_gate_tid2eid); + } + } + + if (owns_output) { + if (!load_hc_weights_cpu(runtime.hc_output_weights, + w.output_hc_fn, + w.output_hc_scale, + w.output_hc_base) || + !ds4_hc_output_weights_ready(runtime.hc_output_weights, w.n_embd, w.n_hc)) { + std::fprintf(stderr, "[deepseek4] missing or invalid HC output weights\n"); + runtime.reset(); + return false; + } + } + + runtime.owner_weights = &w; + runtime.owner_ctx = w.ctx; + runtime.backend = backend; + runtime.n_layer = w.n_layer; + runtime.layer_begin = layer_begin; + runtime.layer_end = layer_end; + runtime.owns_output = owns_output; + return true; +} + bool deepseek4_step_layer_range( ggml_backend_t backend, @@ -4291,71 +4458,34 @@ bool deepseek4_step_layer_range( } } - // Lazy-load per-layer HC weights on CPU (static to avoid reloading) - static std::vector hc_layer_weights_range; - static HcWeightsCpu hc_output_weights_range; - static std::vector hash_routing_tables_range; - static std::vector cached_attn_allocs; - static std::vector cached_decode_attn_hc_pre_graphs; - static std::vector cached_decode_ffn_hc_pre_graphs; - static DeepSeek4CachedDecodeHcPostGraph cached_decode_hc_post_graph; - static std::vector> cached_decode_attn_graphs; - static std::vector cached_decode_ffn_graphs; - static DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; - static DeepSeek4CachedLayerAlloc cached_dynamic_output_alloc; - static DeepSeek4FusedDecodeCache fused_decode_graph_cache; - static Ds4DecodeSharedInputs decode_shared_inputs; - static int hc_loaded_n_layer = 0; - static const ggml_context * hc_loaded_ctx = nullptr; - if (hc_loaded_n_layer != w.n_layer || hc_loaded_ctx != w.ctx) { - reset_hc_layer_weights_cpu(hc_layer_weights_range); - reset_hc_weights_cpu(hc_output_weights_range); - hc_layer_weights_range.resize((size_t)w.n_layer); - hash_routing_tables_range.assign((size_t)w.n_layer, {}); - for (auto & alloc : cached_attn_allocs) { - alloc.free(); - } - cached_attn_allocs.assign((size_t)w.n_layer, {}); - for (auto & g : cached_decode_attn_hc_pre_graphs) { - g.free(); - } - cached_decode_attn_hc_pre_graphs.assign((size_t)w.n_layer, {}); - for (auto & g : cached_decode_ffn_hc_pre_graphs) { - g.free(); - } - cached_decode_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); - cached_decode_hc_post_graph.free(); - for (auto & per_layer : cached_decode_attn_graphs) { - for (auto & g : per_layer) { - g.free(); - } - } - cached_decode_attn_graphs.assign((size_t)w.n_layer, {}); - for (auto & g : cached_decode_ffn_graphs) { - g.free(); - } - cached_decode_ffn_graphs.assign((size_t)w.n_layer, {}); - cached_decode_output_graph.free(); - cached_dynamic_output_alloc.free(); - fused_decode_graph_cache.destroy(); - decode_shared_inputs.free(); - for (int il = 0; il < w.n_layer; il++) { - const DeepSeek4Layer & L = w.layers[(size_t)il]; - load_hc_weights_cpu(hc_layer_weights_range[il].attn, L.hc_attn_fn, L.hc_attn_scale, L.hc_attn_base); - load_hc_weights_cpu(hc_layer_weights_range[il].ffn, L.hc_ffn_fn, L.hc_ffn_scale, L.hc_ffn_base); - if (ds4_backend_is_gpu(backend) && !ds4_env_flag("DFLASH_DS4_HC_CPU")) { - ensure_hc_fn_device(hc_layer_weights_range[il].attn, L.hc_attn_fn); - ensure_hc_fn_device(hc_layer_weights_range[il].ffn, L.hc_ffn_fn); - } - if (il < w.n_hash_layer && L.ffn_gate_tid2eid) { - load_hash_routing_cpu(hash_routing_tables_range[(size_t)il], L.ffn_gate_tid2eid); - } - } - load_hc_weights_cpu(hc_output_weights_range, w.output_hc_fn, w.output_hc_scale, w.output_hc_base); - hc_loaded_n_layer = w.n_layer; - hc_loaded_ctx = w.ctx; + // Each partial model owns only its assigned layer range. Its HC and graph + // caches live on the shard's DeepSeek4Cache — created with the shard, + // freed with it — instead of one process-global cache being reset as + // local/remote contexts alternate. + if (!cache.layer_range_runtime) { + cache.layer_range_runtime = new DeepSeek4LayerRangeRuntimeCache(); + } + DeepSeek4LayerRangeRuntimeCache & runtime = *cache.layer_range_runtime; + if (!runtime.matches(w, backend, layer_begin, layer_end, is_last_shard) && + !initialize_layer_range_runtime_cache( + runtime, backend, w, layer_begin, layer_end, is_last_shard)) { + return false; } + auto & hc_layer_weights_range = runtime.hc_layer_weights; + auto & hc_output_weights_range = runtime.hc_output_weights; + auto & hash_routing_tables_range = runtime.hash_routing_tables; + auto & cached_attn_allocs = runtime.cached_attn_allocs; + auto & cached_decode_attn_hc_pre_graphs = runtime.cached_decode_attn_hc_pre_graphs; + auto & cached_decode_ffn_hc_pre_graphs = runtime.cached_decode_ffn_hc_pre_graphs; + auto & cached_decode_hc_post_graph = runtime.cached_decode_hc_post_graph; + auto & cached_decode_attn_graphs = runtime.cached_decode_attn_graphs; + auto & cached_decode_ffn_graphs = runtime.cached_decode_ffn_graphs; + auto & cached_decode_output_graph = runtime.cached_decode_output_graph; + auto & cached_dynamic_output_alloc = runtime.cached_dynamic_output_alloc; + auto & fused_decode_graph_cache = runtime.fused_decode_graph_cache; + auto & decode_shared_inputs = runtime.decode_shared_inputs; + // Per-layer execution with CPU-side HC static thread_local DeepSeek4LayerRangeScratch scratch; scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, w.n_expert_used); @@ -5076,6 +5206,8 @@ bool create_deepseek4_cache(ggml_backend_t backend, } void free_deepseek4_cache(DeepSeek4Cache & c) { + delete c.layer_range_runtime; + c.layer_range_runtime = nullptr; if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } c.layers.clear(); diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index a5d28d23c..a13831a92 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -254,6 +254,11 @@ struct DeepSeek4LayerCache { DeepSeek4CompressorState indexer_compressor; }; +// Per-shard runtime state for deepseek4_step_layer_range (host-side HC weight +// cache + cached decode graphs). Defined in deepseek4_graph.cpp; owned by the +// DeepSeek4Cache below and released by free_deepseek4_cache(). +struct DeepSeek4LayerRangeRuntimeCache; + struct DeepSeek4Cache { int cur_pos = 0; int max_ctx = 0; @@ -264,6 +269,9 @@ struct DeepSeek4Cache { // HC residual streams: [n_hc * n_embd] persistent state ggml_tensor * hc_state = nullptr; // [n_hc * n_embd] + // Lazily created on the first deepseek4_step_layer_range call. + DeepSeek4LayerRangeRuntimeCache * layer_range_runtime = nullptr; + ggml_context * ctx = nullptr; ggml_backend_buffer_t buf = nullptr; }; From ed21e25b6f3cbf0ee0db74606e4e9b936b9e1388 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 13 Jul 2026 16:54:41 +0200 Subject: [PATCH 3/4] fix(ds4): scope HC scratch by GPU device --- server/src/deepseek4/deepseek4_backend.cpp | 4 +- server/src/deepseek4/deepseek4_graph.cpp | 35 +++- server/src/deepseek4/deepseek4_hc_cuda.cu | 178 ++++++++++++------ server/src/deepseek4/deepseek4_hc_cuda.h | 9 +- server/src/deepseek4/deepseek4_internal.h | 1 + .../deepseek4_layer_split_adapter.cpp | 4 +- .../deepseek4_target_shard_ipc_daemon.cpp | 2 +- server/tests/test_deepseek4_unit.cpp | 87 +++++++++ 8 files changed, 252 insertions(+), 68 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 5a3edbf9b..ed1fb88fb 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -586,7 +586,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, routing_stats_.get()); } else { std::vector hc_state; - ok = deepseek4_step_layer_range(backend_, w_, cache_, hc_state, + ok = deepseek4_step_layer_range(backend_, cfg_.device.gpu, w_, cache_, hc_state, embed.data(), n_tok, pos, 0, w_.n_layer, &logits, tokens.data() + i, @@ -659,7 +659,7 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, routing_stats_.get()); } else { std::vector hc_state; - ok = deepseek4_step_layer_range(backend_, w_, cache_, hc_state, + ok = deepseek4_step_layer_range(backend_, cfg_.device.gpu, w_, cache_, hc_state, embed.data(), 1, pos, 0, w_.n_layer, &logits, &tok_to_eval, diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index c80365823..9e2f91656 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2209,6 +2209,7 @@ struct HcWeightsCpu { std::vector base_data; // [2*n_hc + n_hc*n_hc] void * fn_f16_device = nullptr; // Persistent F16 mirror for quantized HC fn tensors. size_t fn_f16_device_bytes = 0; + int fn_f16_device_id = -1; bool loaded = false; }; @@ -2848,11 +2849,12 @@ static bool load_hc_weights_cpu(HcWeightsCpu & dst, ggml_tensor * fn, static void release_hc_fn_device(HcWeightsCpu & w) { #if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) if (w.fn_f16_device) { - deepseek4_cuda_hc_free(w.fn_f16_device); + deepseek4_cuda_hc_free(w.fn_f16_device_id, w.fn_f16_device); } #endif w.fn_f16_device = nullptr; w.fn_f16_device_bytes = 0; + w.fn_f16_device_id = -1; } static void reset_hc_weights_cpu(HcWeightsCpu & w) { @@ -2871,7 +2873,7 @@ static void reset_hc_layer_weights_cpu(std::vector & weights) weights.clear(); } -static bool ensure_hc_fn_device(HcWeightsCpu & w, ggml_tensor * fn) { +static bool ensure_hc_fn_device(HcWeightsCpu & w, ggml_tensor * fn, int device) { if (!fn || !fn->data) return false; if (fn->type == GGML_TYPE_F16) return true; if (!ds4_rocmfpx_hc_gpu_enabled()) return false; @@ -2879,16 +2881,19 @@ static bool ensure_hc_fn_device(HcWeightsCpu & w, ggml_tensor * fn) { #if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) const size_t bytes = w.fn_data.size() * sizeof(uint16_t); - if (w.fn_f16_device && w.fn_f16_device_bytes == bytes) { + if (w.fn_f16_device && w.fn_f16_device_bytes == bytes && + w.fn_f16_device_id == device) { return true; } release_hc_fn_device(w); - if (!deepseek4_cuda_hc_upload_f16(w.fn_data.data(), bytes, &w.fn_f16_device)) { + if (!deepseek4_cuda_hc_upload_f16(device, w.fn_data.data(), bytes, &w.fn_f16_device)) { w.fn_f16_device = nullptr; w.fn_f16_device_bytes = 0; + w.fn_f16_device_id = -1; return false; } w.fn_f16_device_bytes = bytes; + w.fn_f16_device_id = device; return true; #else return false; @@ -4234,6 +4239,7 @@ struct DeepSeek4LayerRangeRuntimeCache { const DeepSeek4Weights * owner_weights = nullptr; const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; + int device = -1; int n_layer = 0; int layer_begin = -1; int layer_end = -1; @@ -4255,12 +4261,14 @@ struct DeepSeek4LayerRangeRuntimeCache { bool matches(const DeepSeek4Weights & w, ggml_backend_t candidate_backend, + int candidate_device, int candidate_begin, int candidate_end, bool candidate_owns_output) const { return owner_weights == &w && owner_ctx == w.ctx && backend == candidate_backend && + device == candidate_device && n_layer == w.n_layer && layer_begin == candidate_begin && layer_end == candidate_end && @@ -4291,6 +4299,7 @@ struct DeepSeek4LayerRangeRuntimeCache { owner_weights = nullptr; owner_ctx = nullptr; backend = nullptr; + device = -1; n_layer = 0; layer_begin = -1; layer_end = -1; @@ -4322,6 +4331,7 @@ static bool ds4_hc_output_weights_ready(const HcWeightsCpu & weights, static bool initialize_layer_range_runtime_cache( DeepSeek4LayerRangeRuntimeCache & runtime, ggml_backend_t backend, + int device, const DeepSeek4Weights & w, int layer_begin, int layer_end, @@ -4362,8 +4372,8 @@ static bool initialize_layer_range_runtime_cache( return false; } if (ds4_backend_is_gpu(backend) && !ds4_env_flag("DFLASH_DS4_HC_CPU")) { - ensure_hc_fn_device(cached.attn, layer.hc_attn_fn); - ensure_hc_fn_device(cached.ffn, layer.hc_ffn_fn); + ensure_hc_fn_device(cached.attn, layer.hc_attn_fn, device); + ensure_hc_fn_device(cached.ffn, layer.hc_ffn_fn, device); } if (il < w.n_hash_layer && layer.ffn_gate_tid2eid) { load_hash_routing_cpu(runtime.hash_routing_tables[(size_t)il], @@ -4386,6 +4396,7 @@ static bool initialize_layer_range_runtime_cache( runtime.owner_weights = &w; runtime.owner_ctx = w.ctx; runtime.backend = backend; + runtime.device = device; runtime.n_layer = w.n_layer; runtime.layer_begin = layer_begin; runtime.layer_end = layer_end; @@ -4396,6 +4407,7 @@ static bool initialize_layer_range_runtime_cache( bool deepseek4_step_layer_range( ggml_backend_t backend, + int device, const DeepSeek4Weights & w, DeepSeek4Cache & cache, std::vector & hc_state, @@ -4410,6 +4422,13 @@ bool deepseek4_step_layer_range( bool allow_decode_graph_reuse) { const auto step_t0 = Ds4TimingClock::now(); + if (!deepseek4_cuda_hc_set_device(device)) { + std::fprintf(stderr, + "[deepseek4] failed to select HC device %d for layer range [%d,%d)\n", + device, layer_begin, layer_end); + return false; + } + // NOTE: The old deepseek4_step() lacks HC implementation. // Always use the HC-enabled layer_range path below. @@ -4466,9 +4485,9 @@ bool deepseek4_step_layer_range( cache.layer_range_runtime = new DeepSeek4LayerRangeRuntimeCache(); } DeepSeek4LayerRangeRuntimeCache & runtime = *cache.layer_range_runtime; - if (!runtime.matches(w, backend, layer_begin, layer_end, is_last_shard) && + if (!runtime.matches(w, backend, device, layer_begin, layer_end, is_last_shard) && !initialize_layer_range_runtime_cache( - runtime, backend, w, layer_begin, layer_end, is_last_shard)) { + runtime, backend, device, w, layer_begin, layer_end, is_last_shard)) { return false; } diff --git a/server/src/deepseek4/deepseek4_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index b5e592f50..9776feec9 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -21,6 +22,7 @@ constexpr int kMaxHc = 8; constexpr int kMaxMixDim = 2 * kMaxHc + kMaxHc * kMaxHc; struct HcCudaScratch { + int owner_device = -1; float * d_state = nullptr; float * d_sums = nullptr; float * d_mix = nullptr; @@ -31,7 +33,12 @@ struct HcCudaScratch { float * d_comb = nullptr; size_t state_cap = 0; + explicit HcCudaScratch(int device = -1) : owner_device(device) {} + ~HcCudaScratch() { + if (owner_device >= 0) { + (void) cudaSetDevice(owner_device); + } if (d_state) cudaFree(d_state); if (d_sums) cudaFree(d_sums); if (d_mix) cudaFree(d_mix); @@ -63,8 +70,38 @@ struct HcCudaScratch { } }; -std::mutex g_mu; -HcCudaScratch g_scratch; +struct HcScratchSlot { + explicit HcScratchSlot(int device = -1) : scratch(device) {} + + std::mutex mutex; + HcCudaScratch scratch; +}; + +void hc_log_cuda_error(const char * label, cudaError_t err); + +// CUDA and HIP allocations belong to the device that was current when they +// were created. Keep one independently locked slot per logical device so local +// shards cannot reuse another device's scratch pointers. +std::mutex g_scratch_slots_mutex; +std::vector> g_scratch_slots; + +HcScratchSlot * current_scratch_slot() { + int device = -1; + const cudaError_t err = cudaGetDevice(&device); + if (err != cudaSuccess || device < 0) { + hc_log_cuda_error("get current device", err); + return nullptr; + } + + std::lock_guard lock(g_scratch_slots_mutex); + if (g_scratch_slots.size() <= (size_t) device) { + g_scratch_slots.resize((size_t) device + 1); + } + if (!g_scratch_slots[(size_t) device]) { + g_scratch_slots[(size_t) device] = std::make_unique(device); + } + return g_scratch_slots[(size_t) device].get(); +} bool hc_direct_no_sync_enabled() { static int enabled = -1; @@ -256,7 +293,8 @@ __global__ void hc_finish_kernel(const float * hc_state, } } -bool hc_pre_device_locked(const void * hc_state_device, +bool hc_pre_device_locked(HcCudaScratch & scratch, + const void * hc_state_device, const void * fn_device, const void * scale_device, const void * base_device, @@ -277,8 +315,8 @@ bool hc_pre_device_locked(const void * hc_state_device, return false; }; - if (hc_state_device != g_scratch.d_state) { - cudaError_t err = cudaMemcpy(g_scratch.d_state, hc_state_device, + if (hc_state_device != scratch.d_state) { + cudaError_t err = cudaMemcpy(scratch.d_state, hc_state_device, sizeof(float) * (size_t) hc_dim, cudaMemcpyDeviceToDevice); if (err != cudaSuccess) { @@ -287,28 +325,28 @@ bool hc_pre_device_locked(const void * hc_state_device, } hc_mix_norm_kernel<<>>( - g_scratch.d_state, + scratch.d_state, static_cast(fn_device), hc_dim, mix_dim, eps, - g_scratch.d_mix); + scratch.d_mix); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { return fail("mix kernel", err); } hc_finish_kernel<<<1, kThreads>>>( - g_scratch.d_state, - g_scratch.d_mix, + scratch.d_state, + scratch.d_mix, static_cast(scale_device), static_cast(base_device), n_embd, n_hc, sinkhorn_iters, - g_scratch.d_working, - g_scratch.d_post, - g_scratch.d_comb); + scratch.d_working, + scratch.d_post, + scratch.d_comb); err = cudaGetLastError(); if (err != cudaSuccess) { return fail("finish kernel", err); @@ -322,24 +360,24 @@ bool hc_pre_device_locked(const void * hc_state_device, } #endif - if (working_device != g_scratch.d_working) { - err = cudaMemcpy(working_device, g_scratch.d_working, + if (working_device != scratch.d_working) { + err = cudaMemcpy(working_device, scratch.d_working, sizeof(float) * (size_t) n_embd, cudaMemcpyDeviceToDevice); if (err != cudaSuccess) { return fail("copy working d2d", err); } } - if (post_device != g_scratch.d_post) { - err = cudaMemcpy(post_device, g_scratch.d_post, + if (post_device != scratch.d_post) { + err = cudaMemcpy(post_device, scratch.d_post, sizeof(float) * (size_t) n_hc, cudaMemcpyDeviceToDevice); if (err != cudaSuccess) { return fail("copy post d2d", err); } } - if (comb_device != g_scratch.d_comb) { - err = cudaMemcpy(comb_device, g_scratch.d_comb, + if (comb_device != scratch.d_comb) { + err = cudaMemcpy(comb_device, scratch.d_comb, sizeof(float) * (size_t) n_hc * (size_t) n_hc, cudaMemcpyDeviceToDevice); if (err != cudaSuccess) { @@ -356,6 +394,18 @@ bool hc_pre_device_locked(const void * hc_state_device, } // namespace +bool deepseek4_cuda_hc_set_device(int device) { + if (device < 0) { + return false; + } + const cudaError_t err = cudaSetDevice(device); + if (err != cudaSuccess) { + hc_log_cuda_error("set device", err); + return false; + } + return true; +} + bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, const void * fn_device, int n_embd, @@ -366,31 +416,34 @@ bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, return false; } const int hc_dim = n_embd * n_hc; - std::lock_guard lock(g_mu); - if (!g_scratch.ensure((size_t)hc_dim, (size_t)n_embd, (size_t)n_hc)) { + HcScratchSlot * slot = current_scratch_slot(); + if (!slot) return false; + std::lock_guard lock(slot->mutex); + HcCudaScratch & scratch = slot->scratch; + if (!scratch.ensure((size_t)hc_dim, (size_t)n_embd, (size_t)n_hc)) { return false; } - if (cudaMemcpy(g_scratch.d_state, hc_state_host, sizeof(float) * (size_t)hc_dim, + if (cudaMemcpy(scratch.d_state, hc_state_host, sizeof(float) * (size_t)hc_dim, cudaMemcpyHostToDevice) != cudaSuccess) { return false; } - hc_sumsq_kernel<<>>(g_scratch.d_state, hc_dim, g_scratch.d_sums); + hc_sumsq_kernel<<>>(scratch.d_state, hc_dim, scratch.d_sums); if (cudaGetLastError() != cudaSuccess) return false; std::vector sums(kSums); - if (cudaMemcpy(sums.data(), g_scratch.d_sums, sizeof(float) * sums.size(), + if (cudaMemcpy(sums.data(), scratch.d_sums, sizeof(float) * sums.size(), cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } float ss = 0.0f; for (float v : sums) ss += v; const float inv_rms = 1.0f / std::sqrt(ss / (float)hc_dim + eps); - hc_mix_kernel<<>>(g_scratch.d_state, + hc_mix_kernel<<>>(scratch.d_state, static_cast(fn_device), hc_dim, inv_rms, - g_scratch.d_mix); + scratch.d_mix); if (cudaGetLastError() != cudaSuccess) return false; - if (cudaMemcpy(mix_host, g_scratch.d_mix, sizeof(float) * kMixDim, + if (cudaMemcpy(mix_host, scratch.d_mix, sizeof(float) * kMixDim, cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } @@ -419,47 +472,51 @@ bool deepseek4_cuda_hc_pre(const float * hc_state_host, return false; } - std::lock_guard lock(g_mu); - if (!g_scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + HcScratchSlot * slot = current_scratch_slot(); + if (!slot) return false; + std::lock_guard lock(slot->mutex); + HcCudaScratch & scratch = slot->scratch; + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { return false; } - if (cudaMemcpy(g_scratch.d_state, hc_state_host, sizeof(float) * (size_t) hc_dim, + if (cudaMemcpy(scratch.d_state, hc_state_host, sizeof(float) * (size_t) hc_dim, cudaMemcpyHostToDevice) != cudaSuccess) { return false; } - if (cudaMemcpy(g_scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, + if (cudaMemcpy(scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, cudaMemcpyHostToDevice) != cudaSuccess) { return false; } - if (cudaMemcpy(g_scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, + if (cudaMemcpy(scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, cudaMemcpyHostToDevice) != cudaSuccess) { return false; } if (!hc_pre_device_locked( - g_scratch.d_state, + scratch, + scratch.d_state, fn_device, - g_scratch.d_scale, - g_scratch.d_base, + scratch.d_scale, + scratch.d_base, n_embd, n_hc, sinkhorn_iters, eps, - g_scratch.d_working, - g_scratch.d_post, - g_scratch.d_comb, + scratch.d_working, + scratch.d_post, + scratch.d_comb, false)) { return false; } - if (cudaMemcpy(working_host, g_scratch.d_working, sizeof(float) * (size_t) n_embd, + if (cudaMemcpy(working_host, scratch.d_working, sizeof(float) * (size_t) n_embd, cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } - if (cudaMemcpy(post_host, g_scratch.d_post, sizeof(float) * (size_t) n_hc, + if (cudaMemcpy(post_host, scratch.d_post, sizeof(float) * (size_t) n_hc, cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } - if (cudaMemcpy(comb_host, g_scratch.d_comb, sizeof(float) * (size_t) n_hc * (size_t) n_hc, + if (cudaMemcpy(comb_host, scratch.d_comb, sizeof(float) * (size_t) n_hc * (size_t) n_hc, cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } @@ -489,11 +546,15 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, return false; } - std::lock_guard lock(g_mu); - if (!g_scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + HcScratchSlot * slot = current_scratch_slot(); + if (!slot) return false; + std::lock_guard lock(slot->mutex); + HcCudaScratch & scratch = slot->scratch; + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { return false; } - return hc_pre_device_locked(hc_state_device, + return hc_pre_device_locked(scratch, + hc_state_device, fn_device, scale_device, base_device, @@ -529,30 +590,34 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, return false; } - std::lock_guard lock(g_mu); - if (!g_scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + HcScratchSlot * slot = current_scratch_slot(); + if (!slot) return false; + std::lock_guard lock(slot->mutex); + HcCudaScratch & scratch = slot->scratch; + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { std::fprintf(stderr, "[deepseek4-hc-direct] ensure failed\n"); return false; } - if (cudaMemcpy(g_scratch.d_state, hc_state_device, sizeof(float) * (size_t) hc_dim, + if (cudaMemcpy(scratch.d_state, hc_state_device, sizeof(float) * (size_t) hc_dim, cudaMemcpyDeviceToDevice) != cudaSuccess) { hc_log_cuda_error("copy state d2d", cudaGetLastError()); return false; } - if (cudaMemcpy(g_scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, + if (cudaMemcpy(scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, cudaMemcpyHostToDevice) != cudaSuccess) { hc_log_cuda_error("copy scale", cudaGetLastError()); return false; } - if (cudaMemcpy(g_scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, + if (cudaMemcpy(scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, cudaMemcpyHostToDevice) != cudaSuccess) { hc_log_cuda_error("copy base", cudaGetLastError()); return false; } - return hc_pre_device_locked(g_scratch.d_state, + return hc_pre_device_locked(scratch, + scratch.d_state, fn_device, - g_scratch.d_scale, - g_scratch.d_base, + scratch.d_scale, + scratch.d_base, n_embd, n_hc, sinkhorn_iters, @@ -563,12 +628,16 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, true); } -bool deepseek4_cuda_hc_upload_f16(const void * host_f16, +bool deepseek4_cuda_hc_upload_f16(int device, + const void * host_f16, size_t bytes, void ** device_out) { if (!host_f16 || bytes == 0 || !device_out) { return false; } + if (!deepseek4_cuda_hc_set_device(device)) { + return false; + } void * ptr = nullptr; cudaError_t err = cudaMalloc(&ptr, bytes); if (err != cudaSuccess) { @@ -585,8 +654,11 @@ bool deepseek4_cuda_hc_upload_f16(const void * host_f16, return true; } -void deepseek4_cuda_hc_free(void * device_ptr) { +void deepseek4_cuda_hc_free(int device, void * device_ptr) { if (device_ptr) { + if (!deepseek4_cuda_hc_set_device(device)) { + return; + } cudaFree(device_ptr); } } diff --git a/server/src/deepseek4/deepseek4_hc_cuda.h b/server/src/deepseek4/deepseek4_hc_cuda.h index b4ac596fe..b806f4c3a 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.h +++ b/server/src/deepseek4/deepseek4_hc_cuda.h @@ -5,6 +5,10 @@ namespace dflash::common { +// CUDA and HIP select devices per host thread. Call this at worker/shard entry +// before invoking the direct HC helpers. +bool deepseek4_cuda_hc_set_device(int device); + bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, const void * fn_device, int n_embd, @@ -48,10 +52,11 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, void * post_device, void * comb_device); -bool deepseek4_cuda_hc_upload_f16(const void * host_f16, +bool deepseek4_cuda_hc_upload_f16(int device, + const void * host_f16, size_t bytes, void ** device_out); -void deepseek4_cuda_hc_free(void * device_ptr); +void deepseek4_cuda_hc_free(int device, void * device_ptr); } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index a13831a92..9d76b86b6 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -333,6 +333,7 @@ bool deepseek4_step( bool deepseek4_step_layer_range( ggml_backend_t backend, + int device, const DeepSeek4Weights & w, DeepSeek4Cache & cache, std::vector & hc_state, diff --git a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp index 28a4d2157..b5126c85b 100644 --- a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp +++ b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp @@ -474,7 +474,7 @@ bool DeepSeek4LayerSplitAdapter::run_forward( const float * shard_input = local_shard_input(si, embed, hc_state_); if (!deepseek4_step_layer_range( - shard.backend, shard.weights, shard.cache, + shard.backend, shard.gpu, shard.weights, shard.cache, hc_state_, shard_input, n_tokens, base_pos, shard.layer_begin, shard.layer_end, shard_logits, tokens.data(), timing ? &step_tel : nullptr)) { @@ -524,7 +524,7 @@ bool DeepSeek4LayerSplitAdapter::run_mixed_forward( // Run only the local shard's layer range, getting hidden state output std::vector hidden_out; - if (!deepseek4_step_layer_range(local_shard.backend, local_shard.weights, + if (!deepseek4_step_layer_range(local_shard.backend, local_shard.gpu, local_shard.weights, local_shard.cache, hc_state_, embed.data(), n_tokens, base_pos, local_shard.layer_begin, local_shard.layer_end, diff --git a/server/src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp b/server/src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp index d6876db0b..d0d3b6dc2 100644 --- a/server/src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp +++ b/server/src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp @@ -238,7 +238,7 @@ int run_deepseek4_target_shard_ipc_daemon( "[deepseek4-target-shard] forward shard %zu: n_tokens=%d base_pos=%d layers=[%d,%d)\n", i, n_tokens, req.base_pos, shard.layer_begin, shard.layer_end); const bool ok = deepseek4_step_layer_range( - shard.backend, shard.weights, shard.cache, hc_state, + shard.backend, shard.gpu, shard.weights, shard.cache, hc_state, shard_input, n_tokens, req.base_pos, shard.layer_begin, shard.layer_end, is_last ? &logits : nullptr, diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index f98652d6c..c7d8db417 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include #include @@ -1556,6 +1557,91 @@ static void test_hc_pre_kernel_gpu() { ggml_free(ctx); ggml_backend_free(backend); } + +#if defined(GGML_USE_CUDA) || defined(GGML_USE_HIP) +static void test_hc_scratch_per_device() { + std::fprintf(stderr, " test_hc_scratch_per_device ..."); + if (ggml_backend_cuda_get_device_count() < 2) { + std::fprintf(stderr, " skipped (requires two GPU devices)\n"); + return; + } + + constexpr int n_embd = 32; + constexpr int n_hc = 4; + constexpr int sinkhorn_iters = 6; + constexpr float hc_eps = 1.0e-6f; + constexpr int mix_dim = 2 * n_hc + n_hc * n_hc; + constexpr int hc_dim = n_embd * n_hc; + + std::vector hc_state((size_t) hc_dim); + std::iota(hc_state.begin(), hc_state.end(), -16.0f); + for (float & value : hc_state) value *= 0.01f; + std::vector fn_f16((size_t) mix_dim * (size_t) hc_dim, + ggml_fp32_to_fp16(0.0f)); + std::vector scale((size_t) mix_dim, 0.0f); + scale[0] = 1.0f; + scale[1] = 1.0f; + scale[2] = 1.0f; + std::vector base((size_t) mix_dim, 0.0f); + + struct Result { + bool ok = false; + std::vector working; + std::vector post; + std::vector comb; + + Result(int embd, int hc) + : working((size_t) embd), + post((size_t) hc), + comb((size_t) hc * (size_t) hc) {} + }; + + auto run_on_worker = [&](int device, Result & result) { + std::thread worker([&, device]() { + void * fn_device = nullptr; + if (!deepseek4_cuda_hc_set_device(device) || + !deepseek4_cuda_hc_upload_f16( + device, fn_f16.data(), fn_f16.size() * sizeof(ggml_fp16_t), &fn_device)) { + return; + } + result.ok = deepseek4_cuda_hc_pre( + hc_state.data(), fn_device, scale.data(), base.data(), + n_embd, n_hc, sinkhorn_iters, hc_eps, + result.working.data(), result.post.data(), result.comb.data()); + deepseek4_cuda_hc_free(device, fn_device); + }); + worker.join(); + }; + + Result first_device(n_embd, n_hc); + Result second_device(n_embd, n_hc); + Result first_device_again(n_embd, n_hc); + run_on_worker(0, first_device); + run_on_worker(1, second_device); + run_on_worker(0, first_device_again); + + TEST_ASSERT_MSG(first_device.ok, "HC direct call failed on GPU device 0"); + TEST_ASSERT_MSG(second_device.ok, "HC direct call failed on GPU device 1"); + TEST_ASSERT_MSG(first_device_again.ok, "HC direct call failed after returning to GPU device 0"); + if (first_device.ok && second_device.ok && first_device_again.ok) { + for (int i = 0; i < n_embd; ++i) { + TEST_ASSERT(nearly_equal(first_device.working[(size_t) i], + second_device.working[(size_t) i])); + TEST_ASSERT(nearly_equal(first_device.working[(size_t) i], + first_device_again.working[(size_t) i])); + } + for (int i = 0; i < n_hc; ++i) { + TEST_ASSERT(nearly_equal(first_device.post[(size_t) i], + second_device.post[(size_t) i])); + } + for (int i = 0; i < n_hc * n_hc; ++i) { + TEST_ASSERT(nearly_equal(first_device.comb[(size_t) i], + second_device.comb[(size_t) i])); + } + } + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} +#endif #endif int main() { @@ -1588,6 +1674,7 @@ int main() { test_output_graph_reuse_microbench(backend); #if defined(GGML_USE_CUDA) || defined(GGML_USE_HIP) test_hc_pre_kernel_gpu(); + test_hc_scratch_per_device(); #endif ggml_backend_free(backend); From b6a14e629ccd60ae7a15812e226ab0b56d1d2b19 Mon Sep 17 00:00:00 2001 From: Graffioh Date: Mon, 13 Jul 2026 16:56:25 +0200 Subject: [PATCH 4/4] fix(ds4): propagate sampling and use GPU greedy argmax --- server/src/deepseek4/deepseek4_backend.cpp | 2 + server/src/deepseek4/deepseek4_graph.cpp | 129 +++++++++++------- server/src/deepseek4/deepseek4_internal.h | 5 +- .../deepseek4_layer_split_adapter.cpp | 22 ++- .../deepseek4_target_shard_ipc_daemon.cpp | 24 ++-- server/tests/test_deepseek4_unit.cpp | 46 +++++++ 6 files changed, 160 insertions(+), 68 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index ed1fb88fb..3eda43ba5 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -589,6 +589,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, ok = deepseek4_step_layer_range(backend_, cfg_.device.gpu, w_, cache_, hc_state, embed.data(), n_tok, pos, 0, w_.n_layer, &logits, + /*out_argmax=*/nullptr, tokens.data() + i, timing ? &step_tel : nullptr); } @@ -662,6 +663,7 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, ok = deepseek4_step_layer_range(backend_, cfg_.device.gpu, w_, cache_, hc_state, embed.data(), 1, pos, 0, w_.n_layer, &logits, + /*out_argmax=*/nullptr, &tok_to_eval, timing ? &step_tel : nullptr); } diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 9e2f91656..2f91b93b6 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -271,7 +271,8 @@ struct DeepSeek4CachedDecodeOutputGraph { bool valid() const { return owner_ctx && backend && n_tokens > 0 && - sg.ctx && sg.gf && sg.alloc && sg.hidden_input && sg.logits; + sg.ctx && sg.gf && sg.alloc && sg.hidden_input && sg.logits && + sg.argmax_tokens; } void free() { @@ -544,8 +545,11 @@ static bool build_cached_decode_output_graph( ggml_set_input(out.sg.hidden_input); ggml_tensor * normed = build_rms_norm(out.sg.ctx, out.sg.hidden_input, w.out_norm, w.rms_eps); out.sg.logits = ggml_mul_mat(out.sg.ctx, w.output, normed); + out.sg.argmax_tokens = ggml_argmax(out.sg.ctx, out.sg.logits); ggml_set_output(out.sg.logits); + ggml_set_output(out.sg.argmax_tokens); out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 1024, false); + ggml_build_forward_expand(out.sg.gf, out.sg.argmax_tokens); ggml_build_forward_expand(out.sg.gf, out.sg.logits); out.sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); @@ -3675,7 +3679,7 @@ struct DeepSeek4FusedDecodeGraph { } bool built() const { - return sg.ctx && sg.gf && logits; + return sg.ctx && sg.gf && logits && sg.argmax_tokens; } void destroy() { @@ -4022,7 +4026,11 @@ static bool ds4_build_fused_decode_graph( ggml_tensor * final_2d = ggml_reshape_2d(ctx, final_embd, n_embd, 1); ggml_tensor * out_normed = build_rms_norm(ctx, final_2d, w.out_norm, w.rms_eps); fg.logits = ggml_mul_mat(ctx, w.output, out_normed); + fg.sg.logits = fg.logits; + fg.sg.argmax_tokens = ggml_argmax(ctx, fg.logits); ggml_set_output(fg.logits); + ggml_set_output(fg.sg.argmax_tokens); + ggml_build_forward_expand(gf, fg.sg.argmax_tokens); ggml_build_forward_expand(gf, fg.logits); if (!fg.sg.alloc) { @@ -4037,8 +4045,9 @@ static bool ds4_build_fused_decode_graph( return true; } -// Returns 1 on success (out_logits filled), 0 to fall back to the per-layer -// path, -1 on a hard failure after cache state may have been touched. +// Returns 1 on success, 0 to fall back to the per-layer path, or -1 on a hard +// failure after cache state may have been touched. Full logits are copied only +// when requested; greedy decode reads the GPU-side argmax token instead. static int ds4_try_fused_decode_step( DeepSeek4FusedDecodeCache & fc, ggml_backend_t backend, @@ -4050,7 +4059,8 @@ static int ds4_try_fused_decode_step( std::vector & hash_scratch, const float * embed, int kv_start, - std::vector & out_logits, + std::vector * out_logits, + int32_t * out_argmax, const int32_t * token_ids, DeepSeek4StepTelemetry * telemetry) { if (fc.disabled) return 0; @@ -4224,17 +4234,23 @@ static int ds4_try_fused_decode_step( } if (telemetry) telemetry->full_graph_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); - // ── Read logits ───────────────────────────────────────────────── + // ── Read requested output ─────────────────────────────────────── const auto read_t0 = Ds4TimingClock::now(); - out_logits.resize((size_t) w.n_vocab); - ggml_backend_tensor_get(fg->logits, out_logits.data(), 0, - sizeof(float) * (size_t) w.n_vocab); + if (out_logits) { + out_logits->resize((size_t) w.n_vocab); + ggml_backend_tensor_get(fg->logits, out_logits->data(), 0, + sizeof(float) * (size_t) w.n_vocab); + } + if (out_argmax) { + ggml_backend_tensor_get(fg->sg.argmax_tokens, out_argmax, 0, + sizeof(*out_argmax)); + } if (telemetry) telemetry->full_graph_read_us += ds4_elapsed_us(read_t0, Ds4TimingClock::now()); return 1; } -struct DeepSeek4LayerRangeRuntimeCache { - ~DeepSeek4LayerRangeRuntimeCache() { reset(); } +struct DeepSeek4LayerRangeCache { + ~DeepSeek4LayerRangeCache() { reset(); } const DeepSeek4Weights * owner_weights = nullptr; const ggml_context * owner_ctx = nullptr; @@ -4328,8 +4344,8 @@ static bool ds4_hc_output_weights_ready(const HcWeightsCpu & weights, weights.base_data.size() >= (size_t)n_hc; } -static bool initialize_layer_range_runtime_cache( - DeepSeek4LayerRangeRuntimeCache & runtime, +static bool initialize_layer_range_cache( + DeepSeek4LayerRangeCache & runtime, ggml_backend_t backend, int device, const DeepSeek4Weights & w, @@ -4417,6 +4433,7 @@ bool deepseek4_step_layer_range( int layer_begin, int layer_end, std::vector * out_logits, + int32_t * out_argmax, const int32_t * token_ids, DeepSeek4StepTelemetry * telemetry, bool allow_decode_graph_reuse) { @@ -4428,6 +4445,7 @@ bool deepseek4_step_layer_range( device, layer_begin, layer_end); return false; } + if (out_argmax) *out_argmax = -1; // NOTE: The old deepseek4_step() lacks HC implementation. // Always use the HC-enabled layer_range path below. @@ -4446,9 +4464,9 @@ bool deepseek4_step_layer_range( // separate buffer (IPC daemon). Detect the alias before any resize, which // would invalidate embed. const size_t hc_state_elems = (size_t)hc_dim * (size_t)n_tokens; - const bool boundary_in_place = embed != nullptr && embed == hc_state.data(); + const bool embed_points_to_hc_state = embed != nullptr && embed == hc_state.data(); if (hc_state.size() != hc_state_elems) { - if (boundary_in_place) { + if (embed_points_to_hc_state) { std::fprintf(stderr, "[deepseek4] HC boundary state size mismatch for layer range [%d,%d): " "have %zu want %zu\n", @@ -4472,7 +4490,7 @@ bool deepseek4_step_layer_range( layer_begin, layer_end); return false; } - if (!boundary_in_place) { + if (!embed_points_to_hc_state) { memcpy(hc_state.data(), embed, sizeof(float) * hc_state_elems); } } @@ -4481,29 +4499,29 @@ bool deepseek4_step_layer_range( // caches live on the shard's DeepSeek4Cache — created with the shard, // freed with it — instead of one process-global cache being reset as // local/remote contexts alternate. - if (!cache.layer_range_runtime) { - cache.layer_range_runtime = new DeepSeek4LayerRangeRuntimeCache(); + if (!cache.layer_range_cache) { + cache.layer_range_cache = new DeepSeek4LayerRangeCache(); } - DeepSeek4LayerRangeRuntimeCache & runtime = *cache.layer_range_runtime; - if (!runtime.matches(w, backend, device, layer_begin, layer_end, is_last_shard) && - !initialize_layer_range_runtime_cache( - runtime, backend, device, w, layer_begin, layer_end, is_last_shard)) { + DeepSeek4LayerRangeCache & layer_range_cache = *cache.layer_range_cache; + if (!layer_range_cache.matches(w, backend, device, layer_begin, layer_end, is_last_shard) && + !initialize_layer_range_cache( + layer_range_cache, backend, device, w, layer_begin, layer_end, is_last_shard)) { return false; } - auto & hc_layer_weights_range = runtime.hc_layer_weights; - auto & hc_output_weights_range = runtime.hc_output_weights; - auto & hash_routing_tables_range = runtime.hash_routing_tables; - auto & cached_attn_allocs = runtime.cached_attn_allocs; - auto & cached_decode_attn_hc_pre_graphs = runtime.cached_decode_attn_hc_pre_graphs; - auto & cached_decode_ffn_hc_pre_graphs = runtime.cached_decode_ffn_hc_pre_graphs; - auto & cached_decode_hc_post_graph = runtime.cached_decode_hc_post_graph; - auto & cached_decode_attn_graphs = runtime.cached_decode_attn_graphs; - auto & cached_decode_ffn_graphs = runtime.cached_decode_ffn_graphs; - auto & cached_decode_output_graph = runtime.cached_decode_output_graph; - auto & cached_dynamic_output_alloc = runtime.cached_dynamic_output_alloc; - auto & fused_decode_graph_cache = runtime.fused_decode_graph_cache; - auto & decode_shared_inputs = runtime.decode_shared_inputs; + auto & hc_layer_weights_range = layer_range_cache.hc_layer_weights; + auto & hc_output_weights_range = layer_range_cache.hc_output_weights; + auto & hash_routing_tables_range = layer_range_cache.hash_routing_tables; + auto & cached_attn_allocs = layer_range_cache.cached_attn_allocs; + auto & cached_decode_attn_hc_pre_graphs = layer_range_cache.cached_decode_attn_hc_pre_graphs; + auto & cached_decode_ffn_hc_pre_graphs = layer_range_cache.cached_decode_ffn_hc_pre_graphs; + auto & cached_decode_hc_post_graph = layer_range_cache.cached_decode_hc_post_graph; + auto & cached_decode_attn_graphs = layer_range_cache.cached_decode_attn_graphs; + auto & cached_decode_ffn_graphs = layer_range_cache.cached_decode_ffn_graphs; + auto & cached_decode_output_graph = layer_range_cache.cached_decode_output_graph; + auto & cached_dynamic_output_alloc = layer_range_cache.cached_dynamic_output_alloc; + auto & fused_decode_graph_cache = layer_range_cache.fused_decode_graph_cache; + auto & decode_shared_inputs = layer_range_cache.decode_shared_inputs; // Per-layer execution with CPU-side HC static thread_local DeepSeek4LayerRangeScratch scratch; @@ -4511,13 +4529,13 @@ bool deepseek4_step_layer_range( std::vector fused_debug_logits; if (n_tokens == 1 && allow_decode_graph_reuse && layer_begin == 0 && is_last_shard && - out_logits && ds4_backend_is_gpu(backend) && ds4_fused_decode_enabled()) { + (out_logits || out_argmax) && ds4_backend_is_gpu(backend) && ds4_fused_decode_enabled()) { const int rc = ds4_try_fused_decode_step( fused_decode_graph_cache, backend, w, cache, hc_layer_weights_range, hc_output_weights_range, hash_routing_tables_range, scratch.hash_expert_ids, - embed, kv_start, *out_logits, token_ids, telemetry); + embed, kv_start, out_logits, out_argmax, token_ids, telemetry); if (rc < 0) return false; - if (rc > 0 && ds4_env_flag("DFLASH_DS4_FUSED_DEBUG")) { + if (rc > 0 && out_logits && ds4_env_flag("DFLASH_DS4_FUSED_DEBUG")) { // Debug: keep the fused logits, then fall through and run the // per-layer reference for the same token; compare at the end. fused_debug_logits = *out_logits; @@ -5025,7 +5043,7 @@ bool deepseek4_step_layer_range( } // ── Output: HC pre → norm → lm_head (or return hidden state) ──────── - if (is_last_shard && out_logits) { + if (is_last_shard && (out_logits || out_argmax)) { // Final HC pre for output const auto output_t0 = Ds4TimingClock::now(); std::vector & final_embd = scratch.final_embd; @@ -5051,9 +5069,16 @@ bool deepseek4_step_layer_range( if (ggml_backend_graph_compute(backend, cached_decode_output_graph.sg.gf) != GGML_STATUS_SUCCESS) { return false; } - out_logits->resize((size_t)w.n_vocab); - ggml_backend_tensor_get(cached_decode_output_graph.sg.logits, - out_logits->data(), 0, sizeof(float) * (size_t)w.n_vocab); + if (out_logits) { + out_logits->resize((size_t)w.n_vocab); + ggml_backend_tensor_get(cached_decode_output_graph.sg.logits, + out_logits->data(), 0, + sizeof(float) * (size_t)w.n_vocab); + } + if (out_argmax) { + ggml_backend_tensor_get(cached_decode_output_graph.sg.argmax_tokens, + out_argmax, 0, sizeof(*out_argmax)); + } } else { const size_t ctx_size = 16 * 1024 * 1024; ggml_init_params params{}; @@ -5067,8 +5092,11 @@ bool deepseek4_step_layer_range( ggml_set_input(inp); ggml_tensor * normed = build_rms_norm(ctx, inp, w.out_norm, w.rms_eps); ggml_tensor * logits = ggml_mul_mat(ctx, w.output, normed); + ggml_tensor * argmax_tokens = ggml_argmax(ctx, logits); ggml_set_output(logits); + ggml_set_output(argmax_tokens); ggml_cgraph * gf = ggml_new_graph_custom(ctx, 1024, false); + ggml_build_forward_expand(gf, argmax_tokens); ggml_build_forward_expand(gf, logits); if (!cached_dynamic_output_alloc.valid() || @@ -5090,10 +5118,17 @@ bool deepseek4_step_layer_range( return false; } - out_logits->resize((size_t)w.n_vocab); const size_t logits_offset = (size_t)(n_tokens - 1) * (size_t)w.n_vocab * sizeof(float); - ggml_backend_tensor_get(logits, out_logits->data(), logits_offset, - sizeof(float) * (size_t)w.n_vocab); + if (out_logits) { + out_logits->resize((size_t)w.n_vocab); + ggml_backend_tensor_get(logits, out_logits->data(), logits_offset, + sizeof(float) * (size_t)w.n_vocab); + } + if (out_argmax) { + const size_t argmax_offset = (size_t)(n_tokens - 1) * sizeof(*out_argmax); + ggml_backend_tensor_get(argmax_tokens, out_argmax, argmax_offset, + sizeof(*out_argmax)); + } ggml_free(ctx); } if (telemetry) telemetry->output_us += ds4_elapsed_us(output_t0, Ds4TimingClock::now()); @@ -5225,8 +5260,8 @@ bool create_deepseek4_cache(ggml_backend_t backend, } void free_deepseek4_cache(DeepSeek4Cache & c) { - delete c.layer_range_runtime; - c.layer_range_runtime = nullptr; + delete c.layer_range_cache; + c.layer_range_cache = nullptr; if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } c.layers.clear(); diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 9d76b86b6..9e1b61f92 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -257,7 +257,7 @@ struct DeepSeek4LayerCache { // Per-shard runtime state for deepseek4_step_layer_range (host-side HC weight // cache + cached decode graphs). Defined in deepseek4_graph.cpp; owned by the // DeepSeek4Cache below and released by free_deepseek4_cache(). -struct DeepSeek4LayerRangeRuntimeCache; +struct DeepSeek4LayerRangeCache; struct DeepSeek4Cache { int cur_pos = 0; @@ -270,7 +270,7 @@ struct DeepSeek4Cache { ggml_tensor * hc_state = nullptr; // [n_hc * n_embd] // Lazily created on the first deepseek4_step_layer_range call. - DeepSeek4LayerRangeRuntimeCache * layer_range_runtime = nullptr; + DeepSeek4LayerRangeCache * layer_range_cache = nullptr; ggml_context * ctx = nullptr; ggml_backend_buffer_t buf = nullptr; @@ -343,6 +343,7 @@ bool deepseek4_step_layer_range( int layer_begin, int layer_end, std::vector * out_logits, + int32_t * out_argmax = nullptr, const int32_t * token_ids = nullptr, DeepSeek4StepTelemetry * telemetry = nullptr, bool allow_decode_graph_reuse = true); diff --git a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp index b5126c85b..d034bb810 100644 --- a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp +++ b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp @@ -412,7 +412,10 @@ bool DeepSeek4LayerSplitAdapter::init_mixed_target_split() { } void DeepSeek4LayerSplitAdapter::begin_request(const GenerateRequest & req) { - (void)req; + sampler_ = req.sampler; + if (req.do_sample && sampler_.seed != 0) { + sampler_rng_.seed(sampler_.seed); + } } void DeepSeek4LayerSplitAdapter::reset_request_state() { @@ -464,8 +467,10 @@ bool DeepSeek4LayerSplitAdapter::run_forward( return run_mixed_forward(tokens, base_pos, last_tok, logits_out); } - // Local multi-GPU path: run each shard's layers sequentially - // Pass HC state between shards at boundaries + // Local multi-GPU path: run each shard's layers sequentially and pass HC + // state at the boundaries. The final shard computes argmax on the GPU so + // greedy decode returns next_tok without copying the full vocabulary. + last_tok = -1; for (size_t si = 0; si < shards_.size(); ++si) { auto & shard = shards_[si]; const bool is_last = (si == shards_.size() - 1); @@ -477,7 +482,8 @@ bool DeepSeek4LayerSplitAdapter::run_forward( shard.backend, shard.gpu, shard.weights, shard.cache, hc_state_, shard_input, n_tokens, base_pos, shard.layer_begin, shard.layer_end, - shard_logits, tokens.data(), timing ? &step_tel : nullptr)) { + shard_logits, is_last ? &last_tok : nullptr, + tokens.data(), timing ? &step_tel : nullptr)) { std::fprintf(stderr, "[deepseek4-split] forward failed on shard %zu\n", si); return false; } @@ -485,7 +491,10 @@ bool DeepSeek4LayerSplitAdapter::run_forward( } cur_pos_ = base_pos + n_tokens; - last_tok = tokens.back(); + if (last_tok < 0) { + std::fprintf(stderr, "[deepseek4-split] local forward produced no argmax token\n"); + return false; + } last_tok_ = last_tok; if (logits_out && !logits_out->empty()) { @@ -528,7 +537,8 @@ bool DeepSeek4LayerSplitAdapter::run_mixed_forward( local_shard.cache, hc_state_, embed.data(), n_tokens, base_pos, local_shard.layer_begin, local_shard.layer_end, - &hidden_out, tokens.data(), timing ? &local_tel : nullptr)) { + &hidden_out, /*out_argmax=*/nullptr, + tokens.data(), timing ? &local_tel : nullptr)) { std::fprintf(stderr, "[deepseek4-split] local shard forward failed\n"); return false; } diff --git a/server/src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp b/server/src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp index d0d3b6dc2..8eb7d0246 100644 --- a/server/src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp +++ b/server/src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp @@ -228,6 +228,7 @@ int run_deepseek4_target_shard_ipc_daemon( const bool timing = target_shard_timing_enabled(); std::vector logits; + int32_t final_argmax = -1; const float * shard_input = req.boundary_activation->data(); for (size_t i = 0; i < shards.size(); ++i) { auto & shard = shards[i]; @@ -241,7 +242,8 @@ int run_deepseek4_target_shard_ipc_daemon( shard.backend, shard.gpu, shard.weights, shard.cache, hc_state, shard_input, n_tokens, req.base_pos, shard.layer_begin, shard.layer_end, - is_last ? &logits : nullptr, + (is_last && req.want_logits) ? &logits : nullptr, + is_last ? &final_argmax : nullptr, req.token_ids ? req.token_ids->data() : nullptr, timing ? &tel : nullptr); if (!ok) { @@ -261,23 +263,19 @@ int run_deepseek4_target_shard_ipc_daemon( } const int vocab = shards.back().weights.n_vocab; - if (vocab <= 0 || logits.size() != (size_t)vocab) { + if (final_argmax < 0 || final_argmax >= vocab) { + std::fprintf(stderr, + "[deepseek4-target-shard] invalid argmax token: vocab=%d token=%d\n", + vocab, final_argmax); + return false; + } + if (req.want_logits && logits.size() != (size_t)vocab) { std::fprintf(stderr, "[deepseek4-target-shard] invalid logits payload: expected=%d got=%zu\n", vocab, logits.size()); return false; } - - const float * last_logits = logits.data(); - int32_t best = 0; - float best_val = last_logits[0]; - for (int i = 1; i < vocab; ++i) { - if (last_logits[i] > best_val) { - best_val = last_logits[i]; - best = i; - } - } - resp.last_tok = best; + resp.last_tok = final_argmax; if (req.want_logits) { resp.logits = std::move(logits); } diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index c7d8db417..4424a8e89 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -637,6 +637,34 @@ static void test_local_shards_receive_hc_boundary_state() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_layer_split_request_sampler() { + std::fprintf(stderr, " test_layer_split_request_sampler ..."); + + auto adapter = make_test_adapter(); + GenerateRequest req; + req.do_sample = true; + req.sampler.temp = 0.75f; + req.sampler.top_p = 0.9f; + req.sampler.top_k = 32; + req.sampler.rep_pen = 1.1f; + req.sampler.freq_pen = 0.2f; + req.sampler.pres_pen = 0.3f; + req.sampler.seed = 12345; + adapter.begin_request(req); + + TEST_ASSERT(adapter.sampler_.temp == req.sampler.temp); + TEST_ASSERT(adapter.sampler_.top_p == req.sampler.top_p); + TEST_ASSERT(adapter.sampler_.top_k == req.sampler.top_k); + TEST_ASSERT(adapter.sampler_.rep_pen == req.sampler.rep_pen); + TEST_ASSERT(adapter.sampler_.freq_pen == req.sampler.freq_pen); + TEST_ASSERT(adapter.sampler_.pres_pen == req.sampler.pres_pen); + TEST_ASSERT(adapter.sampler_.seed == req.sampler.seed); + std::mt19937_64 expected_rng(req.sampler.seed); + TEST_ASSERT(adapter.sampler_rng_() == expected_rng()); + + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_loader_rejects_missing_required_metadata(ggml_backend_t backend) { std::fprintf(stderr, " test_loader_rejects_missing_required_metadata ..."); @@ -1117,6 +1145,7 @@ static void test_output_graph_reuse_microbench(ggml_backend_t backend) { std::vector lm_head((size_t) n_embd * n_vocab); std::vector rebuild_logits((size_t) n_vocab * n_tokens); std::vector cached_logits((size_t) n_vocab * n_tokens); + std::vector cached_argmax((size_t) n_tokens, -1); std::mt19937 rng(7); std::uniform_real_distribution dist(-0.2f, 0.2f); @@ -1181,8 +1210,11 @@ static void test_output_graph_reuse_microbench(ggml_backend_t backend) { ggml_set_input(lm_head_t); ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, inp_t, 1.0e-6f), norm_w_t); ggml_tensor * logits = ggml_mul_mat(ctx, lm_head_t, norm); + ggml_tensor * argmax_tokens = ggml_argmax(ctx, logits); ggml_set_output(logits); + ggml_set_output(argmax_tokens); ggml_cgraph * gf = ggml_new_graph_custom(ctx, 128, false); + ggml_build_forward_expand(gf, argmax_tokens); ggml_build_forward_expand(gf, logits); ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type()); bool alloc_ok = ggml_gallocr_alloc_graph(alloc, gf); @@ -1205,6 +1237,8 @@ static void test_output_graph_reuse_microbench(ggml_backend_t backend) { TEST_ASSERT(ok); if (ok) { ggml_backend_tensor_get(logits, cached_logits.data(), 0, cached_logits.size() * sizeof(float)); + ggml_backend_tensor_get(argmax_tokens, cached_argmax.data(), 0, + cached_argmax.size() * sizeof(int32_t)); } cached_total_ms += elapsed_ms(t0, t1); } @@ -1216,6 +1250,17 @@ static void test_output_graph_reuse_microbench(ggml_backend_t backend) { TEST_ASSERT_MSG(std::isfinite(cached_logits[i]), "cached output logits must be finite"); TEST_ASSERT_MSG(std::isfinite(rebuild_logits[i]), "rebuilt output logits must be finite"); } + for (int token = 0; token < n_tokens; ++token) { + int32_t expected = 0; + const size_t offset = (size_t) token * n_vocab; + for (int vocab = 1; vocab < n_vocab; ++vocab) { + if (cached_logits[offset + (size_t) vocab] > + cached_logits[offset + (size_t) expected]) { + expected = vocab; + } + } + TEST_ASSERT(cached_argmax[(size_t) token] == expected); + } std::fprintf(stderr, " rebuild_avg=%.3fms cached_avg=%.3fms\n", rebuild_total_ms / iters, cached_total_ms / iters); @@ -1660,6 +1705,7 @@ int main() { test_layer_range_validation(); test_hc_state_dimensions(); test_local_shards_receive_hc_boundary_state(); + test_layer_split_request_sampler(); test_loader_rejects_missing_required_metadata(backend); test_loader_rejects_invalid_compress_ratio_type(backend); test_loader_rejects_zero_vocab_size(backend);