From 5dd8cf598b46bfddedd0309f1fdeb9700f4beb19 Mon Sep 17 00:00:00 2001 From: Howard Su Date: Sun, 12 Jul 2026 17:33:13 +0800 Subject: [PATCH 1/4] Restore DeepSeek4 HIP backend safely Restore the DeepSeek4 full-HIP backend while keeping the active request path stable and preserving batched prefill throughput. Key changes: - Re-enable cfg-driven batched prefill chunks for full-HIP and layer-split paths instead of forcing single-token prefill. - Keep the cached single-token attention graph disabled on HIP because the ggml_set_rows state-update path faults on longer prompts. - Snapshot existing raw SWA rows before writing the current batched chunk into the ring, avoiding stale/future KV reads while retaining batched execution. - Reset full-HIP request state between requests, preserve last logits in snapshots, and restore exact-prefix generation without extra prefill. - Propagate BudgetHook force-close behavior and structured GenerateResult errors through fresh and restored generation paths. Validation: - Local: server/build/test_deepseek4_unit and server/build/test_server_unit passed. - Remote HIP: build-hip/test_deepseek4_unit passed on howardsu@100.126.37.90. - Remote full-HIP perf with --chunk 512: 525 prompt tokens at 105.21 tok/s, 1045 prompt tokens at 109.44 tok/s. --- server/CMakeLists.txt | 49 +- server/src/common/backend_factory.cpp | 17 +- server/src/common/layer_split_backend.cpp | 6 +- server/src/common/layer_split_backend.h | 3 +- server/src/deepseek4/deepseek4_backend.cpp | 792 ++++++------------ server/src/deepseek4/deepseek4_backend.h | 85 +- server/src/deepseek4/deepseek4_daemon.cpp | 21 +- server/src/deepseek4/deepseek4_daemon.h | 10 +- server/src/deepseek4/deepseek4_graph.cpp | 384 +++++++-- server/src/deepseek4/deepseek4_hc_cuda.cu | 344 ++++++++ server/src/deepseek4/deepseek4_hc_cuda.h | 32 + server/src/deepseek4/deepseek4_internal.h | 11 +- .../deepseek4_layer_split_adapter.cpp | 32 +- .../deepseek4/deepseek4_layer_split_adapter.h | 3 +- .../src/gemma4/gemma4_layer_split_adapter.cpp | 4 +- .../src/gemma4/gemma4_layer_split_adapter.h | 3 +- .../src/laguna/laguna_layer_split_adapter.cpp | 4 +- .../src/laguna/laguna_layer_split_adapter.h | 3 +- .../src/qwen35/qwen35_layer_split_adapter.cpp | 4 +- .../src/qwen35/qwen35_layer_split_adapter.h | 3 +- server/test/test_server_unit.cpp | 30 +- server/tests/test_deepseek4_unit.cpp | 256 ++++++ 22 files changed, 1321 insertions(+), 775 deletions(-) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index dacd851bb..aa251a349 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -401,19 +401,6 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip") # rms_norm_hip.cu is needed by the HIP chunk-B graph path regardless of SM80_EQUIV. target_sources(dflash_common PRIVATE src/rms_norm_hip.cu) set_source_files_properties(src/rms_norm_hip.cu PROPERTIES LANGUAGE HIP) - # GPU draft top-K + log-prob kernel (DFlash). The shared body in - # geometric_draft_topk_cuda.cu compiles unchanged for HIP through the - # hip_compat shim, so it is added directly with LANGUAGE HIP - # (same shared-.cu pattern as deepseek4_hc_cuda.cu above — no separate HIP - # translation unit). This makes DFLASH_GPU_DRAFT_TOPK a real path on ROCm - # instead of a no-op, so AMD stops paying the per-step vocab x n_tokens D2H + - # CPU heap extract. The hip_compat include dir is already on dflash_common. - target_sources(dflash_common PRIVATE src/common/geometric_draft_topk_cuda.cu) - set_source_files_properties(src/common/geometric_draft_topk_cuda.cu - PROPERTIES LANGUAGE HIP) - # PUBLIC so test consumers (test_dflash / test_draft_topk_cuda) also take the - # GPU draft top-K path instead of the CPU fallback. - target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK=1) if(DFLASH27B_HIP_SM80_EQUIV) find_path(DFLASH27B_ROCWMMA_INCLUDE_DIR rocwmma/rocwmma.hpp HINTS "${_dflash_rocm_root}/include" /opt/rocm/include @@ -447,9 +434,8 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda") src/flashprefill.cpp src/common/geometric_draft_topk_cuda.cu) # PUBLIC so consumers (e.g. the test_dflash executable) also see the macro - # and take the GPU draft top-K path instead of the CPU fallback. Same macro - # name as the HIP branch above (backend-neutral). - target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK=1) + # and take the GPU draft top-K path instead of the CPU fallback. + target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK_CUDA=1) # GPU port of the sample_logits chain. Compiled in by default; the path is # then opted into at runtime via the DFLASH_GPU_SAMPLE env var. Turn the # whole thing off at configure time with -DDFLASH_GPU_SAMPLER=OFF. @@ -677,40 +663,13 @@ if(DFLASH27B_TESTS) target_link_libraries(test_rms_norm_hip PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) add_test(NAME rms_norm_hip COMMAND test_rms_norm_hip) endif() - # DS4 Hierarchical-Controller pre-mix kernels vs CPU reference. HIP only: - # deepseek4_hc_cuda.cu is compiled into dflash_common on this backend. Validates - # the HC mix matvec + NHC=4 Sinkhorn path against the device's actual wave width. - if(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_deepseek4_hc_cuda.cpp") - add_executable(test_deepseek4_hc_cuda test/test_deepseek4_hc_cuda.cpp) - set_source_files_properties(test/test_deepseek4_hc_cuda.cpp PROPERTIES LANGUAGE HIP) - set_target_properties(test_deepseek4_hc_cuda PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") - target_include_directories(test_deepseek4_hc_cuda PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR}/src - ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) - target_link_libraries(test_deepseek4_hc_cuda PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) - add_test(NAME deepseek4_hc_cuda COMMAND test_deepseek4_hc_cuda) - endif() - # GPU draft top-K kernel vs CPU reference (extract_draft_topk). Built on both - # backends: geometric_draft_topk_cuda.cu is compiled into dflash_common on the - # cuda backend directly and on hip via LANGUAGE HIP + the hip_compat shim. + # GPU draft top-K kernel vs CPU reference (extract_draft_topk). CUDA only: + # geometric_draft_topk_cuda.cu is compiled into dflash_common solely on the cuda backend. if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_topk_cuda.cpp") add_executable(test_draft_topk_cuda test/test_draft_topk_cuda.cpp) target_include_directories(test_draft_topk_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common CUDA::cudart) add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda) - elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_topk_cuda.cpp") - # HIP build of the same GPU-vs-CPU parity test. The test source uses CUDA - # spellings ( + cudaMalloc/cudaMemcpy/...); the hip_compat - # shim maps them onto HIP, exactly as the kernel TU does. Lets the draft - # top-K port be validated on AMD (gfx1100 / gfx1151 / gfx12xx). - add_executable(test_draft_topk_cuda test/test_draft_topk_cuda.cpp) - set_source_files_properties(test/test_draft_topk_cuda.cpp PROPERTIES LANGUAGE HIP) - set_target_properties(test_draft_topk_cuda PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") - target_include_directories(test_draft_topk_cuda PRIVATE - ${CMAKE_CURRENT_SOURCE_DIR}/src - ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) - target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) - add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda) endif() # GPU port of the sample_logits chain vs the CPU reference. CUDA only: # geometric_sampler_cuda.cu is compiled into dflash_common solely on the cuda backend. diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 1383c2166..ef66db129 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -220,29 +220,26 @@ std::unique_ptr create_backend(const BackendArgs & args) { ? compiled_placement_backend() : args.device.backend; - // HIP single-device launches cannot rely on the CUDA/Halo auto-split - // path; use the single-backend loader, which can fall back to hybrid - // expert placement when a full monolithic load does not fit. if (target_backend == PlacementBackend::Hip && !args.device.is_layer_split() && !args.remote_target_shard.enabled()) { DeepSeek4BackendConfig cfg; cfg.model_path = args.model_path; - cfg.device = args.device; - cfg.stream_fd = args.stream_fd; - cfg.max_ctx = args.device.max_ctx; - cfg.chunk = args.chunk; + cfg.device = args.device; + cfg.max_ctx = args.device.max_ctx; + cfg.chunk = args.chunk; auto backend = std::make_unique(cfg); if (!backend->init()) { - std::fprintf(stderr, "[backend_factory] DeepSeek4Backend init failed\n"); + std::fprintf(stderr, "[backend_factory] DeepSeek4Backend full-weight init failed\n"); return nullptr; } return backend; } - // CUDA builds keep the layer-split backend so they can auto-split - // across CUDA and remote HIP target shards. + // CUDA keeps the layer-split backend so prefill stays on the batched + // layer-range path. HIP single-device can still load all weights + // monolithically, without the removed hybrid expert fallback. DeepSeek4LayerSplitAdapterConfig cfg; cfg.target_path = args.model_path; cfg.device = args.device; diff --git a/server/src/common/layer_split_backend.cpp b/server/src/common/layer_split_backend.cpp index 94e7e13a7..229861349 100644 --- a/server/src/common/layer_split_backend.cpp +++ b/server/src/common/layer_split_backend.cpp @@ -85,7 +85,11 @@ GenerateResult LayerSplitBackend::run_from_state(const GenerateRequest & req, } std::vector chunk(req.prompt.begin() + consumed, req.prompt.begin() + consumed + n_tokens); - if (!adapter_->prefill(chunk, base_pos + consumed, last_tok)) { + const bool ends_prompt = consumed + n_tokens == prompt_len; + const bool ends_snap = req.snap_pos >= 0 && req.snap_slot >= 0 && + base_pos + consumed + n_tokens == req.snap_pos; + const bool need_logits = ends_prompt || ends_snap; + if (!adapter_->prefill(chunk, base_pos + consumed, last_tok, need_logits)) { result.fail(GenerateErrorCode::PrefillFailed); return result; } diff --git a/server/src/common/layer_split_backend.h b/server/src/common/layer_split_backend.h index 0af4097b7..c89a83f10 100644 --- a/server/src/common/layer_split_backend.h +++ b/server/src/common/layer_split_backend.h @@ -27,7 +27,8 @@ class LayerSplitAdapter { virtual void reset_request_state() = 0; virtual int prefill_chunk_tokens() const { return 0; } virtual bool prefill(const std::vector & prompt, - int base_pos, int & last_tok) = 0; + int base_pos, int & last_tok, + bool need_logits = true) = 0; virtual bool decode_ar(int last_tok, int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io) = 0; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index e75a49ea9..94f22e353 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1,19 +1,16 @@ -// DeepSeek4Backend implementation — AR-only decode, chunked prefill. +// DeepSeek4Backend implementation — full-weight HIP/CUDA backend. #include "deepseek4_backend.h" -#include "deepseek4_internal.h" -#include "common/sampler.h" -#include "ggml.h" #include "ggml-backend.h" #include "ggml-cuda.h" +#include "ggml-cpu.h" #include #include #include #include #include -#include namespace dflash::common { @@ -28,13 +25,25 @@ static uint64_t elapsed_us(Clock::time_point start, Clock::time_point end) { return (uint64_t)std::chrono::duration_cast(end - start).count(); } +static void reset_full_request_state(DeepSeek4Cache & cache, + std::vector & hc_state, + std::vector & last_logits) { + cache.cur_pos = 0; + for (auto & lc : cache.layers) { + lc.n_comp = 0; + lc.n_index_comp = 0; + } + std::fill(hc_state.begin(), hc_state.end(), 0.0f); + last_logits.clear(); +} + static bool env_flag_enabled(const char * name) { const char * value = std::getenv(name); return value && value[0] && std::strcmp(value, "0") != 0; } -static double gib(uint64_t bytes) { - return (double) bytes / 1024.0 / 1024.0 / 1024.0; +static double ms(uint64_t us) { + return (double)us / 1000.0; } static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTelemetry & src) { @@ -57,24 +66,10 @@ static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTeleme dst.route_read_us += src.route_read_us; dst.route_select_us += src.route_select_us; dst.ffn_eval_us += src.ffn_eval_us; - dst.ffn_hot_us += src.ffn_hot_us; - dst.ffn_cold_us += src.ffn_cold_us; - dst.ffn_combine_us += src.ffn_combine_us; - dst.ffn_partition_us += src.ffn_partition_us; - dst.ffn_hot_graph_builds += src.ffn_hot_graph_builds; - dst.ffn_hot_graph_hits += src.ffn_hot_graph_hits; - dst.ffn_cold_graph_builds += src.ffn_cold_graph_builds; - dst.ffn_cold_graph_hits += src.ffn_cold_graph_hits; dst.hc_post_ffn_us += src.hc_post_ffn_us; dst.output_us += src.output_us; dst.sample_us += src.sample_us; dst.emit_us += src.emit_us; - dst.hot_selected += src.hot_selected; - dst.cold_selected += src.cold_selected; -} - -static double ms(uint64_t us) { - return (double)us / 1000.0; } static void log_step_tel(const char * phase, @@ -88,216 +83,16 @@ static void log_step_tel(const char * phase, "step=%.1fms embed=%.1fms attn_build=%.1fms attn_compute=%.1fms attn_read=%.1fms " "ffn_build=%.1fms ffn_compute=%.1fms ffn_read=%.1fms " "route_build=%.1fms route_compute=%.1fms route_read=%.1fms route_select=%.1fms " - "ffn=%.1fms hot=%.1fms cold=%.1fms combine=%.1fms partition=%.1fms " - "ffn_hot_graph_build=%llu ffn_hot_graph_hit=%llu ffn_cold_graph_build=%llu ffn_cold_graph_hit=%llu " - "hc_pre=%.1fms hc_pre_build=%.1fms hc_pre_input=%.1fms hc_pre_compute=%.1fms " - "hc_post=%.1fms output=%.1fms sample=%.1fms emit=%.1fms " - "hot_sel=%d cold_sel=%d\n", + "ffn=%.1fms hc_pre=%.1fms hc_pre_build=%.1fms hc_pre_input=%.1fms hc_pre_compute=%.1fms " + "hc_post=%.1fms output=%.1fms sample=%.1fms emit=%.1fms\n", phase, tokens, steps, wall_s, tok_s, ms(t.total_us), ms(t.embed_us), ms(t.attn_build_us), ms(t.attn_compute_us), ms(t.attn_read_us), ms(t.ffn_build_us), ms(t.ffn_compute_us), ms(t.ffn_read_us), ms(t.route_build_us), ms(t.route_compute_us), ms(t.route_read_us), ms(t.route_select_us), - ms(t.ffn_eval_us), ms(t.ffn_hot_us), ms(t.ffn_cold_us), ms(t.ffn_combine_us), - ms(t.ffn_partition_us), - (unsigned long long)t.ffn_hot_graph_builds, (unsigned long long)t.ffn_hot_graph_hits, - (unsigned long long)t.ffn_cold_graph_builds, (unsigned long long)t.ffn_cold_graph_hits, - ms(t.hc_pre_attn_us + t.hc_pre_ffn_us), - ms(t.hc_pre_build_us), - ms(t.hc_pre_input_us), - ms(t.hc_pre_compute_us), + ms(t.ffn_eval_us), ms(t.hc_pre_attn_us + t.hc_pre_ffn_us), + ms(t.hc_pre_build_us), ms(t.hc_pre_input_us), ms(t.hc_pre_compute_us), ms(t.hc_post_attn_us + t.hc_post_ffn_us), - ms(t.output_us), ms(t.sample_us), ms(t.emit_us), - t.hot_selected, t.cold_selected); -} - -static uint64_t layer_expert_bytes(const DeepSeek4Layer & layer, int n_expert) { - if (n_expert <= 0) return 0; - uint64_t bytes = 0; - if (layer.ffn_gate_exps) bytes += ggml_nbytes(layer.ffn_gate_exps) / (uint64_t) n_expert; - if (layer.ffn_up_exps) bytes += ggml_nbytes(layer.ffn_up_exps) / (uint64_t) n_expert; - if (layer.ffn_down_exps) bytes += ggml_nbytes(layer.ffn_down_exps) / (uint64_t) n_expert; - return bytes; -} - -struct Ds4ExpertMemoryInfo { - std::vector layer_expert_bytes; - uint64_t total_expert_bytes = 0; - uint64_t bytes_per_uniform_round = 0; - uint64_t hot_bytes = 0; - uint64_t cold_bytes = 0; - int total_hot = 0; - int total_cold = 0; -}; - -struct Ds4HybridBudgetInfo { - Ds4ExpertMemoryInfo mem; - size_t gpu_free = 0; - size_t gpu_total = 0; - uint64_t core_bytes = 0; - uint64_t kv_bytes = 0; - uint64_t warm_bytes = 256ULL * 1024 * 1024; - uint64_t safety_bytes = 512ULL * 1024 * 1024; - uint64_t expert_budget = 0; - int max_hot_per_layer = 0; -}; - -static bool compute_ds4_expert_memory_info(const DeepSeek4Weights & w, - const MoeHybridPlacement * placement, - Ds4ExpertMemoryInfo & out, - std::string * err) { - out = {}; - out.layer_expert_bytes.assign((size_t) w.n_layer, 0); - for (int il = 0; il < w.n_layer; ++il) { - const uint64_t bytes = layer_expert_bytes(w.layers[(size_t) il], w.n_expert); - out.layer_expert_bytes[(size_t) il] = bytes; - out.total_expert_bytes += bytes * (uint64_t) w.n_expert; - out.bytes_per_uniform_round += bytes; - } - if (out.bytes_per_uniform_round == 0) { - if (err) *err = "expert tensor metadata missing after partial load"; - return false; - } - if (!placement) return true; - if (!placement->matches(w.n_layer, w.n_expert, w.n_expert_used)) { - if (err) *err = "placement does not match DS4 dimensions"; - return false; - } - out.total_hot = placement->total_hot; - out.total_cold = w.n_layer * w.n_expert - placement->total_hot; - for (int il = 0; il < w.n_layer; ++il) { - const uint64_t layer_bytes = out.layer_expert_bytes[(size_t) il]; - const uint64_t hot_count = (uint64_t) placement->hot_counts[(size_t) il]; - out.hot_bytes += layer_bytes * hot_count; - out.cold_bytes += layer_bytes * ((uint64_t) w.n_expert - hot_count); - } - return true; -} - -static void log_ds4_expert_memory_info(const char * tag, - const Ds4ExpertMemoryInfo & info, - int n_layer) { - (void) n_layer; - std::fprintf(stderr, - "[deepseek4] %s expert_memory: total=%.2f GiB uniform_round=%.2f MiB hot=%d %.2f GiB cold=%d %.2f GiB\n", - tag, - gib(info.total_expert_bytes), - (double) info.bytes_per_uniform_round / 1024.0 / 1024.0, - info.total_hot, - gib(info.hot_bytes), - info.total_cold, - gib(info.cold_bytes)); -} - -static uint64_t estimate_ds4_cache_bytes(const DeepSeek4Weights & w, int max_ctx) { - size_t total_bytes = 0; - const size_t head_dim = (size_t) w.head_dim; - const size_t swa_size = (size_t) w.n_swa; - - for (int il = 0; il < w.n_layer; ++il) { - total_bytes += swa_size * head_dim * sizeof(uint16_t); - const uint32_t ratio = w.compress_ratios[(size_t) il]; - if (ratio == 0) continue; - - const size_t comp_cap = (size_t) (max_ctx / (int) ratio) + 16; - total_bytes += comp_cap * head_dim * sizeof(uint16_t); - - const size_t window = (ratio == 4) ? 8 : ratio; - total_bytes += window * head_dim * sizeof(float) * 2; - - if (ratio == 4) { - const size_t index_comp_width = (size_t) w.n_indexer_head * (size_t) w.n_indexer_head_dim; - total_bytes += comp_cap * index_comp_width * sizeof(uint16_t); - total_bytes += window * index_comp_width * sizeof(float) * 2; - } - } - - total_bytes += (size_t) w.n_hc * (size_t) w.n_embd * sizeof(float); - return total_bytes; -} - -static void fill_prefix_hot_placement(const DeepSeek4Weights & w, - int hot_per_layer, - MoeHybridPlacement & out) { - out = {}; - out.n_layer = w.n_layer; - out.n_expert = w.n_expert; - out.n_expert_used = w.n_expert_used; - out.hot_counts.assign((size_t) w.n_layer, hot_per_layer); - out.hot_expert_ids.resize((size_t) w.n_layer); - out.total_hot = hot_per_layer * w.n_layer; - for (int il = 0; il < w.n_layer; ++il) { - auto & ids = out.hot_expert_ids[(size_t) il]; - ids.reserve((size_t) hot_per_layer); - for (int ie = 0; ie < hot_per_layer; ++ie) { - ids.push_back((int32_t) ie); - } - } -} - -static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, - int gpu, - int max_ctx, - Ds4HybridBudgetInfo & out, - std::string * err) { - out = {}; - ggml_backend_cuda_get_device_memory(gpu, &out.gpu_free, &out.gpu_total); - if (out.gpu_total == 0) { - if (err) *err = "could not query GPU memory"; - return false; - } - - if (!compute_ds4_expert_memory_info(w, nullptr, out.mem, err)) { - return false; - } - - out.core_bytes = out.gpu_total - out.gpu_free; - out.kv_bytes = estimate_ds4_cache_bytes(w, max_ctx); - - if (out.gpu_total > out.core_bytes + out.kv_bytes + out.warm_bytes + out.safety_bytes) { - out.expert_budget = out.gpu_total - out.core_bytes - out.kv_bytes - out.warm_bytes - out.safety_bytes; - } - if (out.expert_budget > out.mem.total_expert_bytes) { - out.expert_budget = out.mem.total_expert_bytes; - } - if (const char * cap_env = std::getenv("DFLASH_EXPERT_BUDGET_MB")) { - const uint64_t cap_bytes = (uint64_t) std::max(0, std::atoi(cap_env)) * 1024ULL * 1024ULL; - if (cap_bytes > 0 && cap_bytes < out.expert_budget) { - out.expert_budget = cap_bytes; - } - } - if (out.expert_budget == 0) { - if (err) *err = "no VRAM budget available for DS4 experts"; - return false; - } - - out.max_hot_per_layer = std::min(w.n_expert, (int) (out.expert_budget / out.mem.bytes_per_uniform_round)); - if (out.max_hot_per_layer <= 0) { - if (err) *err = "expert budget is smaller than one uniform expert round"; - return false; - } - return true; -} - -static MoeHybridConfig make_ds4_parent_worker_cfg(const DeepSeek4Weights & w) { - MoeHybridConfig cfg; - cfg.n_embd = w.n_embd; - cfg.n_expert = w.n_expert; - cfg.n_expert_used = w.n_expert_used; - cfg.n_ff_exp = w.n_ff_exp; - cfg.n_ff_shexp = w.n_ff_exp; - cfg.n_layer = w.n_layer; - cfg.first_moe_layer = 0; - cfg.swiglu_clamp = w.swiglu_clamp_exp; - cfg.materialize_cold_experts = false; - return cfg; -} - -static MoeHybridConfig make_ds4_parent_cpu_tail_cfg(const DeepSeek4Weights & w) { - MoeHybridConfig cfg = make_ds4_parent_worker_cfg(w); - cfg.materialize_hot_experts = false; - cfg.materialize_cold_experts = true; - cfg.cold_expert_backend = MoeHybridColdBackend::Cpu; - return cfg; + ms(t.output_us), ms(t.sample_us), ms(t.emit_us)); } } // namespace @@ -309,233 +104,68 @@ DeepSeek4Backend::~DeepSeek4Backend() { shutdown(); } -bool DeepSeek4Backend::init() { - backend_ = ggml_backend_cuda_init(cfg_.device.gpu); - if (!backend_) { - std::fprintf(stderr, "[deepseek4] failed to create CUDA backend (gpu=%d)\n", - cfg_.device.gpu); +bool DeepSeek4Backend::load_full_model() { + if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { + std::fprintf(stderr, "[deepseek4] failed to load full model weights\n"); return false; } - - snap_backend_ = ggml_backend_init_by_name("cpu", nullptr); - - const PlacementBackend target_backend = - cfg_.device.backend == PlacementBackend::Auto - ? compiled_placement_backend() - : cfg_.device.backend; - - // HIP single-device launches should avoid the monolithic full-model load: - // a managed ~80 GiB allocation can stall or be killed on integrated UMA - // systems before we ever reach the existing OOM fallback path. - if (target_backend == PlacementBackend::Hip) { - std::fprintf(stderr, - "[deepseek4] HIP target detected; using hybrid expert load path\n"); - if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] hybrid mode failed: %s\n", cfg_.model_path); - return false; - } - } else if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { - std::fprintf(stderr, "[deepseek4] full model load failed, trying hybrid mode...\n"); - if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] hybrid mode also failed: %s\n", cfg_.model_path); - return false; - } - } - const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { - std::fprintf(stderr, "[deepseek4] failed to allocate KV cache (ctx=%d)\n", max_ctx); - return false; - } - - if (moe_hybrid_) { - // Expert IPC removed — layer split replaces expert split. - // The DeepSeek4Backend single-GPU path now runs all experts locally. - } - - std::fprintf(stderr, "[deepseek4] initialized: %d layers, ctx=%d, %d experts (%d used)%s\n", - w_.n_layer, max_ctx, w_.n_expert, w_.n_expert_used, - moe_hybrid_ ? " [hybrid]" : ""); - return true; -} - -bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & w, - int max_ctx, - MoeHybridPlacement & out, - std::string * err) const { - Ds4HybridBudgetInfo budget; - if (!compute_ds4_hybrid_budget_info(w, cfg_.device.gpu, max_ctx, budget, err)) { - return false; - } - - const int hot_per_layer = budget.max_hot_per_layer; - fill_prefix_hot_placement(w, hot_per_layer, out); - - Ds4ExpertMemoryInfo placed_mem; - if (!compute_ds4_expert_memory_info(w, &out, placed_mem, err)) { + std::fprintf(stderr, "[deepseek4] failed to create KV cache (ctx=%d)\n", max_ctx); + free_deepseek4_weights(w_); return false; } - - std::fprintf(stderr, - "[deepseek4] hybrid placement: gpu_total=%.2f GiB gpu_free=%.2f GiB core=%.2f GiB kv=%.2f GiB warm=%.2f GiB safety=%.2f GiB expert_budget=%.2f GiB hot/layer=%d\n", - gib((uint64_t) budget.gpu_total), - gib((uint64_t) budget.gpu_free), - gib(budget.core_bytes), - gib(budget.kv_bytes), - gib(budget.warm_bytes), - gib(budget.safety_bytes), - gib(budget.expert_budget), - hot_per_layer); - log_ds4_expert_memory_info("placement", placed_mem, w.n_layer); + hc_state_.clear(); return true; } -bool DeepSeek4Backend::init_hybrid_model() { - TargetLoadPlan plan; - plan.skip_expert_tensors = true; - if (!load_deepseek4_gguf_partial(cfg_.model_path, backend_, plan, w_)) { - std::fprintf(stderr, "[deepseek4] failed to partially load model for hybrid mode: %s\n", - cfg_.model_path); +bool DeepSeek4Backend::init() { + backend_ = ggml_backend_cuda_init(cfg_.device.gpu); + if (!backend_) { + std::fprintf(stderr, "[deepseek4] failed to init backend gpu=%d\n", cfg_.device.gpu); return false; } - - std::string err; - const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; - if (!compute_uniform_hybrid_placement(w_, max_ctx, moe_placement_, &err)) { - std::fprintf(stderr, "[deepseek4] failed to compute hybrid placement: %s\n", err.c_str()); + snap_backend_ = ggml_backend_cpu_init(); + if (!snap_backend_) { + std::fprintf(stderr, "[deepseek4] failed to init snapshot CPU backend\n"); return false; } - - if (moe_placement_.total_hot >= w_.n_layer * w_.n_expert) { - free_deepseek4_weights(w_); - if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { - std::fprintf(stderr, "[deepseek4] failed to reload full model after placement: %s\n", - cfg_.model_path); - return false; - } - return true; - } - - auto hybrid = std::make_shared(); - const MoeHybridConfig hybrid_cfg = make_ds4_parent_worker_cfg(w_); - if (!build_deepseek4_moe_hybrid_storage_from_file_with_mmap( - cfg_.model_path, backend_, w_, moe_placement_, &hybrid_cfg, *hybrid, &err)) { - std::fprintf(stderr, "[deepseek4] failed to build hybrid expert storage: %s\n", err.c_str()); + if (!load_full_model()) { return false; } - - if (hybrid->has_mmap() && !hybrid->materialized_cold_experts) { - size_t max_expert_bytes = 0; - for (const auto & layer : hybrid->layers) { - const size_t per_expert_bytes = layer.fused_gate_up - ? layer.gate_up_expert_bytes + layer.down_expert_bytes - : layer.gate_expert_bytes + layer.up_expert_bytes + layer.down_expert_bytes; - max_expert_bytes = std::max(max_expert_bytes, per_expert_bytes); - } - if (max_expert_bytes == 0) { - std::fprintf(stderr, "[deepseek4] failed to compute streaming expert size\n"); - return false; - } - if (!stream_engine_.init(backend_, max_expert_bytes, &err)) { - std::fprintf(stderr, "[deepseek4] failed to init cold-expert stream engine: %s\n", - err.c_str()); - return false; - } - std::fprintf(stderr, - "[deepseek4] cold-expert stream engine ready: pinned=%.1f MiB scratch=%.1f MiB\n", - stream_engine_.pinned_bytes() / 1024.0 / 1024.0, - stream_engine_.scratch_bytes() / 1024.0 / 1024.0); - } - - moe_hybrid_ = std::move(hybrid); - w_.moe_hybrid = true; - const int total_cold = w_.n_layer * w_.n_expert - moe_placement_.total_hot; - const char * cold_backend = - moe_hybrid_->cold_backend_kind == MoeHybridColdBackend::Gpu ? "gpu" : "cpu"; - std::fprintf(stderr, "[deepseek4] hybrid experts ready: hot=%d cold=%d cold_backend=%s%s\n", - moe_placement_.total_hot, total_cold, cold_backend, ""); return true; } void DeepSeek4Backend::print_ready_banner() const { - std::printf("[deepseek4-daemon] ready layers=%d ctx=%d experts=%d/%d\n", - w_.n_layer, cache_.max_ctx, w_.n_expert_used, w_.n_expert); + std::printf("[deepseek4] ready (layers=%d, embd=%d, vocab=%d, full-weight backend)\n", + w_.n_layer, w_.n_embd, w_.n_vocab); std::fflush(stdout); } bool DeepSeek4Backend::park(const std::string & what) { - const bool want_target = (what.empty() || what == "all" || what == "target"); - if (!want_target || parked_) return true; - - maybe_save_routing_stats(); - for (int i = 0; i < PREFIX_SLOTS; ++i) { - free_deepseek4_snapshot(snapshots_[i]); - } - last_logits_.clear(); + if (what != "target") return false; free_deepseek4_cache(cache_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; free_deepseek4_weights(w_); + hc_state_.clear(); + last_logits_.clear(); parked_ = true; - std::printf("[deepseek4] parked (VRAM released)\n"); - std::fflush(stdout); return true; } bool DeepSeek4Backend::unpark(const std::string & what) { - const bool want_target = (what.empty() || what == "all" || what == "target"); - if (!want_target || !parked_) return true; - - const PlacementBackend target_backend = - cfg_.device.backend == PlacementBackend::Auto - ? compiled_placement_backend() - : cfg_.device.backend; - - if (target_backend == PlacementBackend::Hip) { - if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] unpark: failed to restore hybrid mode\n"); - free_deepseek4_weights(w_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; - return false; - } - } else if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { - std::fprintf(stderr, "[deepseek4] unpark: full model reload failed, trying hybrid mode...\n"); - if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); - free_deepseek4_weights(w_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; - return false; - } - } - - const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; - if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { - std::fprintf(stderr, "[deepseek4] unpark: failed to recreate KV cache (ctx=%d)\n", max_ctx); - free_deepseek4_cache(cache_); - free_deepseek4_weights(w_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; - return false; - } - + if (what != "target") return false; + if (!parked_) return true; + if (!load_full_model()) return false; parked_ = false; - std::printf("[deepseek4] unparked (VRAM restored)\n"); - std::fflush(stdout); return true; } int DeepSeek4Backend::do_prefill(const std::vector & tokens, - const DaemonIO & io, - int kv_offset) { - // Hybrid currently implements HC for single-token steps only; keep prefill - // token-by-token so the first sampled token is seeded from the correct HC state. - const int chunk = moe_hybrid_ ? 1 : (cfg_.chunk > 0 ? cfg_.chunk : 512); + const DaemonIO & io, + int kv_offset, + int snap_pos, + int snap_slot) { + const int chunk = cfg_.chunk > 0 ? cfg_.chunk : 512; const int n_total = (int)tokens.size(); int pos = kv_offset; last_logits_.clear(); @@ -544,34 +174,43 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, DeepSeek4StepTelemetry tel_acc; int steps = 0; - for (int i = 0; i < n_total; i += chunk) { + for (int i = 0; i < n_total; ) { if (io.cancelled) return pos; + int n_tok = std::min(chunk, n_total - i); + if (snap_slot >= 0 && snap_pos > pos && snap_pos < pos + n_tok) { + n_tok = snap_pos - pos; + } - const int n_tok = std::min(chunk, n_total - i); - - // Embed tokens - std::vector embed(w_.n_embd * n_tok); + std::vector embed((size_t)w_.n_embd * (size_t)n_tok); const auto embed_t0 = Clock::now(); - w_.embedder.embed(tokens.data() + i, n_tok, embed.data()); + if (!w_.embedder.embed(tokens.data() + i, n_tok, embed.data())) { + std::fprintf(stderr, "[deepseek4] embedding failed at pos=%d\n", pos); + return -1; + } DeepSeek4StepTelemetry step_tel; if (timing) step_tel.embed_us = elapsed_us(embed_t0, Clock::now()); - // Run forward pass std::vector logits; - if (!deepseek4_step(backend_, w_, cache_, embed.data(), n_tok, pos, logits, - moe_hybrid_.get(), tokens.data() + i, - moe_hybrid_ ? &stream_engine_ : nullptr, - timing ? &step_tel : nullptr, - routing_stats_.get())) { - std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d\n", pos); + if (!deepseek4_step_layer_range( + backend_, w_, cache_, hc_state_, + embed.data(), n_tok, pos, 0, w_.n_layer, + &logits, tokens.data() + i, timing ? &step_tel : nullptr)) { + std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d n_tokens=%d\n", pos, n_tok); return -1; } + last_logits_ = std::move(logits); + pos += n_tok; + if (snap_slot >= 0 && snap_pos >= 0 && pos == snap_pos) { + if (snapshot_save(snap_slot)) { + std::printf("[snap] inline slot=%d cur_pos=%d\n", snap_slot, snap_pos); + std::fflush(stdout); + } + } if (timing) { add_step_tel(tel_acc, step_tel); steps++; } - last_logits_ = std::move(logits); - pos += n_tok; + i += n_tok; } if (timing) { log_step_tel("prefill", n_total, steps, elapsed_s(phase_t0), tel_acc); @@ -579,134 +218,142 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, return pos; } -bool DeepSeek4Backend::do_decode(int committed, int n_gen, - std::vector & out_tokens, - const DaemonIO & io, - const BudgetHook & budget_hook, - bool * forced_close_out) { +bool DeepSeek4Backend::do_decode(int committed, + int n_gen, + std::vector & out_tokens, + const DaemonIO & io, + const BudgetHook & budget_hook, + bool * forced_close_out) { if (forced_close_out) *forced_close_out = false; const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); const auto phase_t0 = Clock::now(); DeepSeek4StepTelemetry tel_acc; int steps = 0; - - for (int generated = 0; generated < n_gen; generated++) { - if (io.cancelled) break; - - // Budget hook: force-close if remaining budget hits threshold - if (!budget_hook.close_token_ids.empty() && - (n_gen - generated) <= budget_hook.hard_limit_remaining) { - // Inject close-tag tokens - for (int32_t close_tok : budget_hook.close_token_ids) { - out_tokens.push_back(close_tok); - io.emit(close_tok); - if (io.cancelled) break; + int last_tok = -1; + bool budget_close_started = false; + int close_inject_pos = 0; + const int committed_at_entry = committed; + auto maybe_force_close = [&](int32_t & tok, int committed_now) { + if (budget_hook.close_token_ids.empty()) return; + if (budget_close_started && + close_inject_pos < (int)budget_hook.close_token_ids.size()) { + tok = budget_hook.close_token_ids[(size_t)close_inject_pos++]; + return; + } + if (budget_close_started) return; + const int generated = committed_now - committed_at_entry; + const int remaining = n_gen - generated; + if (remaining <= budget_hook.hard_limit_remaining) { + const int32_t first_close = budget_hook.close_token_ids.front(); + if (tok == first_close) { + budget_close_started = true; + close_inject_pos = 1; + return; } + tok = first_close; + budget_close_started = true; + close_inject_pos = 1; if (forced_close_out) *forced_close_out = true; - break; } - - // Get last logits and sample - std::vector logits; - if (generated == 0 && !last_logits_.empty()) { - logits = last_logits_; - } else { - std::vector embed(w_.n_embd); - int32_t tok_to_eval = out_tokens.empty() ? 0 : out_tokens.back(); - const auto embed_t0 = Clock::now(); - w_.embedder.embed(&tok_to_eval, 1, embed.data()); - DeepSeek4StepTelemetry step_tel; - if (timing) step_tel.embed_us = elapsed_us(embed_t0, Clock::now()); - - const int pos = std::max(0, committed + generated - 1); - if (!deepseek4_step(backend_, w_, cache_, embed.data(), 1, - pos, logits, - moe_hybrid_.get(), &tok_to_eval, - moe_hybrid_ ? &stream_engine_ : nullptr, - timing ? &step_tel : nullptr, - routing_stats_.get())) { - std::fprintf(stderr, "[deepseek4] decode step failed\n"); - return false; - } - if (timing) { - add_step_tel(tel_acc, step_tel); - steps++; + }; + + for (int step = 0; step < n_gen; ++step) { + if (io.cancelled) return true; + + int32_t tok_to_eval; + if (step == 0) { + if (last_logits_.empty()) return false; + if (sampler_.needs_logit_processing()) { + tok_to_eval = sample_logits(last_logits_.data(), w_.n_vocab, sampler_, out_tokens, sampler_rng_); + } else { + tok_to_eval = (int)std::distance(last_logits_.begin(), + std::max_element(last_logits_.begin(), last_logits_.end())); } + } else { + tok_to_eval = last_tok; } + maybe_force_close(tok_to_eval, committed + step); - // Sample (argmax for now) - int32_t next_token = 0; - { - const auto sample_t0 = Clock::now(); - float max_val = logits[0]; - for (int i = 1; i < w_.n_vocab; i++) { - if (logits[i] > max_val) { - max_val = logits[i]; - next_token = i; - } - } - if (timing) tel_acc.sample_us += elapsed_us(sample_t0, Clock::now()); + out_tokens.push_back(tok_to_eval); + io.emit(tok_to_eval); + if (deepseek4_is_eos_tok(tok_to_eval, w_)) { + io.emit(-1); + return true; } - out_tokens.push_back(next_token); - const auto emit_t0 = Clock::now(); - io.emit(next_token); - if (timing) tel_acc.emit_us += elapsed_us(emit_t0, Clock::now()); - if (deepseek4_is_eos_tok(next_token, w_)) { - break; + std::vector embed((size_t)w_.n_embd); + if (!w_.embedder.embed(&tok_to_eval, 1, embed.data())) return false; + DeepSeek4StepTelemetry step_tel; + std::vector logits; + if (!deepseek4_step_layer_range( + backend_, w_, cache_, hc_state_, + embed.data(), 1, committed + step, 0, w_.n_layer, + &logits, &tok_to_eval, timing ? &step_tel : nullptr)) { + return false; + } + last_logits_ = std::move(logits); + if (timing) { + add_step_tel(tel_acc, step_tel); + steps++; + } + + if (sampler_.needs_logit_processing()) { + last_tok = sample_logits(last_logits_.data(), w_.n_vocab, sampler_, out_tokens, sampler_rng_); + } else { + last_tok = (int)std::distance(last_logits_.begin(), + std::max_element(last_logits_.begin(), last_logits_.end())); } } + if (timing) { - log_step_tel("decode", (int)out_tokens.size(), steps, elapsed_s(phase_t0), tel_acc); + log_step_tel("decode", n_gen, steps, elapsed_s(phase_t0), tel_acc); } + io.emit(-1); return true; } GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, - const DaemonIO & io) { + const DaemonIO & io) { GenerateResult result; - auto t0 = Clock::now(); + sampler_ = req.sampler; + if (sampler_.seed != 0) sampler_rng_.seed(sampler_.seed); - // Prefill - int committed = do_prefill(req.prompt, io); - if (committed < 0) { - result.fail(GenerateErrorCode::PrefillFailed); + if (parked_ && !unpark("target")) { + result.fail(GenerateErrorCode::ModelParked, "target_parked"); return result; } - result.prefill_s = elapsed_s(t0); - - if (req.n_gen <= 0) { - result.succeed(); - maybe_save_routing_stats(); + if ((int)req.prompt.size() + req.n_gen + 1 > cache_.max_ctx) { + result.fail(GenerateErrorCode::ContextOverflow); return result; } + reset_full_request_state(cache_, hc_state_, last_logits_); - // Decode - auto t1 = Clock::now(); - std::vector gen_tokens; - gen_tokens.reserve(req.n_gen); + const auto t0 = Clock::now(); + const int committed = do_prefill(req.prompt, io, 0, req.snap_pos, req.snap_slot); + if (committed < 0) { + result.fail(GenerateErrorCode::PrefillFailed); + return result; + } + result.prefill_s = elapsed_s(t0); bool forced_close = false; - if (!do_decode(committed, req.n_gen, gen_tokens, io, - req.budget_hook, &forced_close)) { + if (!do_decode(committed, req.n_gen, result.tokens, io, req.budget_hook, &forced_close)) { result.fail(GenerateErrorCode::DecodeFailed); return result; } - + result.decode_s = elapsed_s(t0) - result.prefill_s; result.succeed(); - result.tokens = std::move(gen_tokens); - result.decode_s = elapsed_s(t1); - result.budget_forced_close = forced_close; - maybe_save_routing_stats(); + if (forced_close) result.budget_forced_close = true; return result; } -// ── Snapshots ─────────────────────────────────────────────────────────── - bool DeepSeek4Backend::snapshot_save(int slot) { if (slot < 0 || slot >= PREFIX_SLOTS) return false; - // TODO: Implement snapshot save (copy KV cache + HC state to CPU) - return false; + if (!deepseek4_snapshot_save(cache_, snap_backend_, snapshots_[slot])) { + return false; + } + snapshots_[slot].last_logits = last_logits_; + return true; } void DeepSeek4Backend::snapshot_free(int slot) { @@ -715,56 +362,87 @@ void DeepSeek4Backend::snapshot_free(int slot) { } bool DeepSeek4Backend::snapshot_used(int slot) const { - if (slot < 0 || slot >= PREFIX_SLOTS) return false; - return snapshots_[slot].ctx != nullptr; + return slot >= 0 && slot < PREFIX_SLOTS && snapshots_[slot].ctx != nullptr; } int DeepSeek4Backend::snapshot_cur_pos(int slot) const { - if (slot < 0 || slot >= PREFIX_SLOTS) return 0; - return snapshots_[slot].cur_pos; + return snapshot_used(slot) ? snapshots_[slot].cur_pos : 0; } -GenerateResult DeepSeek4Backend::restore_and_generate_impl( - int slot, const GenerateRequest & req, const DaemonIO & io) { - // TODO: Implement snapshot restore + generate - (void)slot; - return generate_impl(req, io); +GenerateResult DeepSeek4Backend::restore_and_generate_impl(int slot, + const GenerateRequest & req, + const DaemonIO & io) { + GenerateResult result; + if (parked_ && !unpark("target")) { + result.fail(GenerateErrorCode::ModelParked, "target_parked"); + return result; + } + if (!snapshot_used(slot) || !deepseek4_snapshot_restore(snapshots_[slot], cache_)) { + result.fail(GenerateErrorCode::InvalidSnapshotSlot, "restore"); + return result; + } + last_logits_ = snapshots_[slot].last_logits; + + sampler_ = req.sampler; + if (sampler_.seed != 0) sampler_rng_.seed(sampler_.seed); + const int snap_pos = cache_.cur_pos; + if (snap_pos < 0) { + result.fail(GenerateErrorCode::InvalidSnapshotSlot, "restore_pos"); + return result; + } + if (snap_pos > (int)req.prompt.size()) { + std::fprintf(stderr, + "[pc] deepseek4 snapshot longer than prompt (snap=%d > prompt=%zu) — " + "fresh prefill fallback\n", snap_pos, req.prompt.size()); + reset_full_request_state(cache_, hc_state_, last_logits_); + return generate_impl(req, io); + } + if ((int)req.prompt.size() + req.n_gen + 1 > cache_.max_ctx) { + result.fail(GenerateErrorCode::ContextOverflow); + return result; + } + std::vector remaining(req.prompt.begin() + snap_pos, req.prompt.end()); + const auto t0 = Clock::now(); + const int committed = remaining.empty() ? snap_pos : do_prefill(remaining, io, snap_pos); + if (committed < 0) { + result.fail(GenerateErrorCode::PrefillFailed); + return result; + } + result.prefill_s = elapsed_s(t0); + bool forced_close = false; + if (!do_decode(committed, req.n_gen, result.tokens, io, req.budget_hook, &forced_close)) { + result.fail(GenerateErrorCode::DecodeFailed); + return result; + } + result.decode_s = elapsed_s(t0) - result.prefill_s; + result.succeed(); + if (forced_close) result.budget_forced_close = true; + return result; } bool DeepSeek4Backend::handle_compress(const std::string & line, - const DaemonIO & io) { - (void)line; (void)io; - std::fprintf(stderr, "[deepseek4] compress not yet supported\n"); + const DaemonIO & io) { + (void)line; + (void)io; return false; } -void DeepSeek4Backend::free_drafter() { - // No drafter in AR-only mode -} - -void DeepSeek4Backend::maybe_save_routing_stats() { - if (!routing_stats_ || routing_stats_out_path_.empty()) return; - std::string err; - if (!routing_stats_->save_csv(routing_stats_out_path_, &err)) { - std::fprintf(stderr, "[deepseek4] failed to save routing stats %s: %s\n", - routing_stats_out_path_.c_str(), err.c_str()); - } -} - void DeepSeek4Backend::shutdown() { - maybe_save_routing_stats(); - for (int i = 0; i < PREFIX_SLOTS; i++) { - free_deepseek4_snapshot(snapshots_[i]); + for (auto & snap : snapshots_) { + free_deepseek4_snapshot(snap); } free_deepseek4_cache(cache_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - routing_stats_.reset(); - routing_stats_out_path_.clear(); - moe_placement_ = {}; free_deepseek4_weights(w_); - if (snap_backend_) { ggml_backend_free(snap_backend_); snap_backend_ = nullptr; } - if (backend_) { ggml_backend_free(backend_); backend_ = nullptr; } + if (snap_backend_) { + ggml_backend_free(snap_backend_); + snap_backend_ = nullptr; + } + if (backend_) { + ggml_backend_free(backend_); + backend_ = nullptr; + } + hc_state_.clear(); + last_logits_.clear(); } } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 508637867..d287bfaca 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -1,29 +1,26 @@ -// DeepSeek4Backend — ModelBackend for DeepSeek V4 Flash MLA+MoE models. -// -// Architecture: Multi-head Latent Attention (MLA), KV compression with -// learned compressors, Hierarchical Controller (HC), MoE with hash routing -// (first 3 layers) + top-k routing + shared expert. +// DeepSeek4Backend — full-weight single-backend DeepSeek4 runner. #pragma once #include "common/model_backend.h" #include "common/sampler.h" -#include "../common/moe_hybrid_placement.h" -#include "../common/moe_hybrid_routing_stats.h" -#include "../common/moe_hybrid_storage.h" -#include "../common/moe_hybrid_stream.h" #include "deepseek4_internal.h" -#include "ggml.h" #include "ggml-backend.h" -#include #include #include #include namespace dflash::common { +struct DeepSeek4BackendConfig { + const char * model_path = nullptr; + DevicePlacement device; + int chunk = 512; + int max_ctx = 0; +}; + class DeepSeek4Backend : public ModelBackend { public: explicit DeepSeek4Backend(const DeepSeek4BackendConfig & cfg); @@ -34,74 +31,52 @@ class DeepSeek4Backend : public ModelBackend { bool init(); - // ModelBackend interface void print_ready_banner() const override; - bool park(const std::string & what) override; bool unpark(const std::string & what) override; bool is_target_parked() const override { return parked_; } - GenerateResult generate_impl(const GenerateRequest & req, const DaemonIO & io) override; - bool snapshot_save(int slot) override; void snapshot_free(int slot) override; bool snapshot_used(int slot) const override; - int snapshot_cur_pos(int slot) const override; - + int snapshot_cur_pos(int slot) const override; GenerateResult restore_and_generate_impl(int slot, const GenerateRequest & req, const DaemonIO & io) override; - bool handle_compress(const std::string & line, const DaemonIO & io) override; - void free_drafter() override; - + void free_drafter() override {} void shutdown() override; private: - DeepSeek4BackendConfig cfg_; - ggml_backend_t backend_ = nullptr; - ggml_backend_t snap_backend_ = nullptr; - DeepSeek4Weights w_; - DeepSeek4Cache cache_; - bool parked_ = false; - - // Sampler - SamplerCfg sampler_; - std::mt19937_64 sampler_rng_{std::random_device{}()}; - - // Snapshots static constexpr int PREFIX_SLOTS = 64; - DeepSeek4Snapshot snapshots_[PREFIX_SLOTS]; - std::vector last_logits_; - - // Prefill prompt tokens in chunks, return absolute committed position. - int do_prefill(const std::vector & tokens, const DaemonIO & io, - int kv_offset = 0); - // Autoregressive decode loop. - bool do_decode(int committed, int n_gen, + int do_prefill(const std::vector & tokens, + const DaemonIO & io, + int kv_offset = 0, + int snap_pos = -1, + int snap_slot = -1); + bool do_decode(int committed, + int n_gen, std::vector & out_tokens, const DaemonIO & io, const BudgetHook & budget_hook = {}, bool * forced_close_out = nullptr); + bool load_full_model(); - bool init_hybrid_model(); - bool compute_uniform_hybrid_placement(const DeepSeek4Weights & w, - int max_ctx, - MoeHybridPlacement & out, - std::string * err) const; - void maybe_save_routing_stats(); - - std::shared_ptr moe_hybrid_; - MoeHybridPlacement moe_placement_; - MoeHybridStreamEngine stream_engine_; - // Expert IPC removed — layer split replaces expert split. - // Kept for compilation compatibility; init_hybrid_model() is no longer called - // from the layer-split path. - std::shared_ptr routing_stats_; - std::string routing_stats_out_path_; + DeepSeek4BackendConfig cfg_; + ggml_backend_t backend_ = nullptr; + ggml_backend_t snap_backend_ = nullptr; + DeepSeek4Weights w_; + DeepSeek4Cache cache_; + bool parked_ = false; + + SamplerCfg sampler_; + std::mt19937_64 sampler_rng_{std::random_device{}()}; + DeepSeek4Snapshot snapshots_[PREFIX_SLOTS]; + std::vector last_logits_; + std::vector hc_state_; }; } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_daemon.cpp b/server/src/deepseek4/deepseek4_daemon.cpp index fabc1c184..bb0a084bc 100644 --- a/server/src/deepseek4/deepseek4_daemon.cpp +++ b/server/src/deepseek4/deepseek4_daemon.cpp @@ -1,24 +1,25 @@ // DeepSeek4 daemon entry point implementation. #include "deepseek4_daemon.h" -#include "deepseek4_backend.h" + #include "common/daemon_loop.h" +#include "deepseek4_backend.h" #include +#include namespace dflash::common { int run_deepseek4_daemon(const char * model_path, - int gpu, - int stream_fd, - int max_ctx, - int chunk) { + int gpu, + int stream_fd, + int max_ctx, + int chunk) { DeepSeek4BackendConfig cfg; cfg.model_path = model_path; cfg.device.gpu = gpu; - cfg.stream_fd = stream_fd; - cfg.max_ctx = max_ctx; - cfg.chunk = chunk > 0 ? chunk : 512; + cfg.max_ctx = max_ctx; + cfg.chunk = chunk > 0 ? chunk : 512; auto backend = std::make_unique(cfg); if (!backend->init()) { @@ -28,8 +29,8 @@ int run_deepseek4_daemon(const char * model_path, DaemonLoopArgs loop_args; loop_args.stream_fd = stream_fd; - loop_args.chunk = cfg.chunk; - loop_args.max_ctx = max_ctx; + loop_args.chunk = cfg.chunk; + loop_args.max_ctx = max_ctx; return run_daemon(*backend, loop_args); } diff --git a/server/src/deepseek4/deepseek4_daemon.h b/server/src/deepseek4/deepseek4_daemon.h index d6b660cb0..d5b515c65 100644 --- a/server/src/deepseek4/deepseek4_daemon.h +++ b/server/src/deepseek4/deepseek4_daemon.h @@ -2,16 +2,14 @@ #pragma once -#include - namespace dflash::common { // Run the deepseek4 daemon loop. Called from main() when arch == "deepseek4". // Reads commands from stdin, writes tokens to stream_fd. int run_deepseek4_daemon(const char * model_path, - int gpu, - int stream_fd, - int max_ctx, - int chunk); + int gpu, + int stream_fd, + int max_ctx, + int chunk); } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 06fbe8a98..60bc7f3d7 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -34,6 +34,9 @@ namespace dflash::common { namespace { +static constexpr uint16_t F16_ZERO = 0x0000; +static constexpr uint16_t F16_NEG_INF = 0xFC00; + using Ds4TimingClock = std::chrono::steady_clock; static uint64_t ds4_elapsed_us(Ds4TimingClock::time_point start, @@ -114,6 +117,11 @@ struct DeepSeek4I64ArrayBinding { std::vector values; }; +struct DeepSeek4F16ArrayBinding { + ggml_tensor * tensor = nullptr; + std::vector values; +}; + static ggml_tensor * build_rms_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * weight, float eps); static ggml_tensor * build_clamped_swiglu(ggml_context * ctx, @@ -288,6 +296,96 @@ struct DeepSeek4LayerRangeScratch { } }; +struct DeepSeek4BatchedHcWorkspace { + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + ggml_backend_t backend = nullptr; + int n_tokens = 0; + int n_embd = 0; + int n_hc = 0; + ggml_tensor * hc_state = nullptr; + ggml_tensor * next_hc = nullptr; + ggml_tensor * working = nullptr; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + ggml_tensor * final_embd = nullptr; + + ~DeepSeek4BatchedHcWorkspace() { + free(); + } + + void free() { + if (buf) { + ggml_backend_buffer_free(buf); + buf = nullptr; + } + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } + backend = nullptr; + n_tokens = 0; + n_embd = 0; + n_hc = 0; + hc_state = nullptr; + next_hc = nullptr; + working = nullptr; + post = nullptr; + comb = nullptr; + final_embd = nullptr; + } + + bool valid(ggml_backend_t b, int tokens, int embd, int hc) const { + return ctx && buf && backend == b && n_tokens == tokens && + n_embd == embd && n_hc == hc && hc_state && next_hc && + working && post && comb && final_embd; + } + + bool ensure(ggml_backend_t b, int tokens, int embd, int hc) { + if (valid(b, tokens, embd, hc)) { + return true; + } + free(); + + const int hc_dim = embd * hc; + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 8 + 4096; + params.mem_buffer = nullptr; + params.no_alloc = true; + ctx = ggml_init(params); + if (!ctx) { + return false; + } + + hc_state = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hc_dim, tokens); + next_hc = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hc_dim, tokens); + working = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, embd, tokens); + post = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hc, tokens); + comb = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hc * hc, tokens); + final_embd = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, embd, tokens); + if (!hc_state || !next_hc || !working || !post || !comb || + !final_embd) { + free(); + return false; + } + + buf = ggml_backend_alloc_ctx_tensors(ctx, b); + if (!buf) { + free(); + return false; + } + backend = b; + n_tokens = tokens; + n_embd = embd; + n_hc = hc; + return true; + } + + void swap_hc() { + std::swap(hc_state, next_hc); + } +}; + static bool build_cached_decode_ffn_graph( DeepSeek4CachedDecodeFfnGraph & out, ggml_backend_t backend, @@ -825,7 +923,8 @@ static ggml_tensor * build_mla_attention( const DeepSeek4AttentionGraphInputs * cached_inputs, std::vector & i32_inputs, std::vector & i32_array_inputs, - std::vector & i64_array_inputs) { + std::vector & i64_array_inputs, + std::vector & f16_array_inputs) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; @@ -885,6 +984,24 @@ static ggml_tensor * build_mla_attention( rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + const int token_pos = kv_start + n_tokens - 1; + const int n_past_raw = (n_tokens == 1) + ? std::max(0, std::min(kv_start + 1, w.n_swa) - 1) + : std::min(kv_start, w.n_swa); + const int n_cur_raw = n_tokens; + const int n_raw = n_past_raw + n_cur_raw; + const int n_comp_attn = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + const int n_attn = n_raw + n_comp_attn; + const float kq_scale = 1.0f / sqrtf((float)head_dim); + + ggml_tensor * past_raw_f32 = nullptr; + if (n_past_raw > 0) { + past_raw_f32 = ggml_cast(ctx, + ggml_view_2d(ctx, lc.raw_kv, head_dim, n_past_raw, lc.raw_kv->nb[1], 0), + GGML_TYPE_F32); + ggml_build_forward_expand(gf, past_raw_f32); + } + // ── Store ALL KV rows in the raw SWA ring ───────────────────── // For decode (n_tokens=1): write single row. For prefill: write all rows. if (cached_inputs && cached_inputs->raw_kv_rows) { @@ -901,7 +1018,6 @@ static ggml_tensor * build_mla_attention( ggml_build_forward_expand(gf, ggml_cpy(ctx, ggml_cast(ctx, kv_row, GGML_TYPE_F16), kv_slot)); } } - const int token_pos = kv_start + n_tokens - 1; // ── Learned compression update ────────────────────────────────── ggml_tensor * cur_last = ggml_view_2d( @@ -949,30 +1065,10 @@ static ggml_tensor * build_mla_attention( // q: [head_dim, n_head, n_tokens] (after RoPE) // raw_kv: [head_dim, n_swa] F16 persistent ring buffer (single KV head, shared) // comp_kv: [head_dim, comp_cap] F16 compressed rows. - // n_raw = min(kv_start + n_tokens, n_swa) - const int n_raw = std::min(kv_start + n_tokens, w.n_swa); - const int n_comp_attn = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; - const int n_attn = n_raw + n_comp_attn; - const float kq_scale = 1.0f / sqrtf((float)head_dim); - - // Get valid KV rows. For single-token decode/prefill, include the current - // in-graph KV row directly; otherwise attention can race the side-effecting - // cache write and see the previous contents of the raw KV slot. - ggml_tensor * kv_f32 = nullptr; - if (n_tokens == 1) { - ggml_tensor * cur_kv_f32 = ggml_cast(ctx, kv, GGML_TYPE_F32); - if (n_raw > 1) { - ggml_tensor * prev_f32 = ggml_cast(ctx, - ggml_view_2d(ctx, lc.raw_kv, head_dim, n_raw - 1, - lc.raw_kv->nb[1], 0), - GGML_TYPE_F32); - kv_f32 = ggml_concat(ctx, prev_f32, cur_kv_f32, 1); - } else { - kv_f32 = cur_kv_f32; - } - } else { - kv_f32 = ggml_cast(ctx, ggml_view_2d(ctx, lc.raw_kv, head_dim, n_raw, - lc.raw_kv->nb[1], 0), GGML_TYPE_F32); + ggml_tensor * cur_kv_f32 = ggml_cast(ctx, kv, GGML_TYPE_F32); + ggml_tensor * kv_f32 = cur_kv_f32; + if (past_raw_f32) { + kv_f32 = ggml_concat(ctx, past_raw_f32, cur_kv_f32, 1); } if (n_comp_attn > 0 && lc.comp_kv) { ggml_tensor * comp_f32 = ggml_cast(ctx, @@ -990,6 +1086,41 @@ static ggml_tensor * build_mla_attention( ggml_tensor * scores = ggml_mul_mat(ctx, kv_f32, q_flat); scores = ggml_scale(ctx, scores, kq_scale); + std::vector attn_mask_values; + if (n_tokens > 1 && !cached_inputs) { + const int cols = n_head * n_tokens; + attn_mask_values.assign((size_t)n_attn * (size_t)cols, F16_NEG_INF); + for (int ti = 0; ti < n_tokens; ++ti) { + const int abs_q = kv_start + ti; + for (int h = 0; h < n_head; ++h) { + const int col = ti * n_head + h; + for (int k = 0; k < n_past_raw; ++k) { + int abs_k = k; + if (kv_start > 0) { + const int last_past = kv_start - 1; + const int delta = (last_past - k) % w.n_swa; + abs_k = last_past - (delta < 0 ? delta + w.n_swa : delta); + } + if (abs_k >= 0 && abs_k <= abs_q && abs_k > abs_q - w.n_swa) { + attn_mask_values[(size_t)col * (size_t)n_attn + (size_t)k] = F16_ZERO; + } + } + for (int c = 0; c < n_cur_raw; ++c) { + const int abs_k = kv_start + c; + if (abs_k <= abs_q && abs_k > abs_q - w.n_swa) { + attn_mask_values[(size_t)col * (size_t)n_attn + (size_t)(n_past_raw + c)] = F16_ZERO; + } + } + for (int c = 0; c < n_comp_attn; ++c) { + const int comp_end = (c + 1) * ratio - 1; + if (comp_end <= abs_q) { + attn_mask_values[(size_t)col * (size_t)n_attn + (size_t)(n_raw + c)] = F16_ZERO; + } + } + } + } + } + // Sink-aware softmax: DS4 adds one learned per-head sink logit to the // denominator, but the sink contributes no value vector. ggml_tensor * probs = nullptr; @@ -1000,11 +1131,37 @@ static ggml_tensor * build_mla_attention( sink_scores = ggml_repeat(ctx, sink_scores, sink_shape); } ggml_tensor * scores_with_sink = ggml_concat(ctx, scores, sink_scores, 0); - ggml_tensor * probs_with_sink = ggml_soft_max(ctx, scores_with_sink); + ggml_tensor * sink_mask = nullptr; + if (!attn_mask_values.empty()) { + const int cols = n_head * n_tokens; + sink_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, n_attn + 1, cols); + ggml_set_input(sink_mask); + DeepSeek4F16ArrayBinding binding; + binding.tensor = sink_mask; + binding.values.assign((size_t)(n_attn + 1) * (size_t)cols, F16_NEG_INF); + for (int col = 0; col < cols; ++col) { + std::memcpy(binding.values.data() + (size_t)col * (size_t)(n_attn + 1), + attn_mask_values.data() + (size_t)col * (size_t)n_attn, + sizeof(uint16_t) * (size_t)n_attn); + binding.values[(size_t)col * (size_t)(n_attn + 1) + (size_t)n_attn] = F16_ZERO; + } + f16_array_inputs.push_back(std::move(binding)); + } + ggml_tensor * probs_with_sink = sink_mask + ? ggml_soft_max_ext(ctx, scores_with_sink, sink_mask, 1.0f, 0.0f) + : ggml_soft_max(ctx, scores_with_sink); probs = ggml_view_2d(ctx, probs_with_sink, n_attn, n_head * n_tokens, probs_with_sink->nb[1], 0); } else { - probs = ggml_soft_max(ctx, scores); + ggml_tensor * attn_mask = nullptr; + if (!attn_mask_values.empty()) { + const int cols = n_head * n_tokens; + attn_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, n_attn, cols); + ggml_set_input(attn_mask); + f16_array_inputs.push_back({attn_mask, std::move(attn_mask_values)}); + } + probs = attn_mask ? ggml_soft_max_ext(ctx, scores, attn_mask, 1.0f, 0.0f) + : ggml_soft_max(ctx, scores); } // probs: [n_attn, n_head*n_tokens] @@ -1175,10 +1332,12 @@ static bool build_cached_decode_attn_graph( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; + std::vector f16_array_inputs; ggml_tensor * normed = build_rms_norm(out.sg.ctx, out.sg.inp_embed, L.attn_norm, w.rms_eps); out.sg.hidden_states = build_mla_attention(out.sg.ctx, out.sg.gf, normed, w, L, lc, layer_idx, kv_start, 1, &out.inputs, - i32_inputs, i32_array_inputs, i64_array_inputs); + i32_inputs, i32_array_inputs, i64_array_inputs, + f16_array_inputs); if (!out.sg.hidden_states) { out.free(); return false; @@ -2272,13 +2431,14 @@ static bool deepseek4_step_hybrid( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; + std::vector f16_array_inputs; ggml_cgraph * gf = ggml_new_graph(ctx); ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps); ggml_tensor * attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, kv_start, n_tokens, nullptr, i32_inputs, i32_array_inputs, - i64_array_inputs); + i64_array_inputs, f16_array_inputs); // Output just attn_out (HC post handles the residual mixing) ggml_build_forward_expand(gf, attn_out); ggml_gallocr_t attn_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); @@ -2302,6 +2462,10 @@ static bool deepseek4_step_hybrid( ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, sizeof(int64_t) * binding.values.size()); } + for (const DeepSeek4F16ArrayBinding & binding : f16_array_inputs) { + ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, + sizeof(uint16_t) * binding.values.size()); + } const auto attn_compute_t0 = Ds4TimingClock::now(); bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; if (telemetry) telemetry->attn_compute_us += ds4_elapsed_us(attn_compute_t0, Ds4TimingClock::now()); @@ -2698,6 +2862,7 @@ bool deepseek4_step( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; + std::vector f16_array_inputs; // Layer loop for (int il = 0; il < n_layer; il++) { @@ -2715,7 +2880,7 @@ bool deepseek4_step( ggml_tensor * attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, kv_start, n_tokens, nullptr, i32_inputs, i32_array_inputs, - i64_array_inputs); + i64_array_inputs, f16_array_inputs); // ── Residual ──────────────────────────────────────────────── cur = ggml_add(ctx, cur, attn_out); @@ -2769,6 +2934,10 @@ bool deepseek4_step( ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, sizeof(int64_t) * binding.values.size()); } + for (const DeepSeek4F16ArrayBinding & binding : f16_array_inputs) { + ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, + sizeof(uint16_t) * binding.values.size()); + } // Compute if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { @@ -2913,10 +3082,20 @@ bool deepseek4_step_layer_range( std::vector & ffn_out_host = scratch.ffn_out_host; std::vector & hash_expert_ids_host = scratch.hash_expert_ids; const bool reuse_decode_graphs = (n_tokens == 1); + const bool use_batched_gpu_hc = (n_tokens > 1) && ds4_backend_is_gpu(backend); const bool use_backend_decode_hc = reuse_decode_graphs && ds4_backend_is_gpu(backend); const bool use_backend_decode_hc_direct = use_backend_decode_hc && ds4_backend_is_hip(backend); const bool use_backend_decode_hc_graph = use_backend_decode_hc && !use_backend_decode_hc_direct; + static thread_local DeepSeek4BatchedHcWorkspace batched_hc; + if (use_batched_gpu_hc) { + if (!batched_hc.ensure(backend, n_tokens, n_embd, n_hc)) { + std::fprintf(stderr, "[deepseek4] batched HC workspace alloc failed\n"); + return false; + } + ggml_backend_tensor_set(batched_hc.hc_state, hc_state.data(), 0, + sizeof(float) * hc_state.size()); + } const ggml_tensor * hc_state_backend = nullptr; if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { if (!cached_decode_hc_post_graph.valid() || @@ -2945,7 +3124,27 @@ bool deepseek4_step_layer_range( // ── HC pre (attention) ────────────────────────────────────── const auto hc_pre_attn_t0 = Ds4TimingClock::now(); - if (use_backend_decode_hc_direct) { + if (use_batched_gpu_hc) { + const auto hc_pre_attn_compute_t0 = Ds4TimingClock::now(); + if (!deepseek4_cuda_hc_pre_batch_device( + batched_hc.hc_state->data, + L.hc_attn_fn ? L.hc_attn_fn->data : nullptr, + L.hc_attn_scale ? L.hc_attn_scale->data : nullptr, + L.hc_attn_base ? L.hc_attn_base->data : nullptr, + n_tokens, + n_embd, + n_hc, + w.n_hc_sinkhorn_iter, + w.hc_eps, + batched_hc.working->data, + batched_hc.post->data, + batched_hc.comb->data)) { + std::fprintf(stderr, "[deepseek4] batched hc-pre failed layer %d attn\n", il); + return false; + } + if (telemetry) telemetry->hc_pre_compute_us += ds4_elapsed_us(hc_pre_attn_compute_t0, Ds4TimingClock::now()); + attn_in_backend = batched_hc.working; + } else if (use_backend_decode_hc_direct) { const auto hc_pre_attn_input_t0 = Ds4TimingClock::now(); ggml_backend_tensor_set(cached_decode_hc_post_graph.residual_hc, hc_state.data(), 0, sizeof(float) * hc_state.size()); @@ -3023,7 +3222,11 @@ bool deepseek4_step_layer_range( // ── Build & run attention graph ───────────────────────────── { const int token_pos = kv_start + n_tokens - 1; - const bool reuse_decode_attn = n_tokens == 1; + // The cached single-token attention graph updates KV/compressor state + // with ggml_set_rows. On HIP this path can fault after enough prompt + // tokens; use the dynamic graph until the cached state-update path is + // made safe. + const bool reuse_decode_attn = false; ggml_tensor * attn_out = nullptr; ggml_cgraph * gf = nullptr; ggml_context * ctx = nullptr; @@ -3078,7 +3281,7 @@ bool deepseek4_step_layer_range( if (ratio > 0) { const int pos_mod = token_pos % ratio; const int32_t ape_row = pos_mod; - const int64_t state_row = ratio + pos_mod; + const int64_t state_row = (ratio == 4) ? (ratio + pos_mod) : pos_mod; const int64_t comp_row = token_pos / ratio; const int32_t comp_pos = token_pos + 1 - ratio; const bool flush_boundary = ((token_pos + 1) % ratio) == 0; @@ -3118,6 +3321,7 @@ bool deepseek4_step_layer_range( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; + std::vector f16_array_inputs; const size_t graph_size = n_tokens > 1 ? 32768 : 2048; gf = ggml_new_graph_custom(ctx, graph_size, false); @@ -3125,7 +3329,7 @@ bool deepseek4_step_layer_range( attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, kv_start, n_tokens, nullptr, i32_inputs, i32_array_inputs, - i64_array_inputs); + i64_array_inputs, f16_array_inputs); ggml_set_output(attn_out); ggml_build_forward_expand(gf, attn_out); @@ -3153,6 +3357,8 @@ bool deepseek4_step_layer_range( ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(int32_t) * b.values.size()); for (const auto & b : i64_array_inputs) ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(int64_t) * b.values.size()); + for (const auto & b : f16_array_inputs) + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(uint16_t) * b.values.size()); } const auto attn_compute_t0 = Ds4TimingClock::now(); @@ -3162,7 +3368,24 @@ bool deepseek4_step_layer_range( return false; } if (telemetry) telemetry->attn_compute_us += ds4_elapsed_us(attn_compute_t0, Ds4TimingClock::now()); - if (use_backend_decode_hc_graph) { + if (use_batched_gpu_hc) { + const auto hc_post_attn_t0 = Ds4TimingClock::now(); + if (!deepseek4_cuda_hc_post_batch_device( + batched_hc.hc_state->data, + attn_out->data, + batched_hc.post->data, + batched_hc.comb->data, + n_tokens, + n_embd, + n_hc, + batched_hc.next_hc->data)) { + std::fprintf(stderr, "[deepseek4] batched hc-post failed layer %d attn\n", il); + if (ctx) ggml_free(ctx); + return false; + } + batched_hc.swap_hc(); + if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); + } else if (use_backend_decode_hc_graph) { if (hc_state_backend != cached_decode_hc_post_graph.residual_hc) { ggml_backend_tensor_copy(hc_state_backend, cached_decode_hc_post_graph.residual_hc); } @@ -3189,7 +3412,7 @@ bool deepseek4_step_layer_range( if (ctx) ggml_free(ctx); // ── HC post (attention) ───────────────────────────────── - if (!use_backend_decode_hc_graph) { + if (!use_batched_gpu_hc && !use_backend_decode_hc_graph) { const auto hc_post_attn_t0 = Ds4TimingClock::now(); hc_post_batch(next_hc, attn_out_host.data(), @@ -3206,7 +3429,27 @@ bool deepseek4_step_layer_range( // ── HC pre (FFN) ──────────────────────────────────────────── const auto hc_pre_ffn_t0 = Ds4TimingClock::now(); - if (use_backend_decode_hc_direct) { + if (use_batched_gpu_hc) { + const auto hc_pre_ffn_compute_t0 = Ds4TimingClock::now(); + if (!deepseek4_cuda_hc_pre_batch_device( + batched_hc.hc_state->data, + L.hc_ffn_fn ? L.hc_ffn_fn->data : nullptr, + L.hc_ffn_scale ? L.hc_ffn_scale->data : nullptr, + L.hc_ffn_base ? L.hc_ffn_base->data : nullptr, + n_tokens, + n_embd, + n_hc, + w.n_hc_sinkhorn_iter, + w.hc_eps, + batched_hc.working->data, + batched_hc.post->data, + batched_hc.comb->data)) { + std::fprintf(stderr, "[deepseek4] batched hc-pre failed layer %d ffn\n", il); + return false; + } + if (telemetry) telemetry->hc_pre_compute_us += ds4_elapsed_us(hc_pre_ffn_compute_t0, Ds4TimingClock::now()); + ffn_in_backend = batched_hc.working; + } else if (use_backend_decode_hc_direct) { const auto hc_pre_ffn_input_t0 = Ds4TimingClock::now(); ggml_backend_tensor_set(cached_decode_hc_post_graph.residual_hc, hc_state.data(), 0, sizeof(float) * hc_state.size()); @@ -3335,7 +3578,23 @@ bool deepseek4_step_layer_range( return false; } - if (use_backend_decode_hc_graph) { + if (use_batched_gpu_hc) { + const auto hc_post_ffn_t0 = Ds4TimingClock::now(); + if (!deepseek4_cuda_hc_post_batch_device( + batched_hc.hc_state->data, + ffn_out->data, + batched_hc.post->data, + batched_hc.comb->data, + n_tokens, + n_embd, + n_hc, + batched_hc.next_hc->data)) { + std::fprintf(stderr, "[deepseek4] batched hc-post failed layer %d ffn\n", il); + return false; + } + batched_hc.swap_hc(); + if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); + } else if (use_backend_decode_hc_graph) { if (hc_state_backend != cached_decode_hc_post_graph.residual_hc) { ggml_backend_tensor_copy(hc_state_backend, cached_decode_hc_post_graph.residual_hc); } @@ -3360,7 +3619,7 @@ bool deepseek4_step_layer_range( } // ── HC post (FFN) ─────────────────────────────────────── - if (!use_backend_decode_hc_graph) { + if (!use_batched_gpu_hc && !use_backend_decode_hc_graph) { const auto hc_post_ffn_t0 = Ds4TimingClock::now(); hc_post_batch(next_hc, ffn_out_host.data(), @@ -3379,19 +3638,39 @@ bool deepseek4_step_layer_range( if (use_backend_decode_hc_graph && hc_state_backend) { ggml_backend_tensor_get(hc_state_backend, hc_state.data(), 0, sizeof(float) * hc_state.size()); } + if (use_batched_gpu_hc) { + ggml_backend_tensor_get(batched_hc.hc_state, hc_state.data(), 0, + sizeof(float) * hc_state.size()); + } // ── Output: HC pre → norm → lm_head (or return hidden state) ──────── if (is_last_shard && out_logits) { // Final HC pre for output const auto output_t0 = Ds4TimingClock::now(); std::vector & final_embd = scratch.final_embd; - hc_output_batch(final_embd, - hc_state.data(), - hc_output_weights_range, - n_tokens, - n_embd, - n_hc, - w.hc_eps); + if (use_batched_gpu_hc) { + if (!deepseek4_cuda_hc_output_batch_device( + batched_hc.hc_state->data, + w.output_hc_fn ? w.output_hc_fn->data : nullptr, + w.output_hc_scale ? w.output_hc_scale->data : nullptr, + w.output_hc_base ? w.output_hc_base->data : nullptr, + n_tokens, + n_embd, + n_hc, + w.hc_eps, + batched_hc.final_embd->data)) { + std::fprintf(stderr, "[deepseek4] batched hc-output failed\n"); + return false; + } + } else { + hc_output_batch(final_embd, + hc_state.data(), + hc_output_weights_range, + n_tokens, + n_embd, + n_hc, + w.hc_eps); + } if (reuse_decode_graphs) { if (!cached_decode_output_graph.valid() || @@ -3440,7 +3719,11 @@ bool deepseek4_step_layer_range( ggml_free(ctx); return false; } - ggml_backend_tensor_set(inp, final_embd.data(), 0, sizeof(float) * final_embd.size()); + if (use_batched_gpu_hc) { + ggml_backend_tensor_copy(batched_hc.final_embd, inp); + } else { + ggml_backend_tensor_set(inp, final_embd.data(), 0, sizeof(float) * final_embd.size()); + } if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { ggml_free(ctx); return false; @@ -3455,6 +3738,10 @@ bool deepseek4_step_layer_range( if (telemetry) telemetry->output_us += ds4_elapsed_us(output_t0, Ds4TimingClock::now()); } else if (out_logits) { // Return full HC state for next shard (all n_hc streams) + if (use_batched_gpu_hc) { + ggml_backend_tensor_get(batched_hc.hc_state, hc_state.data(), 0, + sizeof(float) * hc_state.size()); + } out_logits->resize((size_t)hc_dim * n_tokens); memcpy(out_logits->data(), hc_state.data(), sizeof(float) * hc_dim * n_tokens); } @@ -3756,6 +4043,7 @@ void free_deepseek4_snapshot(DeepSeek4Snapshot & s) { if (s.ctx) { ggml_free(s.ctx); s.ctx = nullptr; } s.layers.clear(); s.cur_pos = 0; + s.last_logits.clear(); s.hc_state_snap = nullptr; } diff --git a/server/src/deepseek4/deepseek4_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index 8f784d1f2..8bdf694b0 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -246,6 +247,226 @@ __global__ void hc_finish_kernel(const float * hc_state, } } +__global__ void hc_pre_batch_kernel(const float * hc_state, + const __half * fn, + const float * scale, + const float * base, + int n_tokens, + int n_embd, + int n_hc, + int sinkhorn_iters, + float eps, + float * working, + float * post, + float * comb) { + __shared__ float smem[kThreads]; + __shared__ float mix[kMaxMixDim]; + __shared__ float split[kMaxMixDim]; + const int t = blockIdx.x; + const int tid = threadIdx.x; + if (t >= n_tokens) { + return; + } + + const int hc_dim = n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + const float * token_hc = hc_state + (size_t)t * (size_t)hc_dim; + + float sumsq = 0.0f; + for (int c = tid; c < hc_dim; c += blockDim.x) { + const float v = token_hc[c]; + sumsq += v * v; + } + smem[tid] = sumsq; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] += smem[tid + stride]; + __syncthreads(); + } + const float inv_rms = rsqrtf(smem[0] / (float)hc_dim + eps); + + for (int row = 0; row < mix_dim; ++row) { + float dot = 0.0f; + const __half * w = fn + (size_t)row * (size_t)hc_dim; + for (int c = tid; c < hc_dim; c += blockDim.x) { + dot += __half2float(w[c]) * token_hc[c] * inv_rms; + } + smem[tid] = dot; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] += smem[tid + stride]; + __syncthreads(); + } + if (tid == 0) { + mix[row] = smem[0]; + } + __syncthreads(); + } + + if (tid == 0) { + const float pre_scale = scale[0]; + const float post_scale = scale[1]; + const float comb_scale = scale[2]; + const float sinkhorn_eps = 1.0e-6f; + + for (int i = 0; i < n_hc; ++i) { + split[i] = hc_sigmoid(mix[i] * pre_scale + base[i]) + sinkhorn_eps; + } + for (int i = 0; i < n_hc; ++i) { + const float v = 2.0f * hc_sigmoid(mix[n_hc + i] * post_scale + base[n_hc + i]); + split[n_hc + i] = v; + post[(size_t)t * (size_t)n_hc + (size_t)i] = v; + } + + float c[kMaxHc * kMaxHc]; + for (int dst = 0; dst < n_hc; ++dst) { + float row_max = -1.0e30f; + for (int src = 0; src < n_hc; ++src) { + const int idx = src + dst * n_hc; + const float v = mix[2 * n_hc + idx] * comb_scale + base[2 * n_hc + idx]; + c[idx] = v; + row_max = v > row_max ? v : row_max; + } + float row_sum = 0.0f; + for (int src = 0; src < n_hc; ++src) { + const int idx = src + dst * n_hc; + c[idx] = expf(c[idx] - row_max); + row_sum += c[idx]; + } + const float inv = 1.0f / row_sum; + for (int src = 0; src < n_hc; ++src) { + c[src + dst * n_hc] = c[src + dst * n_hc] * inv + sinkhorn_eps; + } + } + for (int src = 0; src < n_hc; ++src) { + float sum = 0.0f; + for (int dst = 0; dst < n_hc; ++dst) sum += c[src + dst * n_hc]; + const float inv = 1.0f / (sum + sinkhorn_eps); + for (int dst = 0; dst < n_hc; ++dst) c[src + dst * n_hc] *= inv; + } + for (int iter = 1; iter < sinkhorn_iters; ++iter) { + for (int dst = 0; dst < n_hc; ++dst) { + float sum = 0.0f; + for (int src = 0; src < n_hc; ++src) sum += c[src + dst * n_hc]; + const float inv = 1.0f / (sum + sinkhorn_eps); + for (int src = 0; src < n_hc; ++src) c[src + dst * n_hc] *= inv; + } + for (int src = 0; src < n_hc; ++src) { + float sum = 0.0f; + for (int dst = 0; dst < n_hc; ++dst) sum += c[src + dst * n_hc]; + const float inv = 1.0f / (sum + sinkhorn_eps); + for (int dst = 0; dst < n_hc; ++dst) c[src + dst * n_hc] *= inv; + } + } + + float * token_comb = comb + (size_t)t * (size_t)n_hc * (size_t)n_hc; + for (int i = 0; i < n_hc * n_hc; ++i) { + split[2 * n_hc + i] = c[i]; + token_comb[i] = c[i]; + } + } + __syncthreads(); + + float * token_working = working + (size_t)t * (size_t)n_embd; + for (int d = tid; d < n_embd; d += blockDim.x) { + float acc = 0.0f; + for (int h = 0; h < n_hc; ++h) { + acc += split[h] * token_hc[(size_t)h * (size_t)n_embd + (size_t)d]; + } + token_working[d] = acc; + } +} + +__global__ void hc_post_batch_kernel(const float * residual_hc, + const float * block_out, + const float * post, + const float * comb, + int n_total, + int n_embd, + int n_hc, + float * out_hc) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= n_total) { + return; + } + const int d = idx % n_embd; + const int h = (idx / n_embd) % n_hc; + const int t = idx / (n_embd * n_hc); + const size_t hc_dim = (size_t)n_embd * (size_t)n_hc; + const float * token_residual = residual_hc + (size_t)t * hc_dim; + const float * token_block = block_out + (size_t)t * (size_t)n_embd; + const float * token_post = post + (size_t)t * (size_t)n_hc; + const float * token_comb = comb + (size_t)t * (size_t)n_hc * (size_t)n_hc; + + float acc = token_block[d] * token_post[h]; + for (int src = 0; src < n_hc; ++src) { + acc += token_comb[h + src * n_hc] * + token_residual[(size_t)src * (size_t)n_embd + (size_t)d]; + } + out_hc[(size_t)idx] = acc; +} + +__global__ void hc_output_batch_kernel(const float * hc_state, + const __half * fn, + const float * scale, + const float * base, + int n_tokens, + int n_embd, + int n_hc, + float eps, + float * final_embd) { + __shared__ float smem[kThreads]; + __shared__ float pre[kMaxHc]; + const int t = blockIdx.x; + const int tid = threadIdx.x; + if (t >= n_tokens) { + return; + } + + const int hc_dim = n_embd * n_hc; + const float * token_hc = hc_state + (size_t)t * (size_t)hc_dim; + + float sumsq = 0.0f; + for (int c = tid; c < hc_dim; c += blockDim.x) { + const float v = token_hc[c]; + sumsq += v * v; + } + smem[tid] = sumsq; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] += smem[tid + stride]; + __syncthreads(); + } + const float inv_rms = rsqrtf(smem[0] / (float)hc_dim + eps); + + for (int row = 0; row < n_hc; ++row) { + float dot = 0.0f; + const __half * w = fn + (size_t)row * (size_t)hc_dim; + for (int c = tid; c < hc_dim; c += blockDim.x) { + dot += __half2float(w[c]) * token_hc[c] * inv_rms; + } + smem[tid] = dot; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] += smem[tid + stride]; + __syncthreads(); + } + if (tid == 0) { + pre[row] = hc_sigmoid(smem[0] * scale[0] + base[row]) + 1.0e-6f; + } + __syncthreads(); + } + + float * token_out = final_embd + (size_t)t * (size_t)n_embd; + for (int d = tid; d < n_embd; d += blockDim.x) { + float acc = 0.0f; + for (int h = 0; h < n_hc; ++h) { + acc += pre[h] * token_hc[(size_t)h * (size_t)n_embd + (size_t)d]; + } + token_out[d] = acc; + } +} + bool hc_pre_device_locked(const void * hc_state_device, const void * fn_device, const void * scale_device, @@ -551,4 +772,127 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, true); } +bool deepseek4_cuda_hc_pre_batch_device(const void * hc_state_device, + const void * fn_device, + const void * scale_device, + const void * base_device, + int n_tokens, + int n_embd, + int n_hc, + int sinkhorn_iters, + float eps, + void * working_device, + void * post_device, + void * comb_device) { + if (!hc_state_device || !fn_device || !scale_device || !base_device || + !working_device || !post_device || !comb_device || + n_tokens <= 0 || n_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { + return false; + } + const int mix_dim = 2 * n_hc + n_hc * n_hc; + if (mix_dim > kMaxMixDim) { + return false; + } + + hc_pre_batch_kernel<<>>( + static_cast(hc_state_device), + static_cast(fn_device), + static_cast(scale_device), + static_cast(base_device), + n_tokens, + n_embd, + n_hc, + sinkhorn_iters, + eps, + static_cast(working_device), + static_cast(post_device), + static_cast(comb_device)); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + hc_log_cuda_error("batched hc-pre kernel", err); + return false; + } + err = cudaDeviceSynchronize(); + if (err != cudaSuccess) { + hc_log_cuda_error("batched hc-pre sync", err); + return false; + } + return true; +} + +bool deepseek4_cuda_hc_post_batch_device(const void * residual_hc_device, + const void * block_out_device, + const void * post_device, + const void * comb_device, + int n_tokens, + int n_embd, + int n_hc, + void * out_hc_device) { + if (!residual_hc_device || !block_out_device || !post_device || !comb_device || + !out_hc_device || n_tokens <= 0 || n_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { + return false; + } + const int64_t total = (int64_t)n_tokens * (int64_t)n_hc * (int64_t)n_embd; + if (total <= 0 || total > std::numeric_limits::max()) { + return false; + } + const int blocks = (int)((total + kThreads - 1) / kThreads); + hc_post_batch_kernel<<>>( + static_cast(residual_hc_device), + static_cast(block_out_device), + static_cast(post_device), + static_cast(comb_device), + (int)total, + n_embd, + n_hc, + static_cast(out_hc_device)); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + hc_log_cuda_error("batched hc-post kernel", err); + return false; + } + err = cudaDeviceSynchronize(); + if (err != cudaSuccess) { + hc_log_cuda_error("batched hc-post sync", err); + return false; + } + return true; +} + +bool deepseek4_cuda_hc_output_batch_device(const void * hc_state_device, + const void * fn_device, + const void * scale_device, + const void * base_device, + int n_tokens, + int n_embd, + int n_hc, + float eps, + void * final_embd_device) { + if (!hc_state_device || !fn_device || !scale_device || !base_device || + !final_embd_device || n_tokens <= 0 || n_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { + return false; + } + hc_output_batch_kernel<<>>( + static_cast(hc_state_device), + static_cast(fn_device), + static_cast(scale_device), + static_cast(base_device), + n_tokens, + n_embd, + n_hc, + eps, + static_cast(final_embd_device)); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + hc_log_cuda_error("batched hc-output kernel", err); + return false; + } + err = cudaDeviceSynchronize(); + if (err != cudaSuccess) { + hc_log_cuda_error("batched hc-output sync", err); + return false; + } + return true; +} + } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_hc_cuda.h b/server/src/deepseek4/deepseek4_hc_cuda.h index d2268ac12..f7c1a6758 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.h +++ b/server/src/deepseek4/deepseek4_hc_cuda.h @@ -47,4 +47,36 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, void * post_device, void * comb_device); +bool deepseek4_cuda_hc_pre_batch_device(const void * hc_state_device, + const void * fn_device, + const void * scale_device, + const void * base_device, + int n_tokens, + int n_embd, + int n_hc, + int sinkhorn_iters, + float eps, + void * working_device, + void * post_device, + void * comb_device); + +bool deepseek4_cuda_hc_post_batch_device(const void * residual_hc_device, + const void * block_out_device, + const void * post_device, + const void * comb_device, + int n_tokens, + int n_embd, + int n_hc, + void * out_hc_device); + +bool deepseek4_cuda_hc_output_batch_device(const void * hc_state_device, + const void * fn_device, + const void * scale_device, + const void * base_device, + int n_tokens, + int n_embd, + int n_hc, + float eps, + void * final_embd_device); + } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 31c086394..70468f6ee 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -265,16 +265,6 @@ struct DeepSeek4Cache { struct DeepSeek4Snapshot; -// ─── Configuration ────────────────────────────────────────────────────── - -struct DeepSeek4BackendConfig { - const char * model_path = nullptr; - DevicePlacement device; - int stream_fd = -1; - int chunk = 512; // prefill chunk size - int max_ctx = 0; // 0 = auto from SWA + compression capacity -}; - // ─── Function declarations ────────────────────────────────────────────── bool load_deepseek4_gguf(const std::string & path, @@ -362,6 +352,7 @@ bool build_deepseek4_moe_hybrid_storage_from_file_with_mmap( struct DeepSeek4Snapshot { int cur_pos = 0; ggml_tensor * hc_state_snap = nullptr; + std::vector last_logits; // Per-layer: raw KV + compressed KV snapshots struct LayerSnap { ggml_tensor * raw_kv = nullptr; diff --git a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp index a6054aa8b..e3158c19a 100644 --- a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp +++ b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp @@ -239,6 +239,10 @@ bool DeepSeek4LayerSplitAdapter::init() { return false; } hc_state_.resize(hc_state_elements(shard.weights), 0.0f); + auto shard_metas = layer_split_shard_metas(shards_); + if (!init_layer_split_snapshot_backends(shard_metas, snapshot_backends_, "deepseek4-split")) { + return false; + } std::fprintf(stderr, "[deepseek4-split] single shard: gpu=%d layers=[0,43) (+output)\n", shard.gpu); return true; } @@ -440,7 +444,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,24 +456,15 @@ 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); } - // 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. + // Shard 0 consumes token embeddings; later shards consume the full + // token-major HC boundary state emitted by the previous shard. + const float * shard_input = embed.data(); for (size_t si = 0; si < shards_.size(); ++si) { auto & shard = shards_[si]; const bool is_last = (si == shards_.size() - 1); @@ -479,13 +473,14 @@ bool DeepSeek4LayerSplitAdapter::run_forward( 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); return false; } if (timing) add_step_tel(tel_acc, step_tel); + shard_input = hc_state_.data(); } cur_pos_ = base_pos + n_tokens; @@ -575,8 +570,9 @@ bool DeepSeek4LayerSplitAdapter::run_mixed_forward( bool DeepSeek4LayerSplitAdapter::prefill( const std::vector & prompt, int base_pos, - int & last_tok) { - const int chunk_size = cfg_.chunk > 0 ? cfg_.chunk : 512; + int & last_tok, + bool need_logits) { + const int chunk_size = 1; const int n_prompt = (int)prompt.size(); for (int offset = 0; offset < n_prompt; offset += chunk_size) { @@ -585,7 +581,7 @@ bool DeepSeek4LayerSplitAdapter::prefill( prompt.begin() + chunk_end); std::vector * logits_out = - (chunk_end >= n_prompt) ? &prefill_last_logits_ : nullptr; + (need_logits && chunk_end >= n_prompt) ? &prefill_last_logits_ : nullptr; if (!run_forward(chunk, base_pos + offset, last_tok, logits_out)) { return false; diff --git a/server/src/deepseek4/deepseek4_layer_split_adapter.h b/server/src/deepseek4/deepseek4_layer_split_adapter.h index 67072d36a..6cb4102be 100644 --- a/server/src/deepseek4/deepseek4_layer_split_adapter.h +++ b/server/src/deepseek4/deepseek4_layer_split_adapter.h @@ -51,7 +51,8 @@ class DeepSeek4LayerSplitAdapter : public LayerSplitAdapter { void reset_request_state() override; int prefill_chunk_tokens() const override { return cfg_.chunk; } bool prefill(const std::vector & prompt, - int base_pos, int & last_tok) override; + int base_pos, int & last_tok, + bool need_logits = true) override; bool decode_ar(int last_tok, int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io) override; diff --git a/server/src/gemma4/gemma4_layer_split_adapter.cpp b/server/src/gemma4/gemma4_layer_split_adapter.cpp index 5f34c2ab7..673779307 100644 --- a/server/src/gemma4/gemma4_layer_split_adapter.cpp +++ b/server/src/gemma4/gemma4_layer_split_adapter.cpp @@ -1025,7 +1025,9 @@ bool Gemma4LayerSplitAdapter::run_mixed_forward( bool Gemma4LayerSplitAdapter::prefill(const std::vector & prompt, int base_pos, - int & last_tok) { + int & last_tok, + bool need_logits) { + (void)need_logits; const bool ok = use_mixed_target_split() ? run_mixed_forward(prompt, base_pos, last_tok, &prefill_last_logits_) : run_forward(prompt, base_pos, last_tok, &prefill_last_logits_); diff --git a/server/src/gemma4/gemma4_layer_split_adapter.h b/server/src/gemma4/gemma4_layer_split_adapter.h index 3e5331705..8e67559a2 100644 --- a/server/src/gemma4/gemma4_layer_split_adapter.h +++ b/server/src/gemma4/gemma4_layer_split_adapter.h @@ -58,7 +58,8 @@ class Gemma4LayerSplitAdapter : public LayerSplitAdapter { void reset_request_state() override; int prefill_chunk_tokens() const override { return cfg_.chunk > 0 ? cfg_.chunk : 0; } bool prefill(const std::vector & prompt, - int base_pos, int & last_tok) override; + int base_pos, int & last_tok, + bool need_logits = true) override; bool decode_ar(int last_tok, int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io) override; diff --git a/server/src/laguna/laguna_layer_split_adapter.cpp b/server/src/laguna/laguna_layer_split_adapter.cpp index 44580a13a..588f30ee0 100644 --- a/server/src/laguna/laguna_layer_split_adapter.cpp +++ b/server/src/laguna/laguna_layer_split_adapter.cpp @@ -840,7 +840,9 @@ bool LagunaLayerSplitAdapter::run_mixed_forward( bool LagunaLayerSplitAdapter::prefill(const std::vector & prompt, int base_pos, - int & last_tok) { + int & last_tok, + bool need_logits) { + (void)need_logits; const bool ok = use_mixed_target_split() ? run_mixed_forward(prompt, base_pos, last_tok, &prefill_last_logits_) : run_forward(prompt, base_pos, last_tok, &prefill_last_logits_); diff --git a/server/src/laguna/laguna_layer_split_adapter.h b/server/src/laguna/laguna_layer_split_adapter.h index a93e982d7..301edfcf5 100644 --- a/server/src/laguna/laguna_layer_split_adapter.h +++ b/server/src/laguna/laguna_layer_split_adapter.h @@ -57,7 +57,8 @@ class LagunaLayerSplitAdapter : public LayerSplitAdapter { void begin_request(const GenerateRequest & req) override; void reset_request_state() override; bool prefill(const std::vector & prompt, - int base_pos, int & last_tok) override; + int base_pos, int & last_tok, + bool need_logits = true) override; bool decode_ar(int last_tok, int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io) override; diff --git a/server/src/qwen35/qwen35_layer_split_adapter.cpp b/server/src/qwen35/qwen35_layer_split_adapter.cpp index 3c8320f1b..aea59ea5f 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.cpp +++ b/server/src/qwen35/qwen35_layer_split_adapter.cpp @@ -522,7 +522,9 @@ int Qwen35LayerSplitAdapter::prefill_chunk_tokens() const { } bool Qwen35LayerSplitAdapter::prefill(const std::vector & prompt, - int base_pos, int & last_tok) { + int base_pos, int & last_tok, + bool need_logits) { + (void)need_logits; if (prompt.empty()) return false; if (base_pos < 0 || base_pos + (int)prompt.size() > cfg_.device.max_ctx) { std::fprintf(stderr, diff --git a/server/src/qwen35/qwen35_layer_split_adapter.h b/server/src/qwen35/qwen35_layer_split_adapter.h index 6fe0b7765..be9b9a154 100644 --- a/server/src/qwen35/qwen35_layer_split_adapter.h +++ b/server/src/qwen35/qwen35_layer_split_adapter.h @@ -59,7 +59,8 @@ class Qwen35LayerSplitAdapter : public LayerSplitAdapter { void reset_request_state() override; int prefill_chunk_tokens() const override; bool prefill(const std::vector & prompt, - int base_pos, int & last_tok) override; + int base_pos, int & last_tok, + bool need_logits = true) override; bool decode_ar(int last_tok, int committed, int n_gen, std::vector & out_tokens, const DaemonIO & io) override; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index 58bd272f3..67dca5a84 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -27,6 +27,7 @@ #include "common/layer_split_backend.h" #include "common/layer_split_kvflash.h" #include "common/layer_split_utils.h" +#include "deepseek4/deepseek4_layer_split_adapter.h" #include "common/kvflash_pager.h" #include "placement/draft_residency.h" #include "common/gguf_bounds.h" @@ -38,6 +39,7 @@ #include #include +#include #include #include #include @@ -2073,6 +2075,7 @@ struct MockLayerSplitAdapter : LayerSplitAdapter { int current_last = -1; std::vector prefill_bases; std::vector prefill_sizes; + std::vector prefill_need_logits; int dflash_base = -1; int dflash_last = -1; std::vector emitted_tokens; @@ -2095,9 +2098,11 @@ struct MockLayerSplitAdapter : LayerSplitAdapter { } int prefill_chunk_tokens() const override { return prefill_chunk; } bool prefill(const std::vector & prompt, - int base_pos, int & last_tok) override { + int base_pos, int & last_tok, + bool need_logits = true) override { prefill_bases.push_back(base_pos); prefill_sizes.push_back((int)prompt.size()); + prefill_need_logits.push_back(need_logits); current_pos = base_pos + (int)prompt.size(); current_last = prompt.empty() ? current_last : prompt.back(); last_tok = current_last; @@ -2255,12 +2260,15 @@ static void test_layer_split_backend_sampling_capability_gate() { } static void test_layer_split_backend_chunks_prefill_by_adapter_limit() { + DeepSeek4LayerSplitAdapterConfig deepseek4_cfg; + TEST_ASSERT(deepseek4_cfg.chunk == 512); + auto * raw = new MockLayerSplitAdapter(); - raw->prefill_chunk = 3; + raw->prefill_chunk = 4; LayerSplitBackend backend{std::unique_ptr(raw)}; GenerateRequest req; - req.prompt = {1, 2, 3, 4, 5, 6, 7, 8}; + req.prompt = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; req.n_gen = 1; DaemonIO io; GenerateResult result = backend.generate(req, io); @@ -2268,12 +2276,20 @@ static void test_layer_split_backend_chunks_prefill_by_adapter_limit() { TEST_ASSERT(result.ok()); TEST_ASSERT(raw->prefill_bases.size() == 3); TEST_ASSERT(raw->prefill_sizes.size() == 3); + TEST_ASSERT(raw->prefill_need_logits.size() == 3); TEST_ASSERT(raw->prefill_bases[0] == 0); - TEST_ASSERT(raw->prefill_sizes[0] == 3); - TEST_ASSERT(raw->prefill_bases[1] == 3); - TEST_ASSERT(raw->prefill_sizes[1] == 3); - TEST_ASSERT(raw->prefill_bases[2] == 6); + TEST_ASSERT(raw->prefill_sizes[0] == 4); + TEST_ASSERT(!raw->prefill_need_logits[0]); + TEST_ASSERT(raw->prefill_bases[1] == 4); + TEST_ASSERT(raw->prefill_sizes[1] == 4); + TEST_ASSERT(!raw->prefill_need_logits[1]); + TEST_ASSERT(raw->prefill_bases[2] == 8); TEST_ASSERT(raw->prefill_sizes[2] == 2); + TEST_ASSERT(raw->prefill_need_logits[2]); + TEST_ASSERT(std::find(raw->prefill_sizes.begin(), raw->prefill_sizes.end(), 1) == + raw->prefill_sizes.end()); + TEST_ASSERT(std::any_of(raw->prefill_sizes.begin(), raw->prefill_sizes.end(), + [](int n) { return n > 1; })); } static void test_layer_split_compress_nopark_uses_default_drafter_path() { diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index bb32aeed2..d507a6e43 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -1309,6 +1309,104 @@ static void test_reference_hc_pre(const std::vector & hc_state, for (int i = 0; i < n_hc * n_hc; ++i) comb[(size_t) i] = split[2 * n_hc + i]; } +static void test_reference_hc_pre_batch(const std::vector & hc_state, + const std::vector & fn_f16, + const std::vector & scale, + const std::vector & base, + int n_tokens, + int n_embd, + int n_hc, + int sinkhorn_iters, + float hc_eps, + std::vector & working, + std::vector & post, + std::vector & comb) { + const int hc_dim = n_embd * n_hc; + working.assign((size_t)n_tokens * (size_t)n_embd, 0.0f); + post.assign((size_t)n_tokens * (size_t)n_hc, 0.0f); + comb.assign((size_t)n_tokens * (size_t)n_hc * (size_t)n_hc, 0.0f); + for (int t = 0; t < n_tokens; ++t) { + std::vector token_hc( + hc_state.begin() + (ptrdiff_t)t * hc_dim, + hc_state.begin() + (ptrdiff_t)(t + 1) * hc_dim); + std::vector token_working; + std::vector token_post; + std::vector token_comb; + test_reference_hc_pre(token_hc, fn_f16, scale, base, n_embd, n_hc, + sinkhorn_iters, hc_eps, + token_working, token_post, token_comb); + std::copy(token_working.begin(), token_working.end(), + working.begin() + (ptrdiff_t)t * n_embd); + std::copy(token_post.begin(), token_post.end(), + post.begin() + (ptrdiff_t)t * n_hc); + std::copy(token_comb.begin(), token_comb.end(), + comb.begin() + (ptrdiff_t)t * n_hc * n_hc); + } +} + +static void test_reference_hc_post_batch(const std::vector & residual_hc, + const std::vector & block_out, + const std::vector & post, + const std::vector & comb, + int n_tokens, + int n_embd, + int n_hc, + std::vector & out_hc) { + const int hc_dim = n_embd * n_hc; + out_hc.assign((size_t)n_tokens * (size_t)hc_dim, 0.0f); + for (int t = 0; t < n_tokens; ++t) { + for (int dst = 0; dst < n_hc; ++dst) { + for (int d = 0; d < n_embd; ++d) { + float acc = block_out[(size_t)t * n_embd + d] * + post[(size_t)t * n_hc + dst]; + for (int src = 0; src < n_hc; ++src) { + acc += comb[(size_t)t * n_hc * n_hc + dst + src * n_hc] * + residual_hc[(size_t)t * hc_dim + (size_t)src * n_embd + d]; + } + out_hc[(size_t)t * hc_dim + (size_t)dst * n_embd + d] = acc; + } + } + } +} + +static void test_reference_hc_output_batch(const std::vector & hc_state, + const std::vector & fn_f16, + const std::vector & scale, + const std::vector & base, + int n_tokens, + int n_embd, + int n_hc, + float hc_eps, + std::vector & final_embd) { + const int hc_dim = n_embd * n_hc; + final_embd.assign((size_t)n_tokens * (size_t)n_embd, 0.0f); + for (int t = 0; t < n_tokens; ++t) { + const float * token_hc = hc_state.data() + (size_t)t * hc_dim; + float sumsq = 0.0f; + for (int i = 0; i < hc_dim; ++i) { + sumsq += token_hc[i] * token_hc[i]; + } + const float inv_rms = 1.0f / std::sqrt(sumsq / (float)hc_dim + hc_eps); + std::vector pre((size_t)n_hc, 0.0f); + for (int h = 0; h < n_hc; ++h) { + float acc = 0.0f; + for (int c = 0; c < hc_dim; ++c) { + acc += ggml_fp16_to_fp32(fn_f16[(size_t)h * hc_dim + c]) * + token_hc[c] * inv_rms; + } + pre[(size_t)h] = 1.0f / (1.0f + std::exp(-(acc * scale[0] + base[(size_t)h]))) + + 1.0e-6f; + } + for (int d = 0; d < n_embd; ++d) { + float acc = 0.0f; + for (int h = 0; h < n_hc; ++h) { + acc += pre[(size_t)h] * token_hc[(size_t)h * n_embd + d]; + } + final_embd[(size_t)t * n_embd + d] = acc; + } + } +} + static void test_hc_pre_kernel_gpu() { std::fprintf(stderr, " test_hc_pre_kernel_gpu ..."); ggml_backend_t backend = ggml_backend_cuda_init(0); @@ -1541,6 +1639,163 @@ static void test_hc_pre_kernel_gpu() { ggml_free(ctx); ggml_backend_free(backend); } + +static void test_hc_batch_kernels_gpu() { + std::fprintf(stderr, " test_hc_batch_kernels_gpu ..."); + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, " skipped (no GPU backend)\n"); + return; + } + + constexpr int n_tokens = 3; + constexpr int n_embd = 32; + constexpr int n_hc = 4; + constexpr int sinkhorn_iters = 5; + constexpr float hc_eps = 1.0e-6f; + constexpr int mix_dim = 2 * n_hc + n_hc * n_hc; + const int hc_dim = n_embd * n_hc; + + std::mt19937 rng(321); + std::uniform_real_distribution dist(-0.25f, 0.25f); + std::vector hc_state((size_t)n_tokens * (size_t)hc_dim); + std::vector block_out((size_t)n_tokens * (size_t)n_embd); + std::vector fn((size_t)mix_dim * (size_t)hc_dim); + std::vector fn_f16(fn.size()); + std::vector output_fn((size_t)n_hc * (size_t)hc_dim); + std::vector output_fn_f16(output_fn.size()); + std::vector scale((size_t)mix_dim, 0.0f); + std::vector base((size_t)mix_dim); + for (float & v : hc_state) v = dist(rng); + for (float & v : block_out) v = dist(rng); + for (float & v : fn) v = dist(rng); + for (float & v : output_fn) v = dist(rng); + for (size_t i = 0; i < fn.size(); ++i) fn_f16[i] = ggml_fp32_to_fp16(fn[i]); + for (size_t i = 0; i < output_fn.size(); ++i) output_fn_f16[i] = ggml_fp32_to_fp16(output_fn[i]); + scale[0] = 0.9f; + scale[1] = 1.15f; + scale[2] = 0.8f; + for (float & v : base) v = 0.1f * dist(rng); + + std::vector ref_working; + std::vector ref_post; + std::vector ref_comb; + std::vector ref_next_hc; + std::vector ref_final_embd; + test_reference_hc_pre_batch(hc_state, fn_f16, scale, base, n_tokens, n_embd, n_hc, + sinkhorn_iters, hc_eps, + ref_working, ref_post, ref_comb); + test_reference_hc_post_batch(hc_state, block_out, ref_post, ref_comb, + n_tokens, n_embd, n_hc, ref_next_hc); + test_reference_hc_output_batch(hc_state, output_fn_f16, scale, base, + n_tokens, n_embd, n_hc, hc_eps, ref_final_embd); + + ggml_context * ctx = make_test_context(1u << 20); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + ggml_backend_free(backend); + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * state_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hc_dim, n_tokens); + ggml_tensor * block_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_tensor * fn_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hc_dim, mix_dim); + ggml_tensor * output_fn_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hc_dim, n_hc); + ggml_tensor * scale_t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mix_dim); + ggml_tensor * base_t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mix_dim); + ggml_tensor * working_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_tensor * post_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_hc, n_tokens); + ggml_tensor * comb_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_hc * n_hc, n_tokens); + ggml_tensor * next_hc_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hc_dim, n_tokens); + ggml_tensor * final_embd_t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + TEST_ASSERT_MSG(buf != nullptr, "ggml backend buffer alloc failed"); + if (!buf) { + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_backend_tensor_set(state_t, hc_state.data(), 0, hc_state.size() * sizeof(float)); + ggml_backend_tensor_set(block_t, block_out.data(), 0, block_out.size() * sizeof(float)); + ggml_backend_tensor_set(fn_t, fn_f16.data(), 0, fn_f16.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(output_fn_t, output_fn_f16.data(), 0, output_fn_f16.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(scale_t, scale.data(), 0, scale.size() * sizeof(float)); + ggml_backend_tensor_set(base_t, base.data(), 0, base.size() * sizeof(float)); + + bool ok = deepseek4_cuda_hc_pre_batch_device(state_t->data, + fn_t->data, + scale_t->data, + base_t->data, + n_tokens, + n_embd, + n_hc, + sinkhorn_iters, + hc_eps, + working_t->data, + post_t->data, + comb_t->data); + TEST_ASSERT_MSG(ok, "batched HC-pre kernel call failed"); + + ok = ok && deepseek4_cuda_hc_post_batch_device(state_t->data, + block_t->data, + post_t->data, + comb_t->data, + n_tokens, + n_embd, + n_hc, + next_hc_t->data); + TEST_ASSERT_MSG(ok, "batched HC-post kernel call failed"); + + ok = ok && deepseek4_cuda_hc_output_batch_device(state_t->data, + output_fn_t->data, + scale_t->data, + base_t->data, + n_tokens, + n_embd, + n_hc, + hc_eps, + final_embd_t->data); + TEST_ASSERT_MSG(ok, "batched HC-output kernel call failed"); + + std::vector gpu_working(ref_working.size()); + std::vector gpu_post(ref_post.size()); + std::vector gpu_comb(ref_comb.size()); + std::vector gpu_next_hc(ref_next_hc.size()); + std::vector gpu_final_embd(ref_final_embd.size()); + if (ok) { + ggml_backend_tensor_get(working_t, gpu_working.data(), 0, gpu_working.size() * sizeof(float)); + ggml_backend_tensor_get(post_t, gpu_post.data(), 0, gpu_post.size() * sizeof(float)); + ggml_backend_tensor_get(comb_t, gpu_comb.data(), 0, gpu_comb.size() * sizeof(float)); + ggml_backend_tensor_get(next_hc_t, gpu_next_hc.data(), 0, gpu_next_hc.size() * sizeof(float)); + ggml_backend_tensor_get(final_embd_t, gpu_final_embd.data(), 0, gpu_final_embd.size() * sizeof(float)); + } + + constexpr float atol = 3.0e-4f; + constexpr float rtol = 3.0e-4f; + for (size_t i = 0; i < ref_working.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(gpu_working[i], ref_working[i], atol, rtol), "batched working mismatch"); + } + for (size_t i = 0; i < ref_post.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(gpu_post[i], ref_post[i], atol, rtol), "batched post mismatch"); + } + for (size_t i = 0; i < ref_comb.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(gpu_comb[i], ref_comb[i], atol, rtol), "batched comb mismatch"); + } + for (size_t i = 0; i < ref_next_hc.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(gpu_next_hc[i], ref_next_hc[i], 5.0e-4f, 5.0e-4f), "batched hc-post output mismatch"); + } + for (size_t i = 0; i < ref_final_embd.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(gpu_final_embd[i], ref_final_embd[i], atol, rtol), "batched hc-output mismatch"); + } + + ggml_backend_buffer_free(buf); + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} #endif int main() { @@ -1572,6 +1827,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_batch_kernels_gpu(); #endif ggml_backend_free(backend); From 74f9d96f4cf59e2590db0f9ea222d545491114b6 Mon Sep 17 00:00:00 2001 From: Howard Su Date: Mon, 13 Jul 2026 08:52:03 +0800 Subject: [PATCH 2/4] fix(dflash): align draft top-K guard with CUDA-only macro CMake now defines DFLASH27B_HAVE_DRAFT_TOPK_CUDA (CUDA backend only) and no longer defines DFLASH27B_HAVE_DRAFT_TOPK on any backend. Update the consumers in qwen35_dflash_target.cpp and test_dflash.cpp to test the new name so the GPU draft top-K path is actually compiled in on CUDA instead of silently falling back to the CPU heap extract. Fix the stale header comment to match the CUDA-only build. --- server/src/common/geometric_draft_topk_cuda.h | 6 ++---- server/src/qwen35/qwen35_dflash_target.cpp | 2 +- server/test/test_dflash.cpp | 2 +- 3 files changed, 4 insertions(+), 6 deletions(-) diff --git a/server/src/common/geometric_draft_topk_cuda.h b/server/src/common/geometric_draft_topk_cuda.h index b926dbefc..e83c9796e 100644 --- a/server/src/common/geometric_draft_topk_cuda.h +++ b/server/src/common/geometric_draft_topk_cuda.h @@ -14,10 +14,8 @@ // // Returns false (caller must fall back to the CPU path) when the GPU runtime is // unavailable, the pointer is not device memory, K is out of range, or any -// device call fails. Compiled on both CUDA and HIP/ROCm builds — this same .cu -// is compiled directly with LANGUAGE HIP on ROCm, the cuda_runtime.h spellings -// mapped by the hip_compat shim; guarded by DFLASH27B_HAVE_DRAFT_TOPK. See -// CMakeLists.txt. +// device call fails. Compiled on the CUDA backend only, where this .cu is added +// to dflash_common; guarded by DFLASH27B_HAVE_DRAFT_TOPK_CUDA. See CMakeLists.txt. #pragma once diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index dc30ff7c4..4185b5d40 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -692,7 +692,7 @@ bool Qwen35DFlashTarget::project_hidden_to_topk( top_log_probs.assign((size_t)n_tokens * K, 0.0f); top_token_ids.assign((size_t)n_tokens * K, 0); -#ifdef DFLASH27B_HAVE_DRAFT_TOPK +#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA // GPU path: top-K + logsumexp directly on the logits device buffer, skipping // the vocab×n_tokens D2H and the CPU heap extract. Falls back to the CPU path // on any failure. Escape hatch: DFLASH_GPU_DRAFT_TOPK=0. diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 7f5473995..407d23af1 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -3284,7 +3284,7 @@ int main(int argc, char ** argv) { } else { // DDTree K>1: need real log-probs for best-first tree scoring. bool topk_done = false; -#ifdef DFLASH27B_HAVE_DRAFT_TOPK +#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA // GPU path: top-K + logsumexp on the draft logits device buffer // (positions 1..q_len-1), no full-vocab D2H. Escape: DFLASH_GPU_DRAFT_TOPK=0. static const bool kGpuDraftTopk = [](){ From 975cf7d42768872a5cb0b73dab895fe68f22c8d3 Mon Sep 17 00:00:00 2001 From: Howard Su Date: Mon, 13 Jul 2026 08:52:20 +0800 Subject: [PATCH 3/4] fix(deepseek4): restore cfg-driven batched prefill in layer-split adapter The layer-split adapter's prefill() hardcoded chunk_size = 1, forcing single-token prefill and defeating the batched layer-range path. This contradicts both the intended behavior (prefill_chunk_tokens() reports cfg_.chunk) and the restore commit's stated goal of cfg-driven batched prefill. Restore chunk_size to cfg_.chunk (default 512). --- server/src/deepseek4/deepseek4_layer_split_adapter.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp index e3158c19a..749eb3eb9 100644 --- a/server/src/deepseek4/deepseek4_layer_split_adapter.cpp +++ b/server/src/deepseek4/deepseek4_layer_split_adapter.cpp @@ -572,7 +572,7 @@ bool DeepSeek4LayerSplitAdapter::prefill( int base_pos, int & last_tok, bool need_logits) { - const int chunk_size = 1; + const int chunk_size = cfg_.chunk > 0 ? cfg_.chunk : 512; const int n_prompt = (int)prompt.size(); for (int offset = 0; offset < n_prompt; offset += chunk_size) { From 0b371f7e0ad932fc8d6cdcc62eb7c2a5d014de60 Mon Sep 17 00:00:00 2001 From: Howard Su Date: Mon, 13 Jul 2026 09:03:35 +0800 Subject: [PATCH 4/4] fix(deepseek4): run per-token compressor updates in batched prefill The learned KV compressor is a sequential accumulator: every token feeds its ratio-window state and each window boundary flushes one pooled comp_kv row (index_comp_kv too for ratio 4). Batched prefill only ran the compressor for the chunk's last token via cur_last, yet advanced n_comp to next_pos/ratio, so every intermediate window's compressed row was left stale/uninitialized while subsequent decode attended over it. Step the compressor for every token in the chunk (matching the single-token decode path); the per-token side effects serialize by graph insertion order, like the raw SWA ring write loop. Scale the dynamic attention graph node budget and metadata arena with n_tokens to hold the per-token compressor chain. --- server/src/deepseek4/deepseek4_graph.cpp | 108 +++++++++++++++-------- 1 file changed, 69 insertions(+), 39 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 60bc7f3d7..f6a7d001f 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -1020,44 +1020,65 @@ static ggml_tensor * build_mla_attention( } // ── Learned compression update ────────────────────────────────── - ggml_tensor * cur_last = ggml_view_2d( - ctx, cur, n_embd, 1, cur->nb[1], (size_t)(n_tokens - 1) * cur->nb[1]); - ggml_tensor * qr_last = ggml_view_2d( - ctx, qr, n_lora_q, 1, qr->nb[1], (size_t)(n_tokens - 1) * qr->nb[1]); - if (ratio > 0 && L.attn_compressor_kv) { - build_compressor_step(ctx, gf, cur_last, - L.attn_compressor_ape, - L.attn_compressor_kv, - L.attn_compressor_gate, - L.attn_compressor_norm, - lc.attn_compressor, - lc.comp_kv, - ratio, - head_dim, - token_pos, - w.n_rot, - w.rms_eps, - w.compress_rope_freq_base, - w.rope_scale_factor, - w.rope_yarn_beta_fast, - w.rope_yarn_beta_slow, - (int)w.rope_orig_ctx, - cached_inputs ? cached_inputs->attn_ape_row : nullptr, - cached_inputs ? cached_inputs->attn_state_rows : nullptr, - cached_inputs ? cached_inputs->attn_comp_rows : nullptr, - cached_inputs ? cached_inputs->attn_comp_pos : nullptr, - i64_array_inputs, - i32_array_inputs); - } - + // The learned compressor is a sequential accumulator: every token feeds + // its ratio-window state and each window boundary flushes exactly one + // pooled comp_kv row (comp_row = pos / ratio). Batched prefill must + // therefore step the compressor for EVERY token in the chunk, exactly as + // the single-token decode path does one token at a time. Running it only + // for the last token leaves every intermediate window's comp_kv (and, for + // ratio 4, index_comp_kv) row unwritten while n_comp is still advanced to + // next_pos / ratio, so later decode attends over stale/uninitialized + // compressed rows. The per-token side effects are serialized by insertion + // order into the graph, matching the raw SWA ring write loop above. + const bool batched_prefill = (n_tokens > 1 && !cached_inputs); + const int comp_first = batched_prefill ? 0 : (n_tokens - 1); + for (int ti = comp_first; ti < n_tokens; ++ti) { + ggml_tensor * cur_col = ggml_view_2d( + ctx, cur, n_embd, 1, cur->nb[1], (size_t)ti * cur->nb[1]); + const int tok_pos = kv_start + ti; + if (ratio > 0 && L.attn_compressor_kv) { + build_compressor_step(ctx, gf, cur_col, + L.attn_compressor_ape, + L.attn_compressor_kv, + L.attn_compressor_gate, + L.attn_compressor_norm, + lc.attn_compressor, + lc.comp_kv, + ratio, + head_dim, + tok_pos, + w.n_rot, + w.rms_eps, + w.compress_rope_freq_base, + w.rope_scale_factor, + w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, + (int)w.rope_orig_ctx, + cached_inputs ? cached_inputs->attn_ape_row : nullptr, + cached_inputs ? cached_inputs->attn_state_rows : nullptr, + cached_inputs ? cached_inputs->attn_comp_rows : nullptr, + cached_inputs ? cached_inputs->attn_comp_pos : nullptr, + i64_array_inputs, + i32_array_inputs); + } + if (ratio == 4 && L.indexer_compressor_kv) { + build_indexer_compressor_step(ctx, gf, cur_col, w, L, lc, tok_pos, + cached_inputs ? cached_inputs->index_ape_row : nullptr, + cached_inputs ? cached_inputs->index_state_rows : nullptr, + cached_inputs ? cached_inputs->index_comp_rows : nullptr, + cached_inputs ? cached_inputs->index_comp_pos : nullptr, + i64_array_inputs, + i32_array_inputs); + } + } + + // Indexer score uses the final token's query as the sparse-selection + // signal (result currently unused downstream). if (ratio == 4 && L.indexer_compressor_kv) { - build_indexer_compressor_step(ctx, gf, cur_last, w, L, lc, token_pos, - cached_inputs ? cached_inputs->index_ape_row : nullptr, - cached_inputs ? cached_inputs->index_state_rows : nullptr, - cached_inputs ? cached_inputs->index_comp_rows : nullptr, - cached_inputs ? cached_inputs->index_comp_pos : nullptr, - i64_array_inputs, - i32_array_inputs); + ggml_tensor * qr_last = ggml_view_2d( + ctx, qr, n_lora_q, 1, qr->nb[1], (size_t)(n_tokens - 1) * qr->nb[1]); + ggml_tensor * cur_last = ggml_view_2d( + ctx, cur, n_embd, 1, cur->nb[1], (size_t)(n_tokens - 1) * cur->nb[1]); (void)build_indexer_score(ctx, qr_last, cur_last, w, L, lc, token_pos, i32_inputs); } @@ -3308,7 +3329,12 @@ bool deepseek4_step_layer_range( } } else { const auto attn_build_t0 = Ds4TimingClock::now(); - const size_t ctx_size = 48 * 1024 * 1024; + // Batched prefill now emits a per-token compressor chain (see + // build_mla_attention), so both the node budget and the tensor + // metadata arena must scale with the chunk size. + const size_t ctx_size = (n_tokens > 1) + ? (size_t)128 * 1024 * 1024 + : (size_t)48 * 1024 * 1024; ggml_init_params params{}; params.mem_size = ctx_size; params.mem_buffer = nullptr; @@ -3322,7 +3348,11 @@ bool deepseek4_step_layer_range( std::vector i32_array_inputs; std::vector i64_array_inputs; std::vector f16_array_inputs; - const size_t graph_size = n_tokens > 1 ? 32768 : 2048; + size_t graph_size = 2048; + if (n_tokens > 1) { + graph_size = (size_t)n_tokens * 128 + 8192; + if (graph_size < 32768) graph_size = 32768; + } gf = ggml_new_graph_custom(ctx, graph_size, false); ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps);