From 1c2d253db58bcce709bd5f272bca9e4eae7ceb8c Mon Sep 17 00:00:00 2001 From: weicj Date: Tue, 7 Jul 2026 10:24:38 +0800 Subject: [PATCH 1/7] feat(dflash): add direct DeepSeek4 pure-UMA runtime path --- server/src/common/backend_factory.cpp | 49 ++++----- server/src/deepseek4/deepseek4_backend.cpp | 14 ++- server/src/deepseek4/deepseek4_graph.cpp | 112 +++++++++++++++++++-- server/src/deepseek4/deepseek4_internal.h | 5 + 4 files changed, 139 insertions(+), 41 deletions(-) diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 1383c2166..08ca33248 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -215,44 +215,33 @@ std::unique_ptr create_backend(const BackendArgs & args) { return backend; } else if (arch == "deepseek4") { - const PlacementBackend target_backend = - args.device.backend == PlacementBackend::Auto - ? 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; - - auto backend = std::make_unique(cfg); + if (args.device.is_layer_split() || args.remote_target_shard.enabled()) { + DeepSeek4LayerSplitAdapterConfig cfg; + cfg.target_path = args.model_path; + cfg.device = args.device; + cfg.remote_target_shard = args.remote_target_shard; + cfg.chunk = args.chunk; + + auto adapter = std::make_unique(cfg); + auto backend = std::make_unique(std::move(adapter)); if (!backend->init()) { - std::fprintf(stderr, "[backend_factory] DeepSeek4Backend init failed\n"); + std::fprintf(stderr, + "[backend_factory] LayerSplitBackend(deepseek4) 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. - DeepSeek4LayerSplitAdapterConfig cfg; - cfg.target_path = args.model_path; - cfg.device = args.device; - cfg.remote_target_shard = args.remote_target_shard; - cfg.chunk = args.chunk; + DeepSeek4BackendConfig dcfg; + dcfg.model_path = args.model_path; + dcfg.device = args.device; + dcfg.stream_fd = args.stream_fd; + dcfg.chunk = args.chunk; + dcfg.max_ctx = args.device.max_ctx; - auto adapter = std::make_unique(cfg); - auto backend = std::make_unique(std::move(adapter)); + auto backend = std::make_unique(dcfg); if (!backend->init()) { - std::fprintf(stderr, "[backend_factory] LayerSplitBackend(deepseek4) init failed\n"); + std::fprintf(stderr, "[backend_factory] DeepSeek4Backend init failed\n"); return nullptr; } return backend; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 9b4616da1..ddcdfc40f 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -40,6 +40,11 @@ static double gib(uint64_t bytes) { static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTelemetry & src) { dst.total_us += src.total_us; dst.embed_us += src.embed_us; + dst.full_graph_build_us += src.full_graph_build_us; + dst.full_graph_alloc_us += src.full_graph_alloc_us; + dst.full_graph_set_us += src.full_graph_set_us; + dst.full_graph_compute_us += src.full_graph_compute_us; + dst.full_graph_read_us += src.full_graph_read_us; dst.hc_pre_attn_us += src.hc_pre_attn_us; dst.hc_pre_build_us += src.hc_pre_build_us; dst.hc_pre_input_us += src.hc_pre_input_us; @@ -85,7 +90,9 @@ static void log_step_tel(const char * phase, const double tok_s = wall_s > 0.0 ? (double)tokens / wall_s : 0.0; std::fprintf(stderr, "[deepseek4-timing] %s tokens=%d steps=%d wall=%.3fs %.2f tok/s " - "step=%.1fms embed=%.1fms attn_build=%.1fms attn_compute=%.1fms attn_read=%.1fms " + "step=%.1fms embed=%.1fms full_build=%.1fms full_alloc=%.1fms full_set=%.1fms " + "full_compute=%.1fms full_read=%.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 " @@ -94,7 +101,10 @@ static void log_step_tel(const char * phase, "hc_post=%.1fms output=%.1fms sample=%.1fms emit=%.1fms " "hot_sel=%d cold_sel=%d\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.total_us), ms(t.embed_us), + ms(t.full_graph_build_us), ms(t.full_graph_alloc_us), ms(t.full_graph_set_us), + ms(t.full_graph_compute_us), ms(t.full_graph_read_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), diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 06fbe8a98..90be91d85 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -176,6 +176,21 @@ struct DeepSeek4CachedDecodeOutputGraph { } }; +struct DeepSeek4LegacyFullStepCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + StepGraph sg; + std::vector meta_arena; + + void free() { + step_graph_destroy(sg); + meta_arena.clear(); + meta_arena.shrink_to_fit(); + owner_ctx = nullptr; + backend = nullptr; + } +}; + struct DeepSeek4AttentionGraphInputs { ggml_tensor * rope_pos = nullptr; ggml_tensor * neg_pos = nullptr; @@ -2669,6 +2684,7 @@ bool deepseek4_step( MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats) { + const auto step_t0 = Ds4TimingClock::now(); if (w.moe_hybrid && moe_hybrid != nullptr) { return deepseek4_step_hybrid(backend, w, cache, *moe_hybrid, @@ -2681,20 +2697,67 @@ bool deepseek4_step( // Create compute graph context — need large budget for MoE layers const size_t ctx_size = ggml_tensor_overhead() * 65536 + 16 * 1024 * 1024; + const bool reuse_full_step_decode = + n_tokens == 1 && + ds4_backend_is_gpu(backend) && + !ds4_env_flag("DFLASH_DS4_DISABLE_FULL_STEP_DECODE_REUSE"); + static thread_local DeepSeek4LegacyFullStepCache full_step_cache; + StepGraph * cached_sg = nullptr; + if (reuse_full_step_decode) { + if (full_step_cache.owner_ctx != w.ctx || full_step_cache.backend != backend) { + full_step_cache.free(); + full_step_cache.owner_ctx = w.ctx; + full_step_cache.backend = backend; + } else { + step_graph_free(full_step_cache.sg); + } + cached_sg = &full_step_cache.sg; + if (full_step_cache.meta_arena.size() < ctx_size) { + full_step_cache.meta_arena.resize(ctx_size); + } + } + ggml_init_params params{}; - params.mem_size = ctx_size; - params.mem_buffer = nullptr; + params.mem_size = cached_sg ? full_step_cache.meta_arena.size() : ctx_size; + params.mem_buffer = cached_sg ? full_step_cache.meta_arena.data() : nullptr; params.no_alloc = true; + const auto full_build_t0 = Ds4TimingClock::now(); ggml_context * ctx = ggml_init(params); if (!ctx) return false; + if (cached_sg) { + cached_sg->ctx = ctx; + } + + ggml_gallocr_t alloc = nullptr; + bool owns_alloc = false; + auto release_full_step = [&]() { + if (cached_sg) { + step_graph_free(*cached_sg); + return; + } + if (alloc && owns_alloc) { + ggml_gallocr_free(alloc); + alloc = nullptr; + } + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } + }; // Input embeddings ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); ggml_set_name(inp, "inp_embed"); ggml_set_input(inp); + if (cached_sg) { + cached_sg->inp_embed = inp; + } ggml_tensor * cur = inp; ggml_cgraph * gf = ggml_new_graph_custom(ctx, 32768, false); + if (cached_sg) { + cached_sg->gf = gf; + } std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; @@ -2743,20 +2806,39 @@ bool deepseek4_step( ggml_tensor * logits = ggml_mul_mat(ctx, w.output, cur); ggml_set_name(logits, "logits"); ggml_set_output(logits); + if (cached_sg) { + cached_sg->logits = logits; + } // ── Build and run graph ───────────────────────────────────────────── ggml_build_forward_expand(gf, logits); + if (telemetry) { + telemetry->full_graph_build_us += ds4_elapsed_us(full_build_t0, Ds4TimingClock::now()); + } // Allocate - ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (cached_sg) { + if (!cached_sg->alloc) { + cached_sg->alloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + alloc = cached_sg->alloc; + } else { + alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + owns_alloc = true; + } + const auto full_alloc_t0 = Ds4TimingClock::now(); if (!ggml_gallocr_alloc_graph(alloc, gf)) { std::fprintf(stderr, "[deepseek4] graph allocation failed\n"); - ggml_gallocr_free(alloc); - ggml_free(ctx); + release_full_step(); return false; } + if (telemetry) { + telemetry->full_graph_alloc_us += ds4_elapsed_us(full_alloc_t0, Ds4TimingClock::now()); + } // Set input data + const auto full_set_t0 = Ds4TimingClock::now(); ggml_backend_tensor_set(inp, embed, 0, n_embd * n_tokens * sizeof(float)); for (const DeepSeek4I32InputBinding & binding : i32_inputs) { ggml_backend_tensor_set(binding.tensor, &binding.value, 0, sizeof(binding.value)); @@ -2769,23 +2851,32 @@ bool deepseek4_step( ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, sizeof(int64_t) * binding.values.size()); } + if (telemetry) { + telemetry->full_graph_set_us += ds4_elapsed_us(full_set_t0, Ds4TimingClock::now()); + } // Compute + const auto full_compute_t0 = Ds4TimingClock::now(); if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] graph compute failed\n"); - ggml_gallocr_free(alloc); - ggml_free(ctx); + release_full_step(); return false; } + if (telemetry) { + telemetry->full_graph_compute_us += ds4_elapsed_us(full_compute_t0, Ds4TimingClock::now()); + } // Read logits (only last token for generation) + const auto full_read_t0 = Ds4TimingClock::now(); out_logits.resize(w.n_vocab); const size_t logits_offset = (size_t)(n_tokens - 1) * w.n_vocab * sizeof(float); ggml_backend_tensor_get(logits, out_logits.data(), logits_offset, w.n_vocab * sizeof(float)); + if (telemetry) { + telemetry->full_graph_read_us += ds4_elapsed_us(full_read_t0, Ds4TimingClock::now()); + } - ggml_gallocr_free(alloc); - ggml_free(ctx); + release_full_step(); const int next_pos = kv_start + n_tokens; for (int il = 0; il < n_layer; ++il) { @@ -2801,6 +2892,9 @@ bool deepseek4_step( } cache.cur_pos = next_pos; + if (telemetry) { + telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + } return true; } diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 31c086394..d2a12b7d3 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -35,6 +35,11 @@ class MoeHybridStreamEngine; struct DeepSeek4StepTelemetry { uint64_t total_us = 0; uint64_t embed_us = 0; + uint64_t full_graph_build_us = 0; + uint64_t full_graph_alloc_us = 0; + uint64_t full_graph_set_us = 0; + uint64_t full_graph_compute_us = 0; + uint64_t full_graph_read_us = 0; uint64_t hc_pre_attn_us = 0; uint64_t hc_pre_build_us = 0; uint64_t hc_pre_input_us = 0; From 7405cb044f8595d2d774f26f16777ad50ff6d5d2 Mon Sep 17 00:00:00 2001 From: weicj Date: Thu, 9 Jul 2026 01:44:38 +0800 Subject: [PATCH 2/7] fix(dflash): stabilize DeepSeek4 pure UMA runtime --- server/CMakeLists.txt | 199 +- server/src/deepseek4/deepseek4_backend.cpp | 1609 +++++++++- server/src/deepseek4/deepseek4_backend.h | 41 +- server/src/deepseek4/deepseek4_graph.cpp | 3120 +++++++++++++++----- server/src/deepseek4/deepseek4_internal.h | 177 +- 5 files changed, 4269 insertions(+), 877 deletions(-) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index b693e1041..d952461f8 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -258,6 +258,7 @@ add_library(dflash_common STATIC src/deepseek4/deepseek4_daemon.cpp src/deepseek4/deepseek4_layer_split_adapter.cpp src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp + src/deepseek4/deepseek4_expert_ipc_daemon.cpp src/flashprefill_q8.cpp src/kv_cache.cpp src/kv_quant.cpp @@ -271,11 +272,11 @@ add_library(dflash_common STATIC src/laguna/laguna_dflash_target.cpp src/common/backend_ipc.cpp src/common/domino_head.cpp - src/common/dspark_head.cpp src/common/target_shard_ipc.cpp src/common/target_shard_ipc_daemon.cpp src/common/dflash_feature_ring.cpp src/common/dflash_capture.cpp + src/common/expert_ipc.cpp src/common/dflash_draft_ipc.cpp src/common/dflash_draft_ipc_daemon.cpp src/common/pflash_drafter_ipc.cpp @@ -283,6 +284,13 @@ add_library(dflash_common STATIC src/common/dflash_spec_decode.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp + src/common/expert_split_plan.cpp + src/common/expert_split_state.cpp + src/common/expert_split_runtime.cpp + src/common/expert_split_compute_runtime.cpp + src/common/expert_split_materialization.cpp + src/common/expert_split_target_config.cpp + src/common/gguf_tensor_data.cpp src/qwen35/graph_builders.cpp src/qwen35moe/qwen35moe_ffn.cpp src/qwen35moe/qwen35moe_backend.cpp @@ -297,6 +305,7 @@ add_library(dflash_common STATIC src/common/moe_hybrid_stream.cpp src/common/moe_expert_compute.cpp src/common/moe_expert_compute_cpu.cpp + src/common/moe_expert_compute_local_gpu.cpp src/common/moe_expert_compute_ipc.cpp src/common/moe_hybrid_swap_manager.cpp src/common/moe_routing_collector.cpp @@ -315,6 +324,7 @@ add_library(dflash_common STATIC src/common/daemon_loop.cpp src/common/gguf_inspect.cpp src/common/backend_factory.cpp + src/common/host_memory_mode.cpp src/placement/placement_config.cpp src/common/layer_split_utils.cpp src/common/ddtree.cpp @@ -379,7 +389,9 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") target_compile_definitions(dflash_common PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) # hip_compat shim is needed by ALL dflash_common sources (peer_access.cpp, # dflash_feature_ring.cpp, flashprefill.cpp), not just the SM80_EQUIV path. - target_include_directories(dflash_common PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) + target_include_directories(dflash_common BEFORE PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src/hip_compat + ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) endif() if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") @@ -417,8 +429,9 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip") src/flashprefill.cpp) set_source_files_properties(src/flashprefill_kernels.hip.cu PROPERTIES LANGUAGE HIP) - target_include_directories(dflash_common PRIVATE + target_include_directories(dflash_common BEFORE PRIVATE ${DFLASH27B_ROCWMMA_INCLUDE_DIR} + ${CMAKE_CURRENT_SOURCE_DIR}/src/hip_compat ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) target_compile_definitions(dflash_common PRIVATE DFLASH27B_HAVE_FLASHPREFILL=1 @@ -435,16 +448,6 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda") # 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. 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. - option(DFLASH_GPU_SAMPLER "Build the CUDA sample_logits path" ON) - if(DFLASH_GPU_SAMPLER) - target_sources(dflash_common PRIVATE src/common/geometric_sampler_cuda.cu) - # PUBLIC so the backends and test executables that call sample_logits() - # also compile their GPU dispatch branch. - target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_GPU_SAMPLER=1) - endif() # Multi-arch: scan all resolved arches and compile every applicable # flashprefill kernel variant. This lets a single binary run on mixed # GPUs (e.g. Volta sm_70 + Pascal sm_61) without "no kernel image" errors. @@ -598,6 +601,11 @@ find_package(OpenMP) if(OpenMP_CXX_FOUND) target_link_libraries(dflash_common PRIVATE OpenMP::OpenMP_CXX) endif() +if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(dflash_common PUBLIC CUDA::cudart) +elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") + target_link_libraries(dflash_common PUBLIC hip::host) +endif() if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") target_compile_options(dflash_common PRIVATE @@ -643,8 +651,9 @@ if(DFLASH27B_TESTS) add_executable(test_flashprefill_kernels test/test_flashprefill_kernels.cpp) set_source_files_properties(test/test_flashprefill_kernels.cpp PROPERTIES LANGUAGE HIP) set_target_properties(test_flashprefill_kernels PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") - target_include_directories(test_flashprefill_kernels PRIVATE + target_include_directories(test_flashprefill_kernels BEFORE PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/src/hip_compat ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) target_link_libraries(test_flashprefill_kernels PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) endif() @@ -656,14 +665,6 @@ if(DFLASH27B_TESTS) target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common CUDA::cudart) 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. - if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_gpu_sampler_cuda.cpp") - add_executable(test_gpu_sampler_cuda test/test_gpu_sampler_cuda.cpp) - target_include_directories(test_gpu_sampler_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) - target_link_libraries(test_gpu_sampler_cuda PRIVATE dflash_common CUDA::cudart) - add_test(NAME gpu_sampler_cuda COMMAND test_gpu_sampler_cuda) - endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_kv_quant.cpp") add_executable(test_kv_quant test/test_kv_quant.cpp) target_include_directories(test_kv_quant PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) @@ -787,6 +788,38 @@ if(DFLASH27B_TESTS) endif() add_test(NAME deepseek4_unit COMMAND test_deepseek4_unit) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_plan.cpp") + add_executable(test_expert_split_plan test/test_expert_split_plan.cpp) + target_include_directories(test_expert_split_plan PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_expert_split_plan PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_expert_split_plan PRIVATE dflash_common) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_runtime.cpp") + add_executable(test_expert_split_runtime test/test_expert_split_runtime.cpp) + target_include_directories(test_expert_split_runtime PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_expert_split_runtime PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_expert_split_runtime PRIVATE dflash_common) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_expert_split_target_config.cpp") + add_executable(test_expert_split_target_config test/test_expert_split_target_config.cpp) + target_include_directories(test_expert_split_target_config PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_expert_split_target_config PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_expert_split_target_config PRIVATE dflash_common) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_moe_expert_compute_multi_target.cpp") + add_executable(test_moe_expert_compute_multi_target test/test_moe_expert_compute_multi_target.cpp) + target_include_directories(test_moe_expert_compute_multi_target PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_include_directories(test_moe_expert_compute_multi_target PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + endif() + target_link_libraries(test_moe_expert_compute_multi_target PRIVATE dflash_common) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/smoke_load_draft.cpp") add_executable(smoke_load_draft test/smoke_load_draft.cpp) target_include_directories(smoke_load_draft PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) @@ -852,6 +885,128 @@ if(DFLASH27B_TESTS) target_include_directories(bench_moe_stream PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS} ${CUDAToolkit_INCLUDE_DIRS}) target_link_libraries(bench_moe_stream PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET} CUDA::cudart) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_expert_split_dispatch.cpp") + add_executable(bench_expert_split_dispatch test/bench_expert_split_dispatch.cpp) + target_include_directories(bench_expert_split_dispatch PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_expert_split_dispatch PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_multi_target_scatter.cpp") + add_executable(bench_multi_target_scatter test/bench_multi_target_scatter.cpp) + target_include_directories(bench_multi_target_scatter PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_multi_target_scatter PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_multi_target_dispatch_pack.cpp") + add_executable(bench_multi_target_dispatch_pack test/bench_multi_target_dispatch_pack.cpp) + target_include_directories(bench_multi_target_dispatch_pack PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_multi_target_dispatch_pack PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_moe_expert_ipc_dtype.cpp") + add_executable(bench_moe_expert_ipc_dtype test/bench_moe_expert_ipc_dtype.cpp) + target_include_directories(bench_moe_expert_ipc_dtype PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_moe_expert_ipc_dtype PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_moe_expert_ipc_control.cpp") + add_executable(bench_moe_expert_ipc_control test/bench_moe_expert_ipc_control.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_remote_cold_combine.cpp") + add_executable(bench_remote_cold_combine test/bench_remote_cold_combine.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_remote_cold_output_zero.cpp") + add_executable(bench_remote_cold_output_zero test/bench_remote_cold_output_zero.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_remote_cold_result_reuse.cpp") + add_executable(bench_remote_cold_result_reuse test/bench_remote_cold_result_reuse.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_multi_target_parallel_output.cpp") + add_executable(bench_multi_target_parallel_output test/bench_multi_target_parallel_output.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_ffn_output_finalize.cpp") + add_executable(bench_ffn_output_finalize test/bench_ffn_output_finalize.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_reduced_stack_subbatch_copy.cpp") + add_executable(bench_reduced_stack_subbatch_copy test/bench_reduced_stack_subbatch_copy.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_parallel_target_spawn.cpp") + add_executable(bench_parallel_target_spawn test/bench_parallel_target_spawn.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_parallel_target_primary_order.cpp") + add_executable(bench_parallel_target_primary_order test/bench_parallel_target_primary_order.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_parallel_secondary_order.cpp") + add_executable(bench_parallel_secondary_order test/bench_parallel_secondary_order.cpp) + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_backend_transfer_async.cpp") + add_executable(bench_backend_transfer_async test/bench_backend_transfer_async.cpp) + target_include_directories(bench_backend_transfer_async PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_backend_transfer_async PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_backend_transfer_async PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_backend_transfer_async PRIVATE hip::host) + endif() + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_backend_host_upload.cpp") + add_executable(bench_backend_host_upload test/bench_backend_host_upload.cpp) + target_include_directories(bench_backend_host_upload PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_backend_host_upload PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_backend_host_upload PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_backend_host_upload PRIVATE hip::host) + endif() + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_cpu_moe_batch_parallel.cpp") + add_executable(bench_cpu_moe_batch_parallel test/bench_cpu_moe_batch_parallel.cpp) + target_include_directories(bench_cpu_moe_batch_parallel PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_cpu_moe_batch_parallel PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_cpu_moe_batch_parallel PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_cpu_moe_batch_parallel PRIVATE hip::host) + endif() + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_cpu_moe_padded_dispatch.cpp") + add_executable(bench_cpu_moe_padded_dispatch test/bench_cpu_moe_padded_dispatch.cpp) + target_include_directories(bench_cpu_moe_padded_dispatch PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_cpu_moe_padded_dispatch PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_cpu_moe_padded_dispatch PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_cpu_moe_padded_dispatch PRIVATE hip::host) + endif() + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_gpu_peer_transfer.cpp") + add_executable(bench_gpu_peer_transfer test/bench_gpu_peer_transfer.cpp) + target_include_directories(bench_gpu_peer_transfer PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_gpu_peer_transfer PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_gpu_peer_transfer PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_gpu_peer_transfer PRIVATE hip::host) + endif() + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_expert_split_small_batch_resident.cpp") + add_executable(bench_expert_split_small_batch_resident test/bench_expert_split_small_batch_resident.cpp) + target_include_directories(bench_expert_split_small_batch_resident PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(bench_expert_split_small_batch_resident PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_expert_split_small_batch_resident PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_expert_split_small_batch_resident PRIVATE hip::host) + endif() + endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_hip_moe_mul_mat_id_small_batch.cpp") + add_executable(bench_hip_moe_mul_mat_id_small_batch test/bench_hip_moe_mul_mat_id_small_batch.cpp) + target_include_directories(bench_hip_moe_mul_mat_id_small_batch PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_sources(bench_hip_moe_mul_mat_id_small_batch PRIVATE + src/common/backend_precision.cpp) + target_link_libraries(bench_hip_moe_mul_mat_id_small_batch PRIVATE ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(bench_hip_moe_mul_mat_id_small_batch PRIVATE CUDA::cudart) + else() + target_link_libraries(bench_hip_moe_mul_mat_id_small_batch PRIVATE hip::host) + endif() + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_laguna_daemon.cpp") add_executable(test_laguna_daemon test/test_laguna_daemon.cpp) target_include_directories(test_laguna_daemon PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index ddcdfc40f..e6673fad9 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3,6 +3,8 @@ #include "deepseek4_backend.h" #include "deepseek4_internal.h" #include "common/sampler.h" +#include "../common/expert_split_target_config.h" +#include "../common/moe_hybrid_types.h" #include "ggml.h" #include "ggml-backend.h" @@ -10,6 +12,7 @@ #include #include +#include #include #include #include @@ -37,6 +40,573 @@ static double gib(uint64_t bytes) { return (double) bytes / 1024.0 / 1024.0 / 1024.0; } +static int env_int(const char * name, int fallback) { + const char * value = std::getenv(name); + if (!value || !value[0]) return fallback; + return std::atoi(value); +} + +static double env_double(const char * name, double fallback) { + const char * value = std::getenv(name); + if (!value || !value[0]) return fallback; + return std::atof(value); +} + +static const char * env_str(const char * name, const char * fallback) { + const char * value = std::getenv(name); + return (value && value[0]) ? value : fallback; +} + +static int env_pct(const char * name, int fallback) { + const char * value = std::getenv(name); + if (!value || !value[0]) return fallback; + char * end = nullptr; + const long parsed = std::strtol(value, &end, 10); + if (end == value || *end != '\0' || parsed <= 0 || parsed > 100) { + return fallback; + } + return (int) parsed; +} + +static uint64_t apply_pct_bytes(uint64_t total_bytes, int pct) { + if (pct >= 100) return total_bytes; + return (total_bytes / 100ULL) * (uint64_t) pct + + ((total_bytes % 100ULL) * (uint64_t) pct) / 100ULL; +} + +static bool backend_is_cuda(ggml_backend_t backend); +static bool backend_is_hip(ggml_backend_t backend); +static bool backend_is_gpu(ggml_backend_t backend); + +} // namespace + +int deepseek4_expert_split_prefill_chunk_limit_from_memory( + int requested_chunk, + bool expert_split_enabled, + bool parent_is_gpu, + const DeepSeek4ExpertSplitGpuMemoryInfo & memory) { + int chunk = std::max(1, requested_chunk); + if (!expert_split_enabled || !parent_is_gpu) { + return chunk; + } + if (memory.total_bytes == 0) { + return std::min(chunk, 64); + } + + const uint64_t gib1 = 1024ULL * 1024ULL * 1024ULL; + const uint64_t gib24 = 24ULL * gib1; + const uint64_t gib48 = 48ULL * gib1; + const uint64_t gib80 = 80ULL * gib1; + const uint64_t total = memory.total_bytes; + const uint64_t free = memory.free_bytes; + if (total <= gib24 + 512ULL * 1024ULL * 1024ULL) { + if (free >= 12ULL * gib1) { + return std::min(chunk, 256); + } + if (free >= 6ULL * gib1) { + return std::min(chunk, 128); + } + return std::min(chunk, 64); + } + if (total <= gib48 + 512ULL * 1024ULL * 1024ULL) { + if (free >= 12ULL * gib1) { + return std::min(chunk, 256); + } + return std::min(chunk, 128); + } + if (total >= gib80 && free >= 24ULL * gib1) { + return std::min(chunk, 512); + } + return std::min(chunk, 256); +} + +bool deepseek4_expert_split_should_disable_cached_decode_from_memory( + bool expert_split_enabled, + bool parent_is_gpu, + bool parent_is_cuda, + bool parent_is_hip, + const DeepSeek4ExpertSplitGpuMemoryInfo & memory) { + if (!expert_split_enabled || !parent_is_gpu) { + return false; + } + if (parent_is_hip) { + return false; + } + if (!parent_is_cuda) { + return true; + } + if (memory.total_bytes == 0) { + return true; + } + + const uint64_t gib1 = 1024ULL * 1024ULL * 1024ULL; + const uint64_t gib24 = 24ULL * gib1; + if (memory.total_bytes > gib24 + 512ULL * 1024ULL * 1024ULL) { + return false; + } + return memory.free_bytes < 3ULL * gib1; +} + +bool deepseek4_expert_split_requires_parent_cuda_graph_disable( + bool expert_split_enabled, + bool parent_is_cuda, + const std::vector & targets) { + if (!expert_split_enabled || !parent_is_cuda || targets.size() <= 1) { + return false; + } + for (size_t i = 1; i < targets.size(); ++i) { + const ExpertSplitComputeTargetRuntime & target = targets[i]; + if (target.placement.total_hot <= 0) { + continue; + } + if (target.target.backend.empty() || target.target.backend == "cpu") { + continue; + } + return true; + } + return false; +} + +uint64_t deepseek4_expert_split_budget_from_memory( + const DeepSeek4ExpertSplitBudgetMemoryInfo & memory) { + const uint64_t fixed_bytes = + memory.core_bytes + + memory.kv_bytes + + memory.warm_bytes + + memory.safety_bytes; + if (memory.total_bytes <= fixed_bytes) { + return 0; + } + + uint64_t budget_bytes = memory.total_bytes - fixed_bytes; + if (memory.parent_is_igpu && memory.host_budget_cap_bytes > 0) { + budget_bytes = std::min(budget_bytes, memory.host_budget_cap_bytes); + } + if (memory.total_expert_bytes > 0 && + budget_bytes > memory.total_expert_bytes) { + budget_bytes = memory.total_expert_bytes; + } + return budget_bytes; +} + +uint64_t deepseek4_expert_split_primary_capacity_for_targets( + uint64_t expert_budget_bytes, + uint64_t hot_budget_bytes, + size_t configured_targets) { + if (configured_targets <= 1 || hot_budget_bytes == 0) { + return expert_budget_bytes; + } + if (expert_budget_bytes == 0) { + return 0; + } + return std::min(expert_budget_bytes, hot_budget_bytes); +} + +uint64_t deepseek4_expert_split_non_cpu_capacity_bytes( + const std::vector & targets) { + uint64_t total_bytes = 0; + for (const ExpertSplitTarget & target : targets) { + if (target.backend == "cpu") { + continue; + } + if (target.unlimited) { + return std::numeric_limits::max(); + } + const uint64_t usable = target.usable_bytes(); + if (usable == 0) { + continue; + } + if (std::numeric_limits::max() - total_bytes < usable) { + return std::numeric_limits::max(); + } + total_bytes += usable; + } + return total_bytes; +} + +uint64_t deepseek4_expert_split_effective_budget_for_targets( + uint64_t requested_budget_bytes, + uint64_t total_expert_bytes, + const std::vector & targets) { + if (requested_budget_bytes == 0 || total_expert_bytes == 0 || targets.empty()) { + return requested_budget_bytes; + } + + const bool has_cpu_target = std::any_of( + targets.begin(), targets.end(), + [](const ExpertSplitTarget & target) { + return target.backend == "cpu"; + }); + if (!has_cpu_target) { + return requested_budget_bytes; + } + + const uint64_t non_cpu_capacity = + deepseek4_expert_split_non_cpu_capacity_bytes(targets); + if (non_cpu_capacity == 0) { + return requested_budget_bytes; + } + + const uint64_t promoted_budget = + std::min(total_expert_bytes, non_cpu_capacity); + return std::max(requested_budget_bytes, promoted_budget); +} + +uint64_t deepseek4_expert_split_effective_hot_budget( + uint64_t expert_budget_bytes, + uint64_t hot_budget_bytes, + size_t configured_targets) { + if (expert_budget_bytes == 0) { + return 0; + } + if (configured_targets > 1) { + return expert_budget_bytes; + } + if (hot_budget_bytes == 0) { + return expert_budget_bytes; + } + return std::min(expert_budget_bytes, hot_budget_bytes); +} + +bool deepseek4_expert_split_hash_ids_to_hotness_counts( + const int32_t * expert_ids, + int n_token_ids, + int n_expert_used, + int n_expert, + std::vector & out_counts) { + out_counts.clear(); + if (!expert_ids || + n_token_ids <= 0 || + n_expert_used <= 0 || + n_expert <= 0) { + return false; + } + + out_counts.assign((size_t) n_expert, 0); + const size_t total_ids = + (size_t) n_token_ids * (size_t) n_expert_used; + for (size_t i = 0; i < total_ids; ++i) { + const int32_t expert = expert_ids[i]; + if (expert < 0 || expert >= n_expert) { + out_counts.clear(); + return false; + } + out_counts[(size_t) expert] += 1; + } + return true; +} + +bool deepseek4_expert_split_route_bias_to_hotness_counts( + const float * route_bias, + int n_expert, + uint64_t layer_total, + std::vector & out_counts) { + out_counts.clear(); + if (!route_bias || n_expert <= 0 || layer_total == 0) { + return false; + } + + const uint64_t minimum_total = + std::max(layer_total, (uint64_t) n_expert); + float max_bias = -INFINITY; + bool saw_finite = false; + for (int expert = 0; expert < n_expert; ++expert) { + const float bias = route_bias[expert]; + if (!std::isfinite(bias)) { + continue; + } + max_bias = saw_finite ? std::max(max_bias, bias) : bias; + saw_finite = true; + } + if (!saw_finite) { + return false; + } + + out_counts.assign((size_t) n_expert, 1); + uint64_t remaining = minimum_total - (uint64_t) n_expert; + if (remaining == 0) { + return true; + } + + std::vector weights((size_t) n_expert, 0.0L); + long double weight_sum = 0.0L; + for (int expert = 0; expert < n_expert; ++expert) { + double delta = (double) route_bias[expert] - (double) max_bias; + if (!std::isfinite(delta)) { + delta = -16.0; + } + delta = std::max(-16.0, std::min(0.0, delta)); + const long double weight = std::exp(delta); + weights[(size_t) expert] = weight; + weight_sum += weight; + } + if (!(weight_sum > 0.0L)) { + out_counts.clear(); + return false; + } + + struct ShareRemainder { + int expert = 0; + long double remainder = 0.0L; + }; + std::vector remainders; + remainders.reserve((size_t) n_expert); + + uint64_t assigned = 0; + for (int expert = 0; expert < n_expert; ++expert) { + const long double raw = + ((long double) remaining * weights[(size_t) expert]) / weight_sum; + const uint64_t extra = (uint64_t) std::floor(raw); + out_counts[(size_t) expert] += extra; + assigned += extra; + remainders.push_back({expert, raw - (long double) extra}); + } + + std::stable_sort(remainders.begin(), remainders.end(), + [](const ShareRemainder & lhs, const ShareRemainder & rhs) { + if (lhs.remainder != rhs.remainder) { + return lhs.remainder > rhs.remainder; + } + return lhs.expert < rhs.expert; + }); + for (uint64_t i = assigned; i < remaining; ++i) { + out_counts[(size_t) remainders[(size_t) (i - assigned)].expert] += 1; + } + + return true; +} + +int deepseek4_expert_split_effective_cache_slots( + int requested_cache_slots, + size_t configured_targets) { + if (configured_targets > 1) { + return 0; + } + return std::max(0, requested_cache_slots); +} + +static uint64_t ceil_div_u64(uint64_t value, uint64_t denom) { + if (denom == 0) { + return 0; + } + return value / denom + ((value % denom) != 0 ? 1ULL : 0ULL); +} + +DeepSeek4ExpertSplitHotCachePlan deepseek4_expert_split_hot_cache_plan( + uint64_t expert_budget_bytes, + uint64_t total_expert_bytes, + int n_expert, + int n_expert_used, + bool parent_is_igpu, + uint64_t igpu_host_cap_bytes, + uint64_t igpu_free_bytes) { + DeepSeek4ExpertSplitHotCachePlan plan; + plan.hot_bytes = expert_budget_bytes; + if (expert_budget_bytes == 0 || total_expert_bytes == 0 || n_expert <= 0) { + return plan; + } + + const uint64_t bytes_per_slot = + total_expert_bytes / (uint64_t) n_expert; + if (bytes_per_slot == 0) { + return plan; + } + + const int min_hot_slots = + std::max(0, std::min(n_expert, n_expert_used)); + const uint64_t min_hot_bytes = + bytes_per_slot * (uint64_t) min_hot_slots; + if (expert_budget_bytes <= min_hot_bytes) { + return plan; + } + + const int max_cache_slots = (int) std::min( + (expert_budget_bytes - min_hot_bytes) / bytes_per_slot, + (uint64_t) std::max(0, n_expert - min_hot_slots)); + if (max_cache_slots <= 0) { + return plan; + } + + int requested_slots = 0; + const bool explicit_slots = + std::getenv("DFLASH_DEEPSEEK4_CACHE_SLOTS") != nullptr; + if (const char * raw = std::getenv("DFLASH_DEEPSEEK4_CACHE_SLOTS")) { + requested_slots = std::max(0, std::atoi(raw)); + } else { + auto spark = spark_budget_split( + expert_budget_bytes, + total_expert_bytes, + n_expert, + /*core_kv_safety=*/0, + /*target_bytes=*/0); + requested_slots = std::max(0, spark.cache_slots); + + if (parent_is_igpu) { + const int reserve_pct = env_pct( + "DFLASH_DEEPSEEK4_CACHE_RESERVE_PCT", 20); + uint64_t reserve_bytes = + (expert_budget_bytes / 100ULL) * (uint64_t) reserve_pct + + ((expert_budget_bytes % 100ULL) * + (uint64_t) reserve_pct) / 100ULL; + if (const char * raw = + std::getenv("DFLASH_DEEPSEEK4_CACHE_RESERVE_MB")) { + const long long reserve_mb = std::atoll(raw); + if (reserve_mb > 0) { + reserve_bytes = + (uint64_t) reserve_mb * 1024ULL * 1024ULL; + } + } else { + const int reserve_cap_mb = env_int( + "DFLASH_DEEPSEEK4_CACHE_RESERVE_CAP_MB", 16384); + if (reserve_cap_mb > 0) { + const uint64_t reserve_cap_bytes = + (uint64_t) reserve_cap_mb * 1024ULL * 1024ULL; + reserve_bytes = + std::min(reserve_bytes, reserve_cap_bytes); + } + } + const int min_cache_slots = std::min( + n_expert, + std::max(n_expert_used * 2, 12)); + if (reserve_bytes > 0) { + requested_slots = std::max( + requested_slots, + (int) (reserve_bytes / bytes_per_slot)); + } + requested_slots = std::max(requested_slots, min_cache_slots); + } + } + + if (parent_is_igpu && !explicit_slots) { + uint64_t pinned_hot_cap = expert_budget_bytes; + if (igpu_host_cap_bytes > 0) { + pinned_hot_cap = std::min(pinned_hot_cap, igpu_host_cap_bytes); + } + if (igpu_free_bytes > 0) { + pinned_hot_cap = std::min(pinned_hot_cap, igpu_free_bytes); + } + pinned_hot_cap = std::min( + pinned_hot_cap, + apply_pct_bytes( + expert_budget_bytes, + env_pct("DFLASH_DEEPSEEK4_IGPU_BOOTSTRAP_CAP_PCT", 80))); + pinned_hot_cap = std::max(pinned_hot_cap, min_hot_bytes); + if (expert_budget_bytes > pinned_hot_cap) { + const uint64_t extra_slots = ceil_div_u64( + expert_budget_bytes - pinned_hot_cap, bytes_per_slot); + requested_slots = std::max( + requested_slots, + (int) std::min(extra_slots, (uint64_t) max_cache_slots)); + } + } + + plan.cache_slots = std::clamp(requested_slots, 0, max_cache_slots); + plan.hot_bytes = + expert_budget_bytes - + (uint64_t) plan.cache_slots * bytes_per_slot; + return plan; +} + +namespace { + +static int resolve_ds4_expert_split_prefill_chunk( + int requested_chunk, + bool expert_split_enabled, + PlacementBackend device_backend, + ggml_backend_t backend, + int gpu_index) { + int chunk = std::max(1, requested_chunk); + if (!expert_split_enabled || !backend_is_gpu(backend)) { + return chunk; + } + + const int env_cap = env_int("DFLASH_EXPERT_SPLIT_PREFILL_CHUNK_MAX", 0); + if (env_cap > 0) { + return std::min(chunk, env_cap); + } + + uint64_t gpu_free = 0; + uint64_t gpu_total = 0; + if (!query_expert_split_backend_memory( + device_backend, gpu_index, gpu_free, gpu_total)) { + return std::min(chunk, 64); + } + return dflash::common::deepseek4_expert_split_prefill_chunk_limit_from_memory( + chunk, + expert_split_enabled, + /*parent_is_gpu=*/true, + DeepSeek4ExpertSplitGpuMemoryInfo{gpu_free, gpu_total}); +} + +static bool should_disable_ds4_expert_split_cached_decode( + bool expert_split_enabled, + PlacementBackend device_backend, + ggml_backend_t backend, + int gpu_index) { + if (!expert_split_enabled || !backend_is_gpu(backend)) { + return false; + } + if (env_flag_enabled("DFLASH_DS4_FORCE_MULTI_TARGET_DECODE_CACHE")) { + return false; + } + if (env_flag_enabled("DFLASH_DS4_DISABLE_MULTI_TARGET_DECODE_CACHE")) { + return true; + } + if (backend_is_hip(backend)) { + return false; + } + if (!backend_is_cuda(backend)) { + return true; + } + + uint64_t gpu_free = 0; + uint64_t gpu_total = 0; + if (!query_expert_split_backend_memory( + device_backend, gpu_index, gpu_free, gpu_total)) { + return true; + } + return dflash::common::deepseek4_expert_split_should_disable_cached_decode_from_memory( + expert_split_enabled, + /*parent_is_gpu=*/true, + /*parent_is_cuda=*/true, + /*parent_is_hip=*/false, + DeepSeek4ExpertSplitGpuMemoryInfo{gpu_free, gpu_total}); +} + +static PlacementBackend resolved_device_backend(const DevicePlacement & device) { + return device.backend == PlacementBackend::Auto + ? compiled_placement_backend() + : device.backend; +} + +static bool backend_is_cuda(ggml_backend_t backend) { + const char * name = ggml_backend_name(backend); + return name && std::strstr(name, "CUDA") != nullptr; +} + +static bool backend_is_hip(ggml_backend_t backend) { + const char * name = ggml_backend_name(backend); + return name && + (std::strstr(name, "HIP") != nullptr || + std::strstr(name, "ROCm") != nullptr); +} + +static bool backend_is_gpu(ggml_backend_t backend) { + return backend_is_cuda(backend) || backend_is_hip(backend); +} + +static bool resolve_ds4_expert_split_targets_from_env( + uint64_t primary_capacity_bytes, + std::vector & out, + std::string * err) { + return dflash::common::resolve_expert_split_targets_from_env( + "DFLASH_DEEPSEEK4_EXPERT_TARGETS", + "DFLASH_DEEPSEEK4_EXPERT_TARGET_CAPS", + primary_capacity_bytes, + out, + err); +} + static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTelemetry & src) { dst.total_us += src.total_us; dst.embed_us += src.embed_us; @@ -66,10 +636,34 @@ static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTeleme 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_cache_promote_us += src.ffn_cache_promote_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.worker_us += src.worker_us; + dst.worker_parent_write_us += src.worker_parent_write_us; + dst.worker_parent_wait_us += src.worker_parent_wait_us; + dst.worker_parent_read_us += src.worker_parent_read_us; + dst.worker_request_read_us += src.worker_request_read_us; + dst.worker_partition_us += src.worker_partition_us; + dst.worker_resident_eval_us += src.worker_resident_eval_us; + dst.worker_miss_build_us += src.worker_miss_build_us; + dst.worker_miss_eval_us += src.worker_miss_eval_us; + dst.worker_request_bytes += src.worker_request_bytes; + dst.worker_response_bytes += src.worker_response_bytes; + dst.worker_hot_graph_builds += src.worker_hot_graph_builds; + dst.worker_hot_graph_hits += src.worker_hot_graph_hits; + dst.worker_cold_graph_builds += src.worker_cold_graph_builds; + dst.worker_cold_graph_hits += src.worker_cold_graph_hits; + dst.worker_hot_graph_build_us += src.worker_hot_graph_build_us; + dst.worker_hot_input_us += src.worker_hot_input_us; + dst.worker_hot_compute_us += src.worker_hot_compute_us; + dst.worker_hot_read_us += src.worker_hot_read_us; + dst.worker_cold_graph_build_us += src.worker_cold_graph_build_us; + dst.worker_cold_input_us += src.worker_cold_input_us; + dst.worker_cold_compute_us += src.worker_cold_compute_us; + dst.worker_cold_read_us += src.worker_cold_read_us; dst.hc_post_ffn_us += src.hc_post_ffn_us; dst.output_us += src.output_us; dst.sample_us += src.sample_us; @@ -95,8 +689,14 @@ static void log_step_tel(const char * phase, "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=%.1fms hot=%.1fms cold=%.1fms combine=%.1fms partition=%.1fms cache=%.1fms worker=%.1fms " "ffn_hot_graph_build=%llu ffn_hot_graph_hit=%llu ffn_cold_graph_build=%llu ffn_cold_graph_hit=%llu " + "worker_write=%.1fms worker_wait=%.1fms worker_read=%.1fms worker_req_read=%.1fms " + "worker_part=%.1fms worker_resident=%.1fms worker_miss_build=%.1fms worker_miss=%.1fms " + "worker_hot_graph_build=%llu worker_hot_graph_hit=%llu worker_cold_graph_build=%llu worker_cold_graph_hit=%llu " + "worker_hot_build=%.1fms worker_hot_input=%.1fms worker_hot_compute=%.1fms worker_hot_read=%.1fms " + "worker_cold_build=%.1fms worker_cold_input=%.1fms worker_cold_compute=%.1fms worker_cold_read=%.1fms " + "worker_req_kib=%.1f worker_resp_kib=%.1f " "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", @@ -108,9 +708,20 @@ static void log_step_tel(const char * phase, 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), + ms(t.ffn_partition_us), ms(t.ffn_cache_promote_us), ms(t.worker_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.worker_parent_write_us), ms(t.worker_parent_wait_us), ms(t.worker_parent_read_us), + ms(t.worker_request_read_us), ms(t.worker_partition_us), ms(t.worker_resident_eval_us), + ms(t.worker_miss_build_us), ms(t.worker_miss_eval_us), + (unsigned long long)t.worker_hot_graph_builds, (unsigned long long)t.worker_hot_graph_hits, + (unsigned long long)t.worker_cold_graph_builds, (unsigned long long)t.worker_cold_graph_hits, + ms(t.worker_hot_graph_build_us), ms(t.worker_hot_input_us), + ms(t.worker_hot_compute_us), ms(t.worker_hot_read_us), + ms(t.worker_cold_graph_build_us), ms(t.worker_cold_input_us), + ms(t.worker_cold_compute_us), ms(t.worker_cold_read_us), + (double)t.worker_request_bytes / 1024.0, + (double)t.worker_response_bytes / 1024.0, ms(t.hc_pre_attn_us + t.hc_pre_ffn_us), ms(t.hc_pre_build_us), ms(t.hc_pre_input_us), @@ -143,11 +754,17 @@ struct Ds4HybridBudgetInfo { Ds4ExpertMemoryInfo mem; size_t gpu_free = 0; size_t gpu_total = 0; + bool parent_is_igpu = false; + uint64_t host_available_bytes = 0; + uint64_t host_budget_cap_bytes = 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 auto_expert_budget = 0; uint64_t expert_budget = 0; + uint64_t hot_budget = 0; + int cache_slots = 0; int max_hot_per_layer = 0; }; @@ -245,30 +862,52 @@ static void fill_prefix_hot_placement(const DeepSeek4Weights & w, } static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, + PlacementBackend device_backend, 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) { + ExpertSplitBackendMemoryInfo memory_info; + if (!query_expert_split_backend_memory_info( + device_backend, gpu, memory_info) || + memory_info.total_bytes == 0) { if (err) *err = "could not query GPU memory"; return false; } + out.gpu_free = (size_t) memory_info.free_bytes; + out.gpu_total = (size_t) memory_info.total_bytes; + out.parent_is_igpu = + memory_info.device_type == GGML_BACKEND_DEVICE_TYPE_IGPU; if (!compute_ds4_expert_memory_info(w, nullptr, out.mem, err)) { return false; } - out.core_bytes = out.gpu_total - out.gpu_free; + out.core_bytes = + moe_hybrid_core_bytes_from_memory("deepseek4", out.gpu_free, out.gpu_total); 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; + ExpertSplitIgpuHostBudgetInfo host_budget; + if (out.parent_is_igpu) { + (void) query_expert_split_igpu_host_budget_info(host_budget); + out.host_available_bytes = host_budget.available_bytes; + out.host_budget_cap_bytes = host_budget.cap_bytes; } + + out.auto_expert_budget = deepseek4_expert_split_budget_from_memory( + DeepSeek4ExpertSplitBudgetMemoryInfo{ + (uint64_t) out.gpu_total, + (uint64_t) out.gpu_free, + out.core_bytes, + out.kv_bytes, + out.warm_bytes, + out.safety_bytes, + out.mem.total_expert_bytes, + out.parent_is_igpu, + out.host_available_bytes, + out.host_budget_cap_bytes, + }); + out.expert_budget = out.auto_expert_budget; 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) { @@ -280,9 +919,21 @@ static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, return false; } - out.max_hot_per_layer = std::min(w.n_expert, (int) (out.expert_budget / out.mem.bytes_per_uniform_round)); + const auto hot_cache_plan = deepseek4_expert_split_hot_cache_plan( + out.expert_budget, + out.mem.total_expert_bytes, + w.n_expert, + w.n_expert_used, + out.parent_is_igpu, + out.host_budget_cap_bytes, + (uint64_t) out.gpu_free); + out.hot_budget = hot_cache_plan.hot_bytes; + out.cache_slots = hot_cache_plan.cache_slots; + out.max_hot_per_layer = std::min( + w.n_expert, + (int) (out.hot_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"; + if (err) *err = "pinned hot budget is smaller than one uniform expert round"; return false; } return true; @@ -298,6 +949,8 @@ static MoeHybridConfig make_ds4_parent_worker_cfg(const DeepSeek4Weights & w) { cfg.n_layer = w.n_layer; cfg.first_moe_layer = 0; cfg.swiglu_clamp = w.swiglu_clamp_exp; + static const int sm = query_gpu_compute_sm(); + cfg.mmq_safe_full_batch = (sm >= 80); cfg.materialize_cold_experts = false; return cfg; } @@ -310,6 +963,189 @@ static MoeHybridConfig make_ds4_parent_cpu_tail_cfg(const DeepSeek4Weights & w) return cfg; } +static MoeHybridConfig make_ds4_parent_gpu_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::Gpu; + return cfg; +} + +static MoeLayerDesc make_ds4_moe_layer_desc(const DeepSeek4Layer & L) { + MoeLayerDesc desc; + desc.ffn_gate_exps = L.ffn_gate_exps; + desc.ffn_up_exps = L.ffn_up_exps; + desc.ffn_down_exps = L.ffn_down_exps; + desc.ffn_gate_up_exps = nullptr; + desc.ffn_gate_shexp = L.ffn_gate_shexp; + desc.ffn_up_shexp = L.ffn_up_shexp; + desc.ffn_down_shexp = L.ffn_down_shexp; + desc.ffn_gate_inp_shexp = nullptr; + return desc; +} + +static bool deepseek4_target_list_prefers_static_hotness_prior( + const std::vector & targets) { + if (targets.size() > 1) { + return true; + } + return std::any_of( + targets.begin(), targets.end(), + [](const ExpertSplitTarget & target) { + return target.backend == "cpu"; + }); +} + +static bool init_uniform_hotness( + int n_layer, + int n_expert, + int n_expert_used, + MoeHybridRoutingStats & hotness, + std::string * err) { + if (!hotness.init(n_layer, n_expert, n_expert_used)) { + if (err) *err = "failed to initialize uniform hotness"; + return false; + } + std::fill(hotness.counts.begin(), hotness.counts.end(), 1); + hotness.layer_totals.assign((size_t) n_layer, (uint64_t) n_expert); + return true; +} + +static bool build_deepseek4_static_hotness_prior( + const DeepSeek4Weights & w, + MoeHybridRoutingStats & hotness, + std::string & source) { + if (!hotness.init(w.n_layer, w.n_expert, w.n_expert_used)) { + return false; + } + + constexpr uint64_t kRouteBiasLayerTotal = 1ULL << 20; + bool used_hash_prior = false; + bool used_bias_prior = false; + std::vector layer_counts; + std::vector hash_ids; + std::vector route_bias; + + for (int layer = 0; layer < w.n_layer; ++layer) { + const DeepSeek4Layer & L = w.layers[(size_t) layer]; + bool seeded = false; + + if (layer < w.n_hash_layer && + L.ffn_gate_tid2eid && + L.ffn_gate_tid2eid->ne[0] == w.n_expert_used && + L.ffn_gate_tid2eid->ne[1] > 0) { + const int n_token_ids = (int) L.ffn_gate_tid2eid->ne[1]; + hash_ids.resize((size_t) n_token_ids * (size_t) w.n_expert_used); + ggml_backend_tensor_get( + L.ffn_gate_tid2eid, + hash_ids.data(), + 0, + sizeof(int32_t) * hash_ids.size()); + if (deepseek4_expert_split_hash_ids_to_hotness_counts( + hash_ids.data(), + n_token_ids, + w.n_expert_used, + w.n_expert, + layer_counts)) { + used_hash_prior = true; + seeded = true; + } + } + + if (!seeded && + L.ffn_exp_probs_b && + L.ffn_exp_probs_b->ne[0] == w.n_expert) { + route_bias.resize((size_t) w.n_expert); + ggml_backend_tensor_get( + L.ffn_exp_probs_b, + route_bias.data(), + 0, + sizeof(float) * route_bias.size()); + if (deepseek4_expert_split_route_bias_to_hotness_counts( + route_bias.data(), + w.n_expert, + kRouteBiasLayerTotal, + layer_counts)) { + used_bias_prior = true; + seeded = true; + } + } + + if (!seeded) { + layer_counts.assign((size_t) w.n_expert, 1); + } + + uint64_t layer_total = 0; + for (int expert = 0; expert < w.n_expert; ++expert) { + const uint64_t count = layer_counts[(size_t) expert]; + hotness.counts[(size_t) layer * (size_t) w.n_expert + (size_t) expert] = + count; + layer_total += count; + } + hotness.layer_totals[(size_t) layer] = layer_total; + } + + if (used_hash_prior && used_bias_prior) { + source = "static-prior(hash+bias)"; + } else if (used_hash_prior) { + source = "static-prior(hash)"; + } else if (used_bias_prior) { + source = "static-prior(bias)"; + } else { + source = "uniform"; + } + return used_hash_prior || used_bias_prior; +} + +static bool build_ds4_cold_owner_placement( + const MoeHybridPlacement & target_placement, + MoeHybridPlacement & out, + std::vector> & cold_order_by_layer, + std::string * err) { + if (!target_placement.matches(target_placement.n_layer, + target_placement.n_expert, + target_placement.n_expert_used)) { + if (err) *err = "invalid expert split target placement"; + return false; + } + out = target_placement; + out.total_hot = 0; + out.hot_counts.assign((size_t) target_placement.n_layer, 0); + out.hot_expert_ids.resize((size_t) target_placement.n_layer); + cold_order_by_layer.resize((size_t) target_placement.n_layer); + for (int il = 0; il < target_placement.n_layer; ++il) { + const auto & target_owned = + target_placement.hot_expert_ids[(size_t) il]; + std::vector owned((size_t) target_placement.n_expert, 0); + auto & cold_order = cold_order_by_layer[(size_t) il]; + cold_order.clear(); + cold_order.reserve(target_owned.size()); + for (int32_t expert : target_owned) { + if (expert < 0 || expert >= target_placement.n_expert) { + if (err) *err = "expert split target placement expert id out of range"; + return false; + } + if (owned[(size_t) expert]) { + if (err) *err = "expert split target placement contains duplicates"; + return false; + } + owned[(size_t) expert] = 1; + cold_order.push_back(expert); + } + auto & hot_experts = out.hot_expert_ids[(size_t) il]; + hot_experts.clear(); + hot_experts.reserve((size_t) target_placement.n_expert - target_owned.size()); + for (int expert = 0; expert < target_placement.n_expert; ++expert) { + if (!owned[(size_t) expert]) { + hot_experts.push_back(expert); + } + } + out.hot_counts[(size_t) il] = (int) hot_experts.size(); + out.total_hot += (int) hot_experts.size(); + } + return true; +} + } // namespace DeepSeek4Backend::DeepSeek4Backend(const DeepSeek4BackendConfig & cfg) @@ -319,6 +1155,51 @@ DeepSeek4Backend::~DeepSeek4Backend() { shutdown(); } +bool DeepSeek4Backend::run_step_with_runtime_path( + const float * embed, + int n_tokens, + int kv_start, + std::vector & out_logits, + const int32_t * token_ids, + bool want_logits, + bool disable_cached_decode, + DeepSeek4StepTelemetry * telemetry) { + if (!moe_hybrid_) { + return deepseek4_step_layer_range(backend_, w_, cache_, hc_state_, + embed, n_tokens, kv_start, + 0, w_.n_layer, + want_logits ? &out_logits : nullptr, + token_ids, telemetry); + } + (void) disable_cached_decode; + + MoeExpertCompute * expert_compute = + expert_split_runtime_.enabled ? expert_split_runtime_.compute_ptr() + : expert_runtime_.compute_ptr(); + const MoeExpertLayer * expert_layers = + expert_split_runtime_.enabled ? expert_split_runtime_.layer_ptr(0) + : expert_runtime_.layer_ptr(0); + + return deepseek4_step(backend_, w_, cache_, embed, n_tokens, kv_start, out_logits, + moe_hybrid_.get(), token_ids, nullptr, + false, + want_logits, + telemetry, + routing_stats_.get(), + expert_compute, + expert_layers); +} + +void DeepSeek4Backend::reset_request_state() { + cache_.cur_pos = 0; + for (auto & layer : cache_.layers) { + layer.n_comp = 0; + layer.n_index_comp = 0; + } + std::fill(hc_state_.begin(), hc_state_.end(), 0.0f); + last_logits_.clear(); +} + bool DeepSeek4Backend::init() { backend_ = ggml_backend_cuda_init(cfg_.device.gpu); if (!backend_) { @@ -329,25 +1210,21 @@ bool DeepSeek4Backend::init() { snap_backend_ = ggml_backend_init_by_name("cpu", nullptr); - const PlacementBackend target_backend = - cfg_.device.backend == PlacementBackend::Auto - ? compiled_placement_backend() - : cfg_.device.backend; + const bool force_hybrid = env_flag_enabled("DFLASH_DEEPSEEK4_FORCE_HYBRID"); + if (force_hybrid) { + std::fprintf(stderr, "[deepseek4] force hybrid mode requested; skipping full model load\n"); + } - // 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"); + // Try full load first; if GPU OOM, fall back to hybrid mode automatically. + if (!force_hybrid && !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 failed: %s\n", cfg_.model_path); + std::fprintf(stderr, "[deepseek4] hybrid mode also 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"); + } else if (force_hybrid) { if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] hybrid mode also failed: %s\n", cfg_.model_path); + std::fprintf(stderr, "[deepseek4] forced hybrid mode failed: %s\n", cfg_.model_path); return false; } } @@ -357,11 +1234,7 @@ bool DeepSeek4Backend::init() { 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. - } + hc_state_.assign((size_t) w_.n_hc * (size_t) w_.n_embd, 0.0f); 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, @@ -374,7 +1247,9 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & MoeHybridPlacement & out, std::string * err) const { Ds4HybridBudgetInfo budget; - if (!compute_ds4_hybrid_budget_info(w, cfg_.device.gpu, max_ctx, budget, err)) { + if (!compute_ds4_hybrid_budget_info( + w, resolved_device_backend(cfg_.device), + cfg_.device.gpu, max_ctx, budget, err)) { return false; } @@ -387,20 +1262,406 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & } 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", + "[deepseek4] hybrid placement: gpu_total=%.2f GiB gpu_free=%.2f GiB core=%.2f GiB kv=%.2f GiB warm=%.2f GiB safety=%.2f GiB host_cap=%.2f GiB expert_budget=%.2f GiB pinned_hot=%.2f GiB cache_slots=%d 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.host_budget_cap_bytes), gib(budget.expert_budget), + gib(budget.hot_budget), + budget.cache_slots, hot_per_layer); log_ds4_expert_memory_info("placement", placed_mem, w.n_layer); return true; } +bool DeepSeek4Backend::build_expert_split_state_from_hotness( + const MoeHybridRoutingStats & hotness, + uint64_t expert_budget_bytes, + uint64_t hot_budget_bytes, + int max_hot_per_layer, + const std::vector & configured_targets, + const DeepSeek4Weights & w, + std::string * err) { + if ((int)layer_expert_bytes_.size() != w.n_layer) { + if (err) *err = "layer expert bytes not initialized"; + return false; + } + if (expert_budget_bytes == 0) { + if (err) *err = "expert budget must be > 0"; + return false; + } + if (!hotness.matches(w.n_layer, w.n_expert, w.n_expert_used)) { + if (err) *err = "hotness table dimensions do not match DeepSeek4 model"; + return false; + } + + ExpertSplitConfig cfg; + cfg.n_layer = w.n_layer; + cfg.n_expert = w.n_expert; + cfg.allow_implicit_cpu_fallback = true; + cfg.require_full_grid = true; + + std::vector targets = configured_targets; + if (targets.size() > 1 && !targets[0].unlimited) { + targets[0].capacity_bytes = + deepseek4_expert_split_primary_capacity_for_targets( + targets[0].capacity_bytes, + hot_budget_bytes, + targets.size()); + } + if (targets.empty()) { + const PlacementBackend backend = resolved_device_backend(cfg_.device); + targets = { + {std::string(placement_backend_name(backend)) + ":" + + std::to_string(cfg_.device.gpu), + placement_backend_name(backend), cfg_.device.gpu, + expert_budget_bytes, 0, false}, + }; + } + if (!build_capacity_weighted_expert_split_min_per_layer( + targets, std::min(w.n_expert_used, w.n_expert), + cfg.min_per_layer_by_target, err)) { + return false; + } + if (!targets.empty() && + std::none_of(cfg.min_per_layer_by_target.begin(), + cfg.min_per_layer_by_target.end(), + [](int v) { return v > 0; })) { + cfg.min_per_layer_by_target.assign(targets.size(), 0); + cfg.min_per_layer_by_target[0] = std::min(w.n_expert_used, w.n_expert); + } + + std::vector units; + units.reserve((size_t)w.n_layer * (size_t)w.n_expert); + for (int il = 0; il < w.n_layer; ++il) { + for (int ie = 0; ie < w.n_expert; ++ie) { + units.push_back(ExpertSplitUnit{ + il, + ie, + layer_expert_bytes_[(size_t)il], + (double)hotness.count(il, ie), + }); + } + } + + std::vector primary_hot_limits; + if (max_hot_per_layer > 0) { + primary_hot_limits.assign((size_t) w.n_layer, max_hot_per_layer); + } + + ExpertSplitStateComponents state; + if (!build_expert_split_state(cfg, targets, units, + w.n_expert_used, + primary_hot_limits.empty() + ? nullptr + : &primary_hot_limits, + state, err)) { + return false; + } + if (!validate_primary_expert_split_target(state.plan.targets, + resolved_device_backend(cfg_.device), + cfg_.device.gpu, err)) { + return false; + } + expert_split_state_ = std::move(state); + return true; +} + +bool DeepSeek4Backend::load_expert_split_placement( + const char * hotness_path, + const DeepSeek4Weights & w, + int max_ctx, + MoeHybridPlacement & out, + std::string * err) { + Ds4HybridBudgetInfo budget; + if (!compute_ds4_hybrid_budget_info( + w, resolved_device_backend(cfg_.device), + cfg_.device.gpu, max_ctx, budget, err)) { + return false; + } + layer_expert_bytes_ = budget.mem.layer_expert_bytes; + + std::vector configured_targets; + if (!resolve_ds4_expert_split_targets_from_env( + budget.expert_budget, configured_targets, err)) { + return false; + } + + uint64_t placement_budget_bytes = budget.expert_budget; + uint64_t placement_non_cpu_capacity_bytes = 0; + bool promoted_cpu_fallback_only = false; + if (!configured_targets.empty() && + budget.auto_expert_budget > budget.expert_budget) { + std::vector auto_capacity_targets; + if (!resolve_ds4_expert_split_targets_from_env( + budget.auto_expert_budget, auto_capacity_targets, err)) { + return false; + } + const uint64_t promoted_budget = + deepseek4_expert_split_effective_budget_for_targets( + budget.expert_budget, + budget.mem.total_expert_bytes, + auto_capacity_targets); + if (promoted_budget > placement_budget_bytes) { + placement_non_cpu_capacity_bytes = + deepseek4_expert_split_non_cpu_capacity_bytes( + auto_capacity_targets); + placement_budget_bytes = promoted_budget; + configured_targets = std::move(auto_capacity_targets); + promoted_cpu_fallback_only = true; + } + } + + MoeHybridRoutingStats hotness; + std::string hotness_source = "uniform"; + if (hotness_path && hotness_path[0]) { + if (!MoeHybridRoutingStats::load_csv(std::string(hotness_path), hotness, err)) { + return false; + } + if (!hotness.matches(w.n_layer, w.n_expert, w.n_expert_used)) { + if (err) *err = "hotness table dimensions do not match DeepSeek4 model"; + return false; + } + hotness_source = hotness_path; + } else if (deepseek4_target_list_prefers_static_hotness_prior( + configured_targets) && + build_deepseek4_static_hotness_prior( + w, hotness, hotness_source)) { + // hotness_source already describes the static prior that was used + } else if (!init_uniform_hotness( + w.n_layer, w.n_expert, w.n_expert_used, hotness, err)) { + return false; + } + + if (routing_stats_ && hotness_path && hotness_path[0] && + hotness.counts.size() == (size_t)w.n_layer * (size_t)w.n_expert) { + routing_stats_->counts = hotness.counts; + routing_stats_->layer_totals.assign((size_t)w.n_layer, 0); + for (int il = 0; il < w.n_layer; ++il) { + for (int ie = 0; ie < w.n_expert; ++ie) { + routing_stats_->layer_totals[(size_t)il] += + hotness.counts[(size_t)il * (size_t)w.n_expert + (size_t)ie]; + } + } + } + + const size_t configured_target_count = + configured_targets.empty() ? 1 : configured_targets.size(); + const uint64_t effective_hot_budget = + deepseek4_expert_split_effective_hot_budget( + placement_budget_bytes, + budget.hot_budget, + configured_target_count); + int effective_max_hot_per_layer = budget.max_hot_per_layer; + if (budget.mem.bytes_per_uniform_round > 0) { + effective_max_hot_per_layer = std::min( + w.n_expert, + (int) (effective_hot_budget / budget.mem.bytes_per_uniform_round)); + } + if (effective_max_hot_per_layer <= 0) { + if (err) { + *err = + "effective pinned hot budget is smaller than one uniform expert round"; + } + return false; + } + + expert_split_total_budget_bytes_ = budget.expert_budget; + expert_split_hot_budget_bytes_ = effective_hot_budget; + expert_split_cache_slots_ = + deepseek4_expert_split_effective_cache_slots( + budget.cache_slots, + configured_target_count); + + if (!build_expert_split_state_from_hotness( + hotness, placement_budget_bytes, effective_hot_budget, + effective_max_hot_per_layer, configured_targets, w, err)) { + return false; + } + expert_split_total_budget_bytes_ = placement_budget_bytes; + if (expert_split_cache_slots_ != budget.cache_slots) { + std::fprintf(stderr, + "[deepseek4] expert_split cache disabled for multi-target placement: requested_slots=%d targets=%zu\n", + budget.cache_slots, + configured_target_count); + } + if (promoted_cpu_fallback_only) { + std::fprintf(stderr, + "[deepseek4] expert_split promoted non-CPU placement budget: requested=%.2f GiB effective=%.2f GiB non_cpu_cap=%.2f GiB total_experts=%.2f GiB\n", + gib(budget.expert_budget), + gib(placement_budget_bytes), + gib(placement_non_cpu_capacity_bytes), + gib(budget.mem.total_expert_bytes)); + } + out = expert_split_state_.materialization.primary_placement; + + Ds4ExpertMemoryInfo placed_mem; + if (!compute_ds4_expert_memory_info(w, &out, placed_mem, err)) { + return false; + } + std::fprintf(stderr, + "[deepseek4] expert_split placement: gpu_total=%.2f GiB gpu_free=%.2f GiB core=%.2f GiB kv=%.2f GiB warm=%.2f GiB safety=%.2f GiB host_cap=%.2f GiB expert_budget=%.2f GiB placement_budget=%.2f GiB pinned_hot=%.2f GiB cache_slots=%d targets=%zu source=%s\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.host_budget_cap_bytes), + gib(budget.expert_budget), + gib(placement_budget_bytes), + gib(effective_hot_budget), + expert_split_cache_slots_, + configured_target_count, + hotness_source.c_str()); + log_ds4_expert_memory_info("expert_split", placed_mem, w.n_layer); + return true; +} + +bool DeepSeek4Backend::init_expert_split_runtime(std::string * err) { + expert_split_runtime_.reset(); + if (!moe_hybrid_ || expert_split_state_.empty()) { + return true; + } + if (!expert_split_state_.compute_runtime.matches( + w_.n_layer, w_.n_expert, w_.n_expert_used)) { + if (err) *err = "expert split compute runtime does not match DeepSeek4 model"; + return false; + } + + MoeExpertComputeRuntimeConfig runtime_cfg; + runtime_cfg.target_path = cfg_.model_path ? cfg_.model_path : ""; + runtime_cfg.n_layer = w_.n_layer; + runtime_cfg.n_expert = w_.n_expert; + runtime_cfg.n_expert_used = w_.n_expert_used; + runtime_cfg.n_embd = w_.n_embd; + runtime_cfg.n_ff_exp = w_.n_ff_exp; + runtime_cfg.enabled = true; + runtime_cfg.log_prefix = "[deepseek4-expert-split]"; + runtime_cfg.local_target_factory_key = "deepseek4-local-gpu-cold"; + runtime_cfg.local_target_compute_factory = + [this](const ExpertSplitComputeTargetRuntime & split_target, + std::string * local_err) { + return make_expert_split_local_target_compute(split_target, local_err); + }; + + std::vector layer_descs((size_t)w_.n_layer); + for (int il = 0; il < w_.n_layer; ++il) { + layer_descs[(size_t)il] = make_ds4_moe_layer_desc(w_.layers[(size_t)il]); + } + + if (!ensure_multi_target_moe_expert_compute_runtime( + expert_split_runtime_, runtime_cfg, + expert_split_state_.compute_runtime, + *moe_hybrid_, layer_descs, err)) { + return false; + } + if (expert_split_runtime_.enabled && expert_split_runtime_.compute_ptr()) { + std::fprintf(stderr, "[deepseek4] expert_split compute ready: %s\n", + expert_split_runtime_.runtime_key.c_str()); + } + return true; +} + +std::unique_ptr DeepSeek4Backend::make_expert_split_local_target_compute( + const ExpertSplitComputeTargetRuntime & split_target, + std::string * err) const { + if (!cfg_.model_path || !cfg_.model_path[0]) { + if (err) *err = "DeepSeek4 local target compute requires a model path"; + return nullptr; + } + if (split_target.target.backend == "cpu") { + return nullptr; + } + if (split_target.target.device_id < 0) { + if (err) *err = "DeepSeek4 local target compute requires an explicit device id"; + return nullptr; + } + MoeHybridPlacement cold_owner_placement; + std::vector> cold_order_by_layer; + if (!build_ds4_cold_owner_placement(split_target.placement, + cold_owner_placement, + cold_order_by_layer, + err)) { + return nullptr; + } + + ggml_backend_t target_backend = + ggml_backend_cuda_init(std::max(0, split_target.target.device_id)); + if (!target_backend) { + if (err) *err = "failed to initialize DeepSeek4 local backend"; + return nullptr; + } + + MoeHybridConfig local_cfg = make_ds4_parent_gpu_tail_cfg(w_); + auto storage = std::make_unique(); + if (!build_deepseek4_moe_hybrid_storage_from_file( + cfg_.model_path, + target_backend, + w_, + cold_owner_placement, + &local_cfg, + *storage, + err, + /*cache_slots=*/0, + /*load_cold_tensors=*/true, + &cold_order_by_layer)) { + ggml_backend_free(target_backend); + return nullptr; + } + + std::vector layer_descs((size_t) w_.n_layer); + for (int il = 0; il < w_.n_layer; ++il) { + layer_descs[(size_t) il] = make_ds4_moe_layer_desc(w_.layers[(size_t) il]); + } + return make_local_gpu_cold_moe_expert_compute( + target_backend, + std::move(storage), + local_cfg, + std::move(layer_descs)); +} + +bool DeepSeek4Backend::init_single_target_expert_runtime(std::string * err) { + expert_runtime_.reset(); + if (!moe_hybrid_) { + return true; + } + + MoeExpertComputeRuntimeConfig runtime_cfg; + runtime_cfg.target_path = cfg_.model_path ? cfg_.model_path : ""; + runtime_cfg.n_layer = w_.n_layer; + runtime_cfg.n_expert = w_.n_expert; + runtime_cfg.n_expert_used = w_.n_expert_used; + runtime_cfg.n_embd = w_.n_embd; + runtime_cfg.n_ff_exp = w_.n_ff_exp; + runtime_cfg.enabled = true; + runtime_cfg.log_prefix = "[deepseek4-hybrid]"; + + std::vector layer_descs((size_t)w_.n_layer); + for (int il = 0; il < w_.n_layer; ++il) { + layer_descs[(size_t)il] = make_ds4_moe_layer_desc(w_.layers[(size_t)il]); + } + + if (!ensure_moe_expert_compute_runtime( + expert_runtime_, runtime_cfg, *moe_hybrid_, layer_descs, err)) { + return false; + } + if (expert_runtime_.compute_ptr()) { + std::fprintf(stderr, "[deepseek4] hybrid expert compute ready: %s\n", + expert_runtime_.runtime_key.c_str()); + } + return true; +} + bool DeepSeek4Backend::init_hybrid_model() { + expert_split_total_budget_bytes_ = 0; + expert_split_hot_budget_bytes_ = 0; + expert_split_cache_slots_ = 0; + TargetLoadPlan plan; plan.skip_expert_tensors = true; if (!load_deepseek4_gguf_partial(cfg_.model_path, backend_, plan, w_)) { @@ -409,10 +1670,21 @@ bool DeepSeek4Backend::init_hybrid_model() { return false; } + if (const char * stats_path = std::getenv("DFLASH_DEEPSEEK4_RUNTIME_STATS_OUT")) { + routing_stats_ = std::make_shared(); + MoeHybridConfig stats_cfg = make_ds4_parent_worker_cfg(w_); + if (!routing_stats_->init(stats_cfg)) { + std::fprintf(stderr, "[deepseek4] routing stats init failed\n"); + return false; + } + routing_stats_out_path_ = stats_path; + } + 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()); + const char * hotness_path = std::getenv("DFLASH_DEEPSEEK4_HOTNESS"); + if (!load_expert_split_placement(hotness_path, w_, max_ctx, moe_placement_, &err)) { + std::fprintf(stderr, "[deepseek4] failed to compute expert_split placement: %s\n", err.c_str()); return false; } @@ -427,43 +1699,64 @@ bool DeepSeek4Backend::init_hybrid_model() { } 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)) { + const std::vector> * cold_order_by_layer = + expert_split_state_.materialization.ordered_cold_union + ? &expert_split_state_.materialization.cold_expert_ids_by_layer + : nullptr; + if (!build_deepseek4_moe_hybrid_storage_from_file( + cfg_.model_path, backend_, w_, moe_placement_, nullptr, + *hybrid, &err, expert_split_cache_slots_, + /*load_cold_tensors=*/true, + cold_order_by_layer)) { std::fprintf(stderr, "[deepseek4] failed to build hybrid expert storage: %s\n", err.c_str()); 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; - } + moe_hybrid_ = std::move(hybrid); + if (!init_single_target_expert_runtime(&err)) { + std::fprintf(stderr, "[deepseek4] failed to initialize hybrid expert runtime: %s\n", + err.c_str()); + return false; + } + if (!init_expert_split_runtime(&err)) { + std::fprintf(stderr, "[deepseek4] failed to initialize expert_split runtime: %s\n", + err.c_str()); + return false; + } + if (deepseek4_expert_split_requires_parent_cuda_graph_disable( + expert_split_runtime_.enabled, + backend_is_cuda(backend_), + expert_split_state_.compute_runtime.targets)) { + setenv("GGML_CUDA_DISABLE_GRAPHS", "1", 1); 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); + "[deepseek4] expert_split multi-target on CUDA: disabling CUDA graphs for parent runtime\n"); + } + if (expert_split_runtime_.enabled) { + const PlacementBackend device_backend = resolved_device_backend(cfg_.device); + const bool disable_cached_decode = + should_disable_ds4_expert_split_cached_decode( + /*expert_split_enabled=*/true, + device_backend, + backend_, + cfg_.device.gpu); + std::fprintf(stderr, + "[deepseek4] expert_split decode cache: %s parent_backend=%s device=%d\n", + disable_cached_decode ? "eager" : "cached", + placement_backend_name(device_backend), + cfg_.device.gpu); } - - 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, ""); + std::fprintf(stderr, + "[deepseek4] hybrid experts ready: hot=%d cold=%d cold_backend=%s cache_slots=%d hot_budget=%.2f GiB expert_budget=%.2f GiB\n", + moe_placement_.total_hot, + total_cold, + cold_backend, + expert_split_cache_slots_, + gib(expert_split_hot_budget_bytes_), + gib(expert_split_total_budget_bytes_)); return true; } @@ -474,105 +1767,85 @@ void DeepSeek4Backend::print_ready_banner() const { } 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(); - free_deepseek4_cache(cache_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; - free_deepseek4_weights(w_); + (void)what; + // TODO: Release GPU resources 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; - } - + (void)what; 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 bool use_hybrid_batch_prefill = + moe_hybrid_ && + backend_is_gpu(backend_) && + !env_flag_enabled("DFLASH_DS4_DISABLE_HYBRID_BATCH_PREFILL"); + const int requested_chunk = cfg_.chunk > 0 ? cfg_.chunk : 512; + const PlacementBackend device_backend = resolved_device_backend(cfg_.device); + const int chunk = use_hybrid_batch_prefill + ? resolve_ds4_expert_split_prefill_chunk( + requested_chunk, + expert_split_runtime_.enabled, + device_backend, + backend_, + cfg_.device.gpu) + : (moe_hybrid_ ? 1 : requested_chunk); + if (use_hybrid_batch_prefill && + expert_split_runtime_.enabled && + chunk < requested_chunk) { + std::fprintf(stderr, + "[deepseek4] expert_split prefill chunk capped: requested=%d effective=%d\n", + requested_chunk, chunk); + } const int n_total = (int)tokens.size(); int pos = kv_offset; last_logits_.clear(); + std::vector embed((size_t)w_.n_embd * (size_t)std::max(1, chunk)); + std::vector logits; + const bool force_intermediate_logits = + env_flag_enabled("DFLASH_DS4_FORCE_INTERMEDIATE_PREFILL_LOGITS"); const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); const auto phase_t0 = Clock::now(); DeepSeek4StepTelemetry tel_acc; int steps = 0; - + if (timing && moe_hybrid_) { + const char * backend_name = backend_ ? ggml_backend_name(backend_) : "(null)"; + std::fprintf(stderr, + "[deepseek4] prefill plan: tokens=%d requested_chunk=%d effective_chunk=%d hybrid_batch=%s expert_split=%s intermediate_logits=%s backend=%s\n", + n_total, + requested_chunk, + chunk, + use_hybrid_batch_prefill ? "on" : "off", + expert_split_runtime_.enabled ? "on" : "off", + force_intermediate_logits ? "on" : "off", + backend_name ? backend_name : "(null)"); + } for (int i = 0; i < n_total; i += chunk) { if (io.cancelled) return pos; const int n_tok = std::min(chunk, n_total - i); + const bool want_logits = + force_intermediate_logits || (i + n_tok) >= n_total; // Embed tokens - std::vector embed(w_.n_embd * n_tok); const auto embed_t0 = Clock::now(); w_.embedder.embed(tokens.data() + i, n_tok, embed.data()); 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())) { + logits.clear(); + if (!run_step_with_runtime_path(embed.data(), n_tok, pos, logits, + tokens.data() + i, + want_logits, + false, + timing ? &step_tel : nullptr)) { std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d\n", pos); return -1; } @@ -580,12 +1853,19 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, add_step_tel(tel_acc, step_tel); steps++; } - last_logits_ = std::move(logits); + if (want_logits) { + last_logits_ = std::move(logits); + } pos += n_tok; } if (timing) { log_step_tel("prefill", n_total, steps, elapsed_s(phase_t0), tel_acc); } + if (env_flag_enabled("DFLASH_DS4_TRACE_DECODE")) { + std::fprintf(stderr, + "[deepseek4-trace] prefill done committed=%d last_logits=%zu\n", + pos, last_logits_.size()); + } return pos; } @@ -599,10 +1879,24 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, const auto phase_t0 = Clock::now(); DeepSeek4StepTelemetry tel_acc; int steps = 0; - + std::vector embed((size_t)w_.n_embd); + std::vector logits; + const PlacementBackend device_backend = resolved_device_backend(cfg_.device); + const bool disable_cached_decode = + should_disable_ds4_expert_split_cached_decode( + expert_split_runtime_.enabled, + device_backend, + backend_, + cfg_.device.gpu); for (int generated = 0; generated < n_gen; generated++) { if (io.cancelled) break; + if (env_flag_enabled("DFLASH_DS4_TRACE_DECODE")) { + std::fprintf(stderr, + "[deepseek4-trace] decode iter=%d last_logits=%zu out_tokens=%zu\n", + generated, last_logits_.size(), out_tokens.size()); + } + // Budget hook: force-close if remaining budget hits threshold if (!budget_hook.close_token_ids.empty() && (n_gen - generated) <= budget_hook.hard_limit_remaining) { @@ -617,24 +1911,31 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, } // Get last logits and sample - std::vector logits; if (generated == 0 && !last_logits_.empty()) { + if (env_flag_enabled("DFLASH_DS4_TRACE_DECODE")) { + std::fprintf(stderr, + "[deepseek4-trace] decode iter=%d using cached prefill logits=%zu\n", + generated, last_logits_.size()); + } logits = last_logits_; } else { - std::vector embed(w_.n_embd); int32_t tok_to_eval = out_tokens.empty() ? 0 : out_tokens.back(); + if (env_flag_enabled("DFLASH_DS4_TRACE_DECODE")) { + std::fprintf(stderr, + "[deepseek4-trace] decode iter=%d fallback forward tok=%d\n", + generated, tok_to_eval); + } 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())) { + if (!run_step_with_runtime_path(embed.data(), 1, pos, logits, + &tok_to_eval, + true, + disable_cached_decode, + timing ? &step_tel : nullptr)) { std::fprintf(stderr, "[deepseek4] decode step failed\n"); return false; } @@ -658,10 +1959,21 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, if (timing) tel_acc.sample_us += elapsed_us(sample_t0, Clock::now()); } out_tokens.push_back(next_token); + if (env_flag_enabled("DFLASH_DS4_TRACE_DECODE")) { + std::fprintf(stderr, + "[deepseek4-trace] decode iter=%d sampled=%d emit_begin\n", + generated, next_token); + } const auto emit_t0 = Clock::now(); io.emit(next_token); + if (env_flag_enabled("DFLASH_DS4_TRACE_DECODE")) { + std::fprintf(stderr, + "[deepseek4-trace] decode iter=%d sampled=%d emit_done\n", + generated, next_token); + } if (timing) tel_acc.emit_us += elapsed_us(emit_t0, Clock::now()); + // Check EOS using the GGUF tokenizer metadata loaded with the weights. if (deepseek4_is_eos_tok(next_token, w_)) { break; } @@ -675,10 +1987,12 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, const DaemonIO & io) { GenerateResult result; + DaemonIO out_io = io.with_token_callback(req.on_token); auto t0 = Clock::now(); + reset_request_state(); // Prefill - int committed = do_prefill(req.prompt, io); + int committed = do_prefill(req.prompt, out_io); if (committed < 0) { result.error = "prefill"; return result; @@ -686,6 +2000,7 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, result.prefill_s = elapsed_s(t0); if (req.n_gen <= 0) { + out_io.emit(-1); result.ok = true; maybe_save_routing_stats(); return result; @@ -697,12 +2012,13 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, gen_tokens.reserve(req.n_gen); bool forced_close = false; - if (!do_decode(committed, req.n_gen, gen_tokens, io, + if (!do_decode(committed, req.n_gen, gen_tokens, out_io, req.budget_hook, &forced_close)) { result.error = "decode"; return result; } + out_io.emit(-1); result.ok = true; result.tokens = std::move(gen_tokens); result.decode_s = elapsed_s(t1); @@ -767,8 +2083,13 @@ void DeepSeek4Backend::shutdown() { free_deepseek4_snapshot(snapshots_[i]); } free_deepseek4_cache(cache_); - stream_engine_.destroy(); moe_hybrid_.reset(); + expert_runtime_.reset(); + expert_split_runtime_.reset(); + expert_split_state_ = {}; + layer_expert_bytes_.clear(); + hc_state_.clear(); + last_logits_.clear(); routing_stats_.reset(); routing_stats_out_path_.clear(); moe_placement_ = {}; diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 508637867..f83ff4f9c 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -11,7 +11,8 @@ #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 "../common/expert_split_state.h" +#include "../common/moe_expert_compute.h" #include "deepseek4_internal.h" #include "ggml.h" @@ -60,6 +61,15 @@ class DeepSeek4Backend : public ModelBackend { void shutdown() override; private: + bool run_step_with_runtime_path(const float * embed, + int n_tokens, + int kv_start, + std::vector & out_logits, + const int32_t * token_ids, + bool want_logits, + bool disable_cached_decode, + DeepSeek4StepTelemetry * telemetry); + void reset_request_state(); DeepSeek4BackendConfig cfg_; ggml_backend_t backend_ = nullptr; ggml_backend_t snap_backend_ = nullptr; @@ -74,6 +84,7 @@ class DeepSeek4Backend : public ModelBackend { // Snapshots static constexpr int PREFIX_SLOTS = 64; DeepSeek4Snapshot snapshots_[PREFIX_SLOTS]; + std::vector hc_state_; std::vector last_logits_; // Prefill prompt tokens in chunks, return absolute committed position. @@ -92,14 +103,34 @@ class DeepSeek4Backend : public ModelBackend { int max_ctx, MoeHybridPlacement & out, std::string * err) const; + bool load_expert_split_placement(const char * hotness_path, + const DeepSeek4Weights & w, + int max_ctx, + MoeHybridPlacement & out, + std::string * err); + bool build_expert_split_state_from_hotness(const MoeHybridRoutingStats & hotness, + uint64_t expert_budget_bytes, + uint64_t hot_budget_bytes, + int max_hot_per_layer, + const std::vector & configured_targets, + const DeepSeek4Weights & w, + std::string * err); + std::unique_ptr make_expert_split_local_target_compute( + const ExpertSplitComputeTargetRuntime & split_target, + std::string * err) const; + bool init_single_target_expert_runtime(std::string * err); + bool init_expert_split_runtime(std::string * err); 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. + ExpertSplitStateComponents expert_split_state_; + MoeExpertComputeRuntime expert_runtime_; + MoeMultiTargetExpertRuntime expert_split_runtime_; + std::vector layer_expert_bytes_; + uint64_t expert_split_total_budget_bytes_ = 0; + uint64_t expert_split_hot_budget_bytes_ = 0; + int expert_split_cache_slots_ = 0; std::shared_ptr routing_stats_; std::string routing_stats_out_path_; }; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 90be91d85..ddabe28dc 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -11,10 +11,10 @@ #include "deepseek4_internal.h" #include "deepseek4_hc_cuda.h" #include "internal.h" +#include "expert_ipc.h" #include "../common/step_graph.h" #include "../common/moe_hybrid_ffn_eval.h" #include "../common/moe_hybrid_routing_stats.h" -#include "../common/moe_hybrid_stream.h" #include "../common/moe_hybrid_types.h" #include "ggml.h" @@ -28,6 +28,9 @@ #include #include #include +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386) || defined(_M_IX86) +#include +#endif #include #include @@ -46,6 +49,44 @@ static bool ds4_env_flag(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +static size_t ds4_full_prefill_graph_size(int n_tokens) { + if (n_tokens <= 1) { + // The non-hybrid all-hot decode path still builds one full-layer graph + // per token. Recent ggml/operator growth can exceed the old 4096-node + // cap during decode even though prefill remains fine. + return 16384; + } + if (n_tokens <= 128) { + return 65536; + } + return 131072; +} + +// Cached decode subgraphs have grown as DS4 picked up richer route/HC/device- +// resident FFN paths. The older fixed node caps were small enough to trip ggml +// graph-size assertions during first decode on HIP. A larger cgraph only adds +// light CPU-side metadata, so keep generous headroom here. +static constexpr size_t DS4_CACHED_DECODE_FFN_GRAPH_SIZE = 4096; +static constexpr size_t DS4_CACHED_DECODE_ROUTE_GRAPH_SIZE = 4096; +static constexpr size_t DS4_CACHED_DECODE_EXPERT_GRAPH_SIZE = 4096; +static constexpr size_t DS4_CACHED_DECODE_OUTPUT_GRAPH_SIZE = 2048; +static constexpr size_t DS4_CACHED_DECODE_ATTN_GRAPH_SIZE = 4096; +static constexpr size_t DS4_CACHED_DECODE_HC_PRE_GRAPH_SIZE = 4096; +static constexpr size_t DS4_CACHED_DECODE_HC_POST_GRAPH_SIZE = 1024; + +static void ds4_trace_decode_marker(bool enabled, + int kv_start, + int layer, + const char * phase) { + if (!enabled) { + return; + } + std::fprintf(stderr, + "[deepseek4-trace] decode-step kv=%d layer=%d %s\n", + kv_start, layer, phase); + std::fflush(stderr); +} + template static void ds4_parallel_for_tokens(int n_tokens, int min_parallel_tokens, Fn && fn) { if (n_tokens <= min_parallel_tokens) { @@ -55,7 +96,7 @@ static void ds4_parallel_for_tokens(int n_tokens, int min_parallel_tokens, Fn && unsigned nth = std::thread::hardware_concurrency(); if (nth == 0) nth = 4; - if (nth > 8) nth = 8; + if (nth > 16) nth = 16; if ((int)nth <= 1) { fn(0, n_tokens); return; @@ -89,6 +130,7 @@ static void add_ffn_telemetry(DeepSeek4StepTelemetry * dst, dst->ffn_cold_us += src.cold_us; dst->ffn_combine_us += src.combine_us; dst->ffn_partition_us += src.partition_us; + dst->ffn_cache_promote_us += src.cache_promote_us; dst->ffn_hot_graph_builds += src.hot_graph_builds; dst->ffn_hot_graph_hits += src.hot_graph_hits; dst->ffn_cold_graph_builds += src.cold_graph_builds; @@ -97,6 +139,53 @@ static void add_ffn_telemetry(DeepSeek4StepTelemetry * dst, dst->cold_selected += src.cold_selected; } +struct Ds4TensorReadback { + const ggml_tensor * tensor = nullptr; + void * data = nullptr; + size_t size = 0; +}; + +static void ds4_tensor_gets_async_and_sync( + ggml_backend_t backend, + const Ds4TensorReadback * readbacks, + size_t count) { + if (!backend || !readbacks || count == 0) { + return; + } + size_t issued = 0; + for (size_t i = 0; i < count; ++i) { + const Ds4TensorReadback & rb = readbacks[i]; + if (!rb.tensor || !rb.data || rb.size == 0) { + continue; + } + ggml_backend_tensor_get_async( + backend, rb.tensor, rb.data, 0, rb.size); + ++issued; + } + if (issued > 0) { + ggml_backend_synchronize(backend); + } +} + +static const float * ds4_get_or_load_route_bias_host( + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + DeepSeek4LayerCache & lc) { + if (!L.ffn_exp_probs_b) { + return nullptr; + } + if (!lc.route_bias_loaded || (int) lc.route_bias_host.size() != w.n_expert) { + lc.route_bias_host.resize((size_t) w.n_expert); + ggml_backend_tensor_get( + L.ffn_exp_probs_b, + lc.route_bias_host.data(), + 0, + sizeof(float) * (size_t) w.n_expert); + lc.route_bias_loaded = true; + } + return lc.route_bias_host.data(); +} + } // namespace struct DeepSeek4I32InputBinding { @@ -157,6 +246,65 @@ struct DeepSeek4CachedDecodeFfnGraph { } }; +struct DeepSeek4CachedDecodeRouteGraph { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + int layer_idx = -1; + int n_tokens = 0; + bool hash_routed = false; + StepGraph sg; + ggml_tensor * ffn_normed = nullptr; + ggml_tensor * selected = nullptr; + ggml_tensor * weights = nullptr; + ggml_tensor * hash_ids = nullptr; + + bool valid() const { + return owner_ctx && backend && layer_idx >= 0 && n_tokens > 0 && + sg.ctx && sg.gf && sg.alloc && sg.inp_embed && + ffn_normed && weights && + (hash_routed ? hash_ids != nullptr : selected != nullptr); + } + + void free() { + ffn_normed = nullptr; + selected = nullptr; + weights = nullptr; + hash_ids = nullptr; + step_graph_destroy(sg); + owner_ctx = nullptr; + backend = nullptr; + layer_idx = -1; + n_tokens = 0; + hash_routed = false; + } +}; + +struct DeepSeek4CachedDecodeExpertGraph { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + int layer_idx = -1; + int n_tokens = 0; + StepGraph sg; + ggml_tensor * selected = nullptr; + ggml_tensor * weights = nullptr; + + bool valid() const { + return owner_ctx && backend && layer_idx >= 0 && n_tokens > 0 && + sg.ctx && sg.gf && sg.alloc && sg.inp_embed && + sg.hidden_states && selected && weights; + } + + void free() { + selected = nullptr; + weights = nullptr; + step_graph_destroy(sg); + owner_ctx = nullptr; + backend = nullptr; + layer_idx = -1; + n_tokens = 0; + } +}; + struct DeepSeek4CachedDecodeOutputGraph { const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; @@ -176,16 +324,17 @@ struct DeepSeek4CachedDecodeOutputGraph { } }; -struct DeepSeek4LegacyFullStepCache { +struct DeepSeek4CachedDecodeGpuFfnState { const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; - StepGraph sg; - std::vector meta_arena; + GpuResidentState state; + + bool valid() const { + return owner_ctx && backend && state.valid(); + } void free() { - step_graph_destroy(sg); - meta_arena.clear(); - meta_arena.shrink_to_fit(); + state.destroy(); owner_ctx = nullptr; backend = nullptr; } @@ -213,6 +362,7 @@ struct DeepSeek4CachedDecodeAttnGraph { int n_raw = 0; int n_comp_attn = 0; int n_index_comp = 0; + bool flush_boundary = false; bool compressed = false; bool indexed = false; StepGraph sg; @@ -239,6 +389,7 @@ struct DeepSeek4CachedDecodeAttnGraph { n_raw = 0; n_comp_attn = 0; n_index_comp = 0; + flush_boundary = false; compressed = false; indexed = false; } @@ -263,12 +414,25 @@ struct DeepSeek4CachedLayerAlloc { } }; +struct DeepSeek4LegacyFullStepCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + StepGraph sg; + + void free() { + step_graph_destroy(sg); + owner_ctx = nullptr; + backend = nullptr; + } +}; + struct DeepSeek4LayerRangeScratch { const ggml_context * owner_ctx = nullptr; int n_tokens = 0; int n_embd = 0; int n_hc = 0; int n_expert_used = 0; + int n_expert = 0; std::vector cur; std::vector ffn_working; std::vector hc_post; @@ -277,18 +441,25 @@ struct DeepSeek4LayerRangeScratch { std::vector attn_out_host; std::vector ffn_out_host; std::vector final_embd; + std::vector ffn_normed_host; + std::vector probs_host; + std::vector weights_host; + std::vector bias_host; std::vector hash_expert_ids; + std::vector selected_host; void ensure(const ggml_context * ctx, int tokens, int embd, int hc, - int expert_used) { + int expert_used, + int expert_count) { owner_ctx = ctx; n_tokens = tokens; n_embd = embd; n_hc = hc; n_expert_used = expert_used; + n_expert = expert_count; const size_t embd_count = (size_t) tokens * (size_t) embd; const size_t hc_count = embd_count * (size_t) hc; cur.resize(embd_count); @@ -299,7 +470,12 @@ struct DeepSeek4LayerRangeScratch { attn_out_host.resize(embd_count); ffn_out_host.resize(embd_count); final_embd.resize(embd_count); + ffn_normed_host.resize(embd_count); + probs_host.resize((size_t) tokens * (size_t) expert_count); + weights_host.resize((size_t) tokens * (size_t) expert_used); + bias_host.resize((size_t) expert_count); hash_expert_ids.resize((size_t) tokens * (size_t) expert_used); + selected_host.resize((size_t) tokens * (size_t) expert_used); } }; @@ -325,7 +501,8 @@ static bool build_cached_decode_ffn_graph( out.sg.inp_embed = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens); ggml_set_input(out.sg.inp_embed); - out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 2048, false); + out.sg.gf = ggml_new_graph_custom( + out.sg.ctx, DS4_CACHED_DECODE_FFN_GRAPH_SIZE, false); ggml_tensor * ffn_normed = build_rms_norm(out.sg.ctx, out.sg.inp_embed, L.ffn_norm, w.rms_eps); ggml_tensor * ffn_out = nullptr; @@ -392,6 +569,90 @@ static bool build_cached_decode_ffn_graph( return true; } +static bool build_cached_decode_route_graph( + DeepSeek4CachedDecodeRouteGraph & out, + ggml_backend_t backend, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int layer_idx, + int n_tokens, + bool hash_routed) { + out.free(); + + const size_t ctx_size = 16 * 1024 * 1024; + ggml_init_params params{}; + params.mem_size = ctx_size; + params.mem_buffer = nullptr; + params.no_alloc = true; + out.sg.ctx = ggml_init(params); + if (!out.sg.ctx) { + return false; + } + + out.sg.inp_embed = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens); + ggml_set_input(out.sg.inp_embed); + out.sg.gf = ggml_new_graph_custom( + out.sg.ctx, DS4_CACHED_DECODE_ROUTE_GRAPH_SIZE, false); + + out.ffn_normed = build_rms_norm(out.sg.ctx, out.sg.inp_embed, L.ffn_norm, w.rms_eps); + ggml_set_output(out.ffn_normed); + ggml_build_forward_expand(out.sg.gf, out.ffn_normed); + + ggml_tensor * logits = ggml_mul_mat(out.sg.ctx, L.ffn_gate_inp, out.ffn_normed); + ggml_tensor * probs = ggml_sqrt(out.sg.ctx, ggml_softplus(out.sg.ctx, logits)); + + if (hash_routed) { + out.hash_ids = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I32, w.n_expert_used, n_tokens); + ggml_set_input(out.hash_ids); + + ggml_tensor * probs_3d = ggml_reshape_3d(out.sg.ctx, probs, 1, w.n_expert, n_tokens); + out.weights = ggml_get_rows(out.sg.ctx, probs_3d, out.hash_ids); + out.weights = ggml_reshape_2d(out.sg.ctx, out.weights, w.n_expert_used, n_tokens); + + ggml_tensor * w_sum = ggml_sum_rows(out.sg.ctx, out.weights); + w_sum = ggml_clamp(out.sg.ctx, w_sum, 6.103515625e-5f, INFINITY); + out.weights = ggml_div(out.sg.ctx, out.weights, w_sum); + if (w.expert_weight_scale != 1.0f) { + out.weights = ggml_scale(out.sg.ctx, out.weights, w.expert_weight_scale); + } + } else { + ggml_tensor * selection = probs; + if (L.ffn_exp_probs_b) { + selection = ggml_add(out.sg.ctx, selection, L.ffn_exp_probs_b); + } + + out.selected = ggml_top_k(out.sg.ctx, selection, w.n_expert_used); + ggml_tensor * probs_3d = ggml_reshape_3d(out.sg.ctx, probs, 1, w.n_expert, n_tokens); + out.weights = ggml_get_rows(out.sg.ctx, probs_3d, out.selected); + out.weights = ggml_reshape_2d(out.sg.ctx, out.weights, w.n_expert_used, n_tokens); + + ggml_tensor * w_sum = ggml_sum_rows(out.sg.ctx, out.weights); + w_sum = ggml_clamp(out.sg.ctx, w_sum, 6.103515625e-5f, INFINITY); + out.weights = ggml_div(out.sg.ctx, out.weights, w_sum); + if (w.expert_weight_scale != 1.0f) { + out.weights = ggml_scale(out.sg.ctx, out.weights, w.expert_weight_scale); + } + ggml_set_output(out.selected); + ggml_build_forward_expand(out.sg.gf, out.selected); + } + + ggml_set_output(out.weights); + ggml_build_forward_expand(out.sg.gf, out.weights); + + out.sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(out.sg.alloc, out.sg.gf)) { + out.free(); + return false; + } + + out.owner_ctx = w.ctx; + out.backend = backend; + out.layer_idx = layer_idx; + out.n_tokens = n_tokens; + out.hash_routed = hash_routed; + return true; +} + static bool build_cached_decode_output_graph( DeepSeek4CachedDecodeOutputGraph & out, ggml_backend_t backend, @@ -409,12 +670,34 @@ static bool build_cached_decode_output_graph( return false; } - out.sg.hidden_input = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens); + const bool use_output_hc = + w.output_hc_fn && w.output_hc_scale && w.output_hc_base; + const int input_dim = use_output_hc ? (w.n_hc * w.n_embd) : w.n_embd; + out.sg.hidden_input = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, input_dim, n_tokens); ggml_set_input(out.sg.hidden_input); - ggml_tensor * normed = build_rms_norm(out.sg.ctx, out.sg.hidden_input, w.out_norm, w.rms_eps); + ggml_tensor * final_hidden = out.sg.hidden_input; + if (use_output_hc) { + ggml_tensor * flat = ggml_rms_norm(out.sg.ctx, out.sg.hidden_input, w.hc_eps); + ggml_tensor * pre = ggml_mul_mat(out.sg.ctx, w.output_hc_fn, flat); + ggml_tensor * scale0 = ggml_reshape_2d(out.sg.ctx, + ggml_view_1d(out.sg.ctx, w.output_hc_scale, 1, 0), 1, 1); + ggml_tensor * pre_scaled = ggml_mul(out.sg.ctx, pre, ggml_repeat(out.sg.ctx, scale0, pre)); + ggml_tensor * base = ggml_reshape_2d(out.sg.ctx, w.output_hc_base, w.n_hc, 1); + ggml_tensor * hc_weights = ggml_sigmoid(out.sg.ctx, + ggml_add(out.sg.ctx, pre_scaled, base)); + hc_weights = ggml_clamp(out.sg.ctx, hc_weights, 1.0e-6f, INFINITY); + ggml_tensor * hc_weights_3d = ggml_reshape_3d(out.sg.ctx, hc_weights, 1, w.n_hc, n_tokens); + ggml_tensor * hc_state_3d = ggml_reshape_3d(out.sg.ctx, out.sg.hidden_input, w.n_embd, w.n_hc, n_tokens); + ggml_tensor * weighted = ggml_mul(out.sg.ctx, hc_state_3d, hc_weights_3d); + weighted = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, weighted, 1, 0, 2, 3)); + weighted = ggml_sum_rows(out.sg.ctx, weighted); + final_hidden = ggml_reshape_2d(out.sg.ctx, weighted, w.n_embd, n_tokens); + } + ggml_tensor * normed = build_rms_norm(out.sg.ctx, final_hidden, w.out_norm, w.rms_eps); out.sg.logits = ggml_mul_mat(out.sg.ctx, w.output, normed); ggml_set_output(out.sg.logits); - out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 1024, false); + out.sg.gf = ggml_new_graph_custom( + out.sg.ctx, DS4_CACHED_DECODE_OUTPUT_GRAPH_SIZE, false); ggml_build_forward_expand(out.sg.gf, out.sg.logits); out.sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); @@ -429,6 +712,79 @@ static bool build_cached_decode_output_graph( return true; } +static bool build_cached_decode_expert_graph( + DeepSeek4CachedDecodeExpertGraph & out, + ggml_backend_t backend, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int layer_idx, + int n_tokens) { + out.free(); + + const size_t ctx_size = 16 * 1024 * 1024; + ggml_init_params params{}; + params.mem_size = ctx_size; + params.mem_buffer = nullptr; + params.no_alloc = true; + out.sg.ctx = ggml_init(params); + if (!out.sg.ctx) { + return false; + } + + out.sg.inp_embed = ggml_new_tensor_2d( + out.sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens); + ggml_set_input(out.sg.inp_embed); + out.selected = ggml_new_tensor_2d( + out.sg.ctx, GGML_TYPE_I32, w.n_expert_used, n_tokens); + ggml_set_input(out.selected); + out.weights = ggml_new_tensor_2d( + out.sg.ctx, GGML_TYPE_F32, w.n_expert_used, n_tokens); + ggml_set_input(out.weights); + out.sg.gf = ggml_new_graph_custom( + out.sg.ctx, DS4_CACHED_DECODE_EXPERT_GRAPH_SIZE, false); + + ggml_tensor * shared_out = build_shared_ffn(out.sg.ctx, out.sg.inp_embed, w, L); + ggml_tensor * cur_3d = ggml_reshape_3d( + out.sg.ctx, out.sg.inp_embed, w.n_embd, 1, n_tokens); + ggml_tensor * gate_e = + ggml_mul_mat_id(out.sg.ctx, L.ffn_gate_exps, cur_3d, out.selected); + ggml_tensor * up_e = + ggml_mul_mat_id(out.sg.ctx, L.ffn_up_exps, cur_3d, out.selected); + gate_e = ggml_reshape_3d( + out.sg.ctx, gate_e, w.n_ff_exp, w.n_expert_used, n_tokens); + up_e = ggml_reshape_3d( + out.sg.ctx, up_e, w.n_ff_exp, w.n_expert_used, n_tokens); + ggml_tensor * mid_e = build_clamped_swiglu( + out.sg.ctx, gate_e, up_e, w.swiglu_clamp_exp); + ggml_tensor * down_e = + ggml_mul_mat_id(out.sg.ctx, L.ffn_down_exps, mid_e, out.selected); + down_e = ggml_reshape_3d( + out.sg.ctx, down_e, w.n_embd, w.n_expert_used, n_tokens); + + ggml_tensor * weights_3d = ggml_reshape_3d( + out.sg.ctx, out.weights, 1, w.n_expert_used, n_tokens); + ggml_tensor * routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); + routed_out = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(out.sg.ctx, routed_out); + routed_out = ggml_reshape_2d(out.sg.ctx, routed_out, w.n_embd, n_tokens); + + out.sg.hidden_states = ggml_add(out.sg.ctx, shared_out, routed_out); + ggml_set_output(out.sg.hidden_states); + ggml_build_forward_expand(out.sg.gf, out.sg.hidden_states); + + out.sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(out.sg.alloc, out.sg.gf)) { + out.free(); + return false; + } + + out.owner_ctx = w.ctx; + out.backend = backend; + out.layer_idx = layer_idx; + out.n_tokens = n_tokens; + return true; +} + // ─── Helper: RMSNorm ──────────────────────────────────────────────────── static ggml_tensor * build_rms_norm(ggml_context * ctx, ggml_tensor * x, @@ -1076,6 +1432,7 @@ struct DeepSeek4CachedDecodeHcPreGraph { const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; int layer_idx = -1; + int n_tokens = 0; bool ffn = false; StepGraph sg; ggml_tensor * post = nullptr; @@ -1083,6 +1440,7 @@ struct DeepSeek4CachedDecodeHcPreGraph { bool valid() const { return owner_ctx && backend && layer_idx >= 0 && + n_tokens > 0 && sg.ctx && sg.gf && sg.alloc && sg.inp_embed && sg.hidden_states && post && comb; } @@ -1092,6 +1450,7 @@ struct DeepSeek4CachedDecodeHcPreGraph { owner_ctx = nullptr; backend = nullptr; layer_idx = -1; + n_tokens = 0; ffn = false; post = nullptr; comb = nullptr; @@ -1101,6 +1460,7 @@ struct DeepSeek4CachedDecodeHcPreGraph { struct DeepSeek4CachedDecodeHcPostGraph { const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; + int n_tokens = 0; StepGraph sg; ggml_tensor * residual_hc = nullptr; ggml_tensor * block_out = nullptr; @@ -1109,6 +1469,7 @@ struct DeepSeek4CachedDecodeHcPostGraph { bool valid() const { return owner_ctx && backend && + n_tokens > 0 && sg.ctx && sg.gf && sg.alloc && sg.hidden_states && residual_hc && block_out && post && comb; } @@ -1117,6 +1478,7 @@ struct DeepSeek4CachedDecodeHcPostGraph { step_graph_destroy(sg); owner_ctx = nullptr; backend = nullptr; + n_tokens = 0; residual_hc = nullptr; block_out = nullptr; post = nullptr; @@ -1134,7 +1496,8 @@ static bool build_cached_decode_attn_graph( int kv_start, int raw_attn_count, int comp_attn_count, - int index_comp_count) { + int index_comp_count, + bool flush_boundary) { out.free(); const size_t ctx_size = 48 * 1024 * 1024; @@ -1152,12 +1515,14 @@ static bool build_cached_decode_attn_graph( out.n_raw = raw_attn_count; out.n_comp_attn = comp_attn_count; out.n_index_comp = index_comp_count; + out.flush_boundary = flush_boundary; out.compressed = ratio > 0; out.indexed = ratio == 4; out.sg.inp_embed = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_embd, 1); ggml_set_input(out.sg.inp_embed); - out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 2048, false); + out.sg.gf = ggml_new_graph_custom( + out.sg.ctx, DS4_CACHED_DECODE_ATTN_GRAPH_SIZE, false); out.inputs.rope_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); out.inputs.neg_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); @@ -1304,7 +1669,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, return false; } const bool can_use_device_params = - ds4_backend_is_cuda(backend) && + ds4_backend_is_gpu(backend) && scale_tensor && base_tensor && scale_tensor->data && base_tensor->data && scale_tensor->buffer && base_tensor->buffer && @@ -1362,7 +1727,8 @@ static bool build_cached_decode_hc_pre_graph( const DeepSeek4Layer & L, const float * scale_data, int layer_idx, - bool ffn) { + bool ffn, + int n_tokens = 1) { out.free(); const size_t ctx_size = 4 * 1024 * 1024; @@ -1379,20 +1745,23 @@ static bool build_cached_decode_hc_pre_graph( ggml_tensor * hc_fn = ffn ? L.hc_ffn_fn : L.hc_attn_fn; ggml_tensor * hc_base = ffn ? L.hc_ffn_base : L.hc_attn_base; - out.sg.inp_embed = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, hc_dim, 1); + out.sg.inp_embed = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, hc_dim, n_tokens); ggml_set_input(out.sg.inp_embed); - out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 2048, false); + out.sg.gf = ggml_new_graph_custom( + out.sg.ctx, DS4_CACHED_DECODE_HC_PRE_GRAPH_SIZE, false); ggml_tensor * flat = ggml_rms_norm(out.sg.ctx, out.sg.inp_embed, w.hc_eps); ggml_tensor * mix = ggml_mul_mat(out.sg.ctx, hc_fn, flat); - ggml_tensor * pre_mix = ggml_reshape_2d(out.sg.ctx, - ggml_view_1d(out.sg.ctx, mix, w.n_hc, 0), w.n_hc, 1); - ggml_tensor * post_mix = ggml_reshape_2d(out.sg.ctx, - ggml_view_1d(out.sg.ctx, mix, w.n_hc, (size_t) w.n_hc * mix->nb[0]), w.n_hc, 1); - ggml_tensor * comb_mix = ggml_reshape_2d(out.sg.ctx, - ggml_view_1d(out.sg.ctx, mix, w.n_hc * w.n_hc, (size_t) (2 * w.n_hc) * mix->nb[0]), - w.n_hc, w.n_hc); + ggml_tensor * pre_mix = ggml_cont(out.sg.ctx, ggml_view_2d(out.sg.ctx, mix, + w.n_hc, n_tokens, mix->nb[1], 0)); + ggml_tensor * post_mix = ggml_cont(out.sg.ctx, ggml_view_2d(out.sg.ctx, mix, + w.n_hc, n_tokens, mix->nb[1], (size_t) w.n_hc * mix->nb[0])); + ggml_tensor * comb_mix = ggml_cont(out.sg.ctx, ggml_view_3d(out.sg.ctx, mix, + w.n_hc, w.n_hc, n_tokens, + (size_t) w.n_hc * mix->nb[0], + mix->nb[1], + (size_t) (2 * w.n_hc) * mix->nb[0])); ggml_tensor * pre_base = ggml_reshape_2d(out.sg.ctx, ggml_view_1d(out.sg.ctx, hc_base, w.n_hc, 0), w.n_hc, 1); @@ -1405,17 +1774,17 @@ static bool build_cached_decode_hc_pre_graph( ggml_tensor * pre = ggml_sigmoid(out.sg.ctx, ggml_add(out.sg.ctx, ggml_scale(out.sg.ctx, pre_mix, scale_data[0]), - pre_base)); + ggml_repeat(out.sg.ctx, pre_base, pre_mix))); ggml_tensor * post = ggml_scale(out.sg.ctx, ggml_sigmoid(out.sg.ctx, ggml_add(out.sg.ctx, ggml_scale(out.sg.ctx, post_mix, scale_data[1]), - post_base)), + ggml_repeat(out.sg.ctx, post_base, post_mix))), 2.0f); ggml_tensor * comb = ggml_add(out.sg.ctx, ggml_scale(out.sg.ctx, comb_mix, scale_data[2]), - comb_base); + ggml_repeat(out.sg.ctx, comb_base, comb_mix)); comb = ggml_soft_max(out.sg.ctx, comb); comb = ds4_hc_col_normalize(out.sg.ctx, comb); for (int iter = 1; iter < w.n_hc_sinkhorn_iter; ++iter) { @@ -1423,9 +1792,11 @@ static bool build_cached_decode_hc_pre_graph( comb = ds4_hc_col_normalize(out.sg.ctx, comb); } - ggml_tensor * hc_state_2d = ggml_reshape_2d(out.sg.ctx, out.sg.inp_embed, w.n_embd, w.n_hc); - ggml_tensor * hc_state_t = ggml_cont(out.sg.ctx, ggml_transpose(out.sg.ctx, hc_state_2d)); - ggml_tensor * working = ggml_mul_mat(out.sg.ctx, hc_state_t, pre); + ggml_tensor * hc_state_3d = ggml_reshape_3d(out.sg.ctx, out.sg.inp_embed, w.n_embd, w.n_hc, n_tokens); + ggml_tensor * hc_state_t = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, hc_state_3d, 1, 0, 2, 3)); + ggml_tensor * pre_3d = ggml_reshape_3d(out.sg.ctx, pre, w.n_hc, 1, n_tokens); + ggml_tensor * working = ggml_mul_mat(out.sg.ctx, hc_state_t, pre_3d); + working = ggml_reshape_2d(out.sg.ctx, working, w.n_embd, n_tokens); out.sg.hidden_states = working; out.post = post; @@ -1446,6 +1817,7 @@ static bool build_cached_decode_hc_pre_graph( out.owner_ctx = w.ctx; out.backend = backend; out.layer_idx = layer_idx; + out.n_tokens = n_tokens; out.ffn = ffn; return true; } @@ -1453,7 +1825,8 @@ static bool build_cached_decode_hc_pre_graph( static bool build_cached_decode_hc_post_graph( DeepSeek4CachedDecodeHcPostGraph & out, ggml_backend_t backend, - const DeepSeek4Weights & w) { + const DeepSeek4Weights & w, + int n_tokens = 1) { out.free(); const size_t ctx_size = 2 * 1024 * 1024; @@ -1467,29 +1840,30 @@ static bool build_cached_decode_hc_post_graph( } const int hc_dim = w.n_embd * w.n_hc; - out.residual_hc = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, hc_dim, 1); - out.block_out = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_embd, 1); - out.post = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_hc, 1); - out.comb = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_hc, w.n_hc); + out.residual_hc = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, hc_dim, n_tokens); + out.block_out = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens); + out.post = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_F32, w.n_hc, n_tokens); + out.comb = ggml_new_tensor_3d(out.sg.ctx, GGML_TYPE_F32, w.n_hc, w.n_hc, n_tokens); ggml_set_input(out.residual_hc); ggml_set_input(out.block_out); ggml_set_input(out.post); ggml_set_input(out.comb); - out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 256, false); + out.sg.gf = ggml_new_graph_custom( + out.sg.ctx, DS4_CACHED_DECODE_HC_POST_GRAPH_SIZE, false); - ggml_tensor * residual_2d = ggml_reshape_2d(out.sg.ctx, out.residual_hc, w.n_embd, w.n_hc); - ggml_tensor * residual_t = ggml_cont(out.sg.ctx, ggml_transpose(out.sg.ctx, residual_2d)); - ggml_tensor * comb_t = ggml_cont(out.sg.ctx, ggml_transpose(out.sg.ctx, out.comb)); + ggml_tensor * residual_3d = ggml_reshape_3d(out.sg.ctx, out.residual_hc, w.n_embd, w.n_hc, n_tokens); + ggml_tensor * residual_t = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, residual_3d, 1, 0, 2, 3)); + ggml_tensor * comb_t = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, out.comb, 1, 0, 2, 3)); ggml_tensor * mixed_t = ggml_mul_mat(out.sg.ctx, comb_t, residual_t); - ggml_tensor * mixed = ggml_cont(out.sg.ctx, ggml_transpose(out.sg.ctx, mixed_t)); - ggml_tensor * post_t = ggml_cont(out.sg.ctx, ggml_transpose(out.sg.ctx, out.post)); - ggml_tensor * block_rep = ggml_repeat(out.sg.ctx, out.block_out, mixed); - ggml_tensor * post_rep = ggml_repeat(out.sg.ctx, post_t, mixed); - out.sg.hidden_states = ggml_reshape_2d( - out.sg.ctx, + ggml_tensor * mixed = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, mixed_t, 1, 0, 2, 3)); + ggml_tensor * block_3d = ggml_reshape_3d(out.sg.ctx, out.block_out, w.n_embd, 1, n_tokens); + ggml_tensor * post_3d = ggml_reshape_3d(out.sg.ctx, out.post, 1, w.n_hc, n_tokens); + ggml_tensor * block_rep = ggml_repeat(out.sg.ctx, block_3d, mixed); + ggml_tensor * post_rep = ggml_repeat(out.sg.ctx, post_3d, mixed); + out.sg.hidden_states = ggml_reshape_2d(out.sg.ctx, ggml_add(out.sg.ctx, mixed, ggml_mul(out.sg.ctx, block_rep, post_rep)), - hc_dim, 1); + hc_dim, n_tokens); ggml_set_output(out.sg.hidden_states); ggml_build_forward_expand(out.sg.gf, out.sg.hidden_states); @@ -1502,6 +1876,7 @@ static bool build_cached_decode_hc_post_graph( out.owner_ctx = w.ctx; out.backend = backend; + out.n_tokens = n_tokens; return true; } @@ -1522,6 +1897,8 @@ static MoeHybridConfig make_ds4_moe_hybrid_config(const DeepSeek4Weights & w) { cfg.n_layer = w.n_layer; cfg.first_moe_layer = 0; cfg.swiglu_clamp = w.swiglu_clamp_exp; + static const int sm = query_gpu_compute_sm(); + cfg.mmq_safe_full_batch = (sm >= 80); return cfg; } @@ -1549,148 +1926,341 @@ static ggml_tensor * build_shared_ffn( return ggml_mul_mat(ctx, L.ffn_down_shexp, mid_sh); } -static bool eval_ds4_hybrid( +static bool eval_ds4_hybrid_or_worker( ggml_backend_t backend, ggml_backend_t cpu_backend, const MoeHybridConfig & hybrid_cfg, const MoeLayerDesc & desc, - const MoeHybridStorage * hybrid_owner, MoeHybridLayerStorage & storage, - MoeHybridStreamEngine * stream_engine, + ExpertIpcClient * expert_worker, int layer, int n_embd, int n_expert_used, const float * ffn_normed_host, + ggml_tensor * ffn_normed_backend, const int32_t * selected_host, const float * weights_host, int n_tokens, std::vector & ffn_out_host, ggml_gallocr_t * hot_alloc, ggml_gallocr_t * cold_alloc, - DeepSeek4StepTelemetry * step_tel) { + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layer, + bool worker_owns_hot_ids, + DeepSeek4StepTelemetry * step_tel, + ggml_tensor ** ffn_out_backend = nullptr) { const auto ffn_t0 = Ds4TimingClock::now(); - if (!storage.down_cold && !storage.gate_up_cold) { - if (!hybrid_owner || !stream_engine || !stream_engine->is_ready() || - !hybrid_owner->has_mmap() || - layer < 0 || layer >= (int) hybrid_owner->layer_regions.size()) { - std::fprintf(stderr, - "[deepseek4] layer %d requires cold-expert streaming but it is unavailable\n", - layer); - return false; - } - - const LayerExpertRegions & regions = hybrid_owner->layer_regions[(size_t) layer]; - ffn_out_host.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); - std::vector hot_selected; - std::vector hot_weights; - std::vector hot_out; - std::vector cold_out; - for (int ti = 0; ti < n_tokens; ++ti) { - const float * token_inp = ffn_normed_host + (size_t)ti * (size_t)n_embd; - const int32_t * token_selected = selected_host + (size_t)ti * (size_t)n_expert_used; - const float * token_weights = weights_host + (size_t)ti * (size_t)n_expert_used; - hot_selected.clear(); - hot_weights.clear(); - bool has_cold = false; - for (int ei = 0; ei < n_expert_used; ++ei) { - const int32_t gid = token_selected[ei]; + if (ffn_out_backend) { + *ffn_out_backend = nullptr; + } + const bool trace_decode = + ds4_env_flag("DFLASH_DS4_TRACE_DECODE") && n_tokens == 1; + const bool use_worker = expert_worker && expert_worker->active() && + (worker_owns_hot_ids || (!storage.down_cold && !storage.gate_up_cold)); + if (!use_worker) { + const auto all_selected_hot_for_batch = [&]() { + const int total_selected = n_tokens * n_expert_used; + for (int i = 0; i < total_selected; ++i) { + const int32_t gid = selected_host[i]; if (gid < 0 || gid >= (int32_t) storage.hot_local_by_global.size()) { - std::fprintf(stderr, - "[deepseek4] layer %d selected expert id out of range: %d\n", - layer, (int) gid); return false; } - if (storage.hot_local_by_global[(size_t)gid] >= 0) { - hot_selected.push_back(gid); - hot_weights.push_back(token_weights[ei]); - } else { - has_cold = true; + if (storage.hot_local_by_global[(size_t) gid] < 0) { + return false; + } + } + return true; + }; + const bool all_selected_hot = + hybrid_cfg.mmq_safe_full_batch && + n_tokens > 1 && + all_selected_hot_for_batch(); + if (all_selected_hot && ffn_normed_host && !ffn_out_backend) { + const auto hot_only_t0 = Ds4TimingClock::now(); + if (eval_moe_hot_only_batched( + backend, hybrid_cfg, desc, storage, + ffn_normed_host, selected_host, weights_host, + n_tokens, ffn_out_host, nullptr, hot_alloc)) { + if (step_tel) { + MoeHybridFfnTelemetry hot_only_tel; + hot_only_tel.ffn_wall_us = ds4_elapsed_us(hot_only_t0, Ds4TimingClock::now()); + hot_only_tel.hot_us = hot_only_tel.ffn_wall_us; + hot_only_tel.hot_selected = n_tokens * n_expert_used; + step_tel->ffn_eval_us += hot_only_tel.ffn_wall_us; + add_ffn_telemetry(step_tel, hot_only_tel); } + return true; + } + // Fall back to the hybrid batched path if hot-only evaluation fails. + } + if (trace_decode) { + std::fprintf(stderr, + "[deepseek4-trace] decode-step layer=%d batched_hybrid begin\n", + layer); + } + MoeHybridFfnTelemetry ffn_tel; + bool ffn_ok = eval_moe_hybrid_ffn_batched( + backend, cpu_backend, hybrid_cfg, desc, storage, + ffn_normed_host, selected_host, weights_host, + n_tokens, ffn_out_host, nullptr, hot_alloc, cold_alloc, + expert_compute, expert_layer, + step_tel ? &ffn_tel : nullptr, + ffn_normed_backend); + if (ffn_ok) { + if (trace_decode) { + std::fprintf(stderr, + "[deepseek4-trace] decode-step layer=%d batched_hybrid ok\n", + layer); + } + if (step_tel) { + step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); + add_ffn_telemetry(step_tel, ffn_tel); } + return true; + } + + if (trace_decode) { + std::fprintf(stderr, + "[deepseek4-trace] decode-step layer=%d batched_hybrid fallback_single\n", + layer); + } - std::string err; + if (!ffn_normed_host) { + return false; + } + ffn_out_host.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); + std::vector single_out; + for (int ti = 0; ti < n_tokens; ++ti) { MoeHybridFfnTelemetry single_tel; if (!eval_moe_hybrid_ffn_single( backend, hybrid_cfg, desc, storage, cpu_backend, - token_inp, - hot_selected.empty() ? nullptr : hot_selected.data(), - hot_weights.empty() ? nullptr : hot_weights.data(), - (int) hot_selected.size(), - hot_out, - step_tel ? &single_tel : nullptr, - &err)) { - std::fprintf(stderr, - "[deepseek4] layer %d hot/shared eval failed: %s\n", - layer, err.c_str()); + ffn_normed_host + (size_t)ti * (size_t)n_embd, + selected_host + (size_t)ti * (size_t)n_expert_used, + weights_host + (size_t)ti * (size_t)n_expert_used, + n_expert_used, single_out, + step_tel ? &single_tel : nullptr)) { return false; } add_ffn_telemetry(step_tel, single_tel); - - cold_out.assign((size_t)n_embd, 0.0f); - if (has_cold) { - if (!eval_moe_cold_experts_streaming( - *stream_engine, backend, - hybrid_owner->mmap_data, hybrid_owner->mmap_size, - hybrid_cfg, desc, regions, storage, - token_inp, token_selected, token_weights, 1, - cold_out, &err)) { - std::fprintf(stderr, - "[deepseek4] layer %d cold streaming eval failed: %s\n", - layer, err.c_str()); - return false; - } - } - - float * dst = ffn_out_host.data() + (size_t)ti * (size_t)n_embd; - for (int i = 0; i < n_embd; ++i) { - dst[i] = hot_out[(size_t)i] + cold_out[(size_t)i]; - } + std::memcpy(ffn_out_host.data() + (size_t)ti * (size_t)n_embd, + single_out.data(), sizeof(float) * (size_t)n_embd); } if (step_tel) step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); return true; } - MoeHybridFfnTelemetry ffn_tel; - bool ffn_ok = eval_moe_hybrid_ffn_batched( - backend, cpu_backend, hybrid_cfg, desc, storage, - ffn_normed_host, selected_host, weights_host, - n_tokens, ffn_out_host, nullptr, hot_alloc, cold_alloc, - nullptr, nullptr, - step_tel ? &ffn_tel : nullptr); - if (ffn_ok) { - if (step_tel) { - step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); - add_ffn_telemetry(step_tel, ffn_tel); - } - return true; - } - ffn_out_host.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); std::vector single_out; + std::vector worker_out; + std::vector local_ids; + std::vector local_weights; + std::vector remote_ids; + std::vector remote_weights; + const bool async_worker = false; for (int ti = 0; ti < n_tokens; ++ti) { - MoeHybridFfnTelemetry single_tel; - if (!eval_moe_hybrid_ffn_single( + local_ids.clear(); + local_weights.clear(); + remote_ids.clear(); + remote_weights.clear(); + const int32_t * ids = selected_host + (size_t)ti * (size_t)n_expert_used; + const float * weights = weights_host + (size_t)ti * (size_t)n_expert_used; + for (int ei = 0; ei < n_expert_used; ++ei) { + const int32_t gid = ids[ei]; + if (gid < 0 || gid >= (int32_t)storage.hot_local_by_global.size()) return false; + const bool is_hot_id = storage.hot_local_by_global[(size_t)gid] >= 0; + if (worker_owns_hot_ids ? !is_hot_id : is_hot_id) { + local_ids.push_back(gid); + local_weights.push_back(weights[ei]); + } else { + remote_ids.push_back(gid); + remote_weights.push_back(weights[ei]); + } + } + + float * dst = ffn_out_host.data() + (size_t)ti * (size_t)n_embd; + MoeHybridFfnTelemetry local_tel; + ExpertIpcTiming ipc_timing; + ExpertIpcClient::PendingEval pending; + bool worker_pending = false; + Ds4TimingClock::time_point worker_t0{}; + if (!remote_ids.empty() && async_worker) { + worker_t0 = Ds4TimingClock::now(); + if (!expert_worker->eval_begin(layer, 1, n_embd, (int)remote_ids.size(), + ffn_normed_host + (size_t)ti * (size_t)n_embd, + remote_ids.data(), remote_weights.data(), + pending, step_tel ? &ipc_timing : nullptr)) { + return false; + } + worker_pending = true; + } + + const bool local_ok = eval_moe_hybrid_ffn_single( backend, hybrid_cfg, desc, storage, cpu_backend, ffn_normed_host + (size_t)ti * (size_t)n_embd, - selected_host + (size_t)ti * (size_t)n_expert_used, - weights_host + (size_t)ti * (size_t)n_expert_used, - n_expert_used, single_out, - step_tel ? &single_tel : nullptr)) { + local_ids.data(), local_weights.data(), (int)local_ids.size(), + single_out, step_tel ? &local_tel : nullptr); + if (!local_ok) { + if (worker_pending) { + std::vector discard; + expert_worker->eval_end(pending, discard, step_tel ? &ipc_timing : nullptr); + } return false; } - add_ffn_telemetry(step_tel, single_tel); - std::memcpy(ffn_out_host.data() + (size_t)ti * (size_t)n_embd, - single_out.data(), sizeof(float) * (size_t)n_embd); - } - if (step_tel) step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); - return true; -} - -static Ds4MoeRouting build_moe_routing( - ggml_context * ctx, - ggml_tensor * cur, - const DeepSeek4Weights & w, - const DeepSeek4Layer & L, + add_ffn_telemetry(step_tel, local_tel); + std::memcpy(dst, single_out.data(), sizeof(float) * (size_t)n_embd); + if (!remote_ids.empty()) { + bool worker_ok = false; + if (worker_pending) { + worker_ok = expert_worker->eval_end(pending, worker_out, step_tel ? &ipc_timing : nullptr); + } else { + worker_t0 = Ds4TimingClock::now(); + worker_ok = expert_worker->eval(layer, 1, n_embd, (int)remote_ids.size(), + ffn_normed_host + (size_t)ti * (size_t)n_embd, + remote_ids.data(), remote_weights.data(), + worker_out, step_tel ? &ipc_timing : nullptr); + } + if (!worker_ok) return false; + if (step_tel) { + if (worker_pending) { + step_tel->worker_us += ipc_timing.parent_write_us + + ipc_timing.parent_wait_us + + ipc_timing.parent_read_us; + } else { + step_tel->worker_us += ds4_elapsed_us(worker_t0, Ds4TimingClock::now()); + } + } + if (step_tel) { + step_tel->worker_parent_write_us += ipc_timing.parent_write_us; + step_tel->worker_parent_wait_us += ipc_timing.parent_wait_us; + step_tel->worker_parent_read_us += ipc_timing.parent_read_us; + step_tel->worker_request_read_us += ipc_timing.worker_request_read_us; + step_tel->worker_partition_us += ipc_timing.worker_partition_us; + step_tel->worker_resident_eval_us += ipc_timing.worker_resident_eval_us; + step_tel->worker_miss_build_us += ipc_timing.worker_miss_build_us; + step_tel->worker_miss_eval_us += ipc_timing.worker_miss_eval_us; + step_tel->worker_request_bytes += ipc_timing.request_bytes; + step_tel->worker_response_bytes += ipc_timing.response_bytes; + step_tel->worker_hot_graph_builds += ipc_timing.worker_hot_graph_builds; + step_tel->worker_hot_graph_hits += ipc_timing.worker_hot_graph_hits; + step_tel->worker_cold_graph_builds += ipc_timing.worker_cold_graph_builds; + step_tel->worker_cold_graph_hits += ipc_timing.worker_cold_graph_hits; + step_tel->worker_hot_graph_build_us += ipc_timing.worker_hot_graph_build_us; + step_tel->worker_hot_input_us += ipc_timing.worker_hot_input_us; + step_tel->worker_hot_compute_us += ipc_timing.worker_hot_compute_us; + step_tel->worker_hot_read_us += ipc_timing.worker_hot_read_us; + step_tel->worker_cold_graph_build_us += ipc_timing.worker_cold_graph_build_us; + step_tel->worker_cold_input_us += ipc_timing.worker_cold_input_us; + step_tel->worker_cold_compute_us += ipc_timing.worker_cold_compute_us; + step_tel->worker_cold_read_us += ipc_timing.worker_cold_read_us; + if (worker_owns_hot_ids) { + step_tel->hot_selected += (int)remote_ids.size(); + } else { + step_tel->cold_selected += (int)remote_ids.size(); + } + } + for (int i = 0; i < n_embd; ++i) { + dst[i] += worker_out[(size_t)i]; + } + } + } + if (step_tel) step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); + return true; +} + +static bool ensure_cached_decode_gpu_ffn_state( + DeepSeek4CachedDecodeGpuFfnState & out, + ggml_backend_t backend, + const DeepSeek4Weights & w) { + if (out.valid() && out.owner_ctx == w.ctx && out.backend == backend) { + return true; + } + + out.free(); + if (!init_gpu_resident_state(out.state, backend, w.n_embd)) { + return false; + } + + std::vector zeros((size_t)w.n_embd, 0.0f); + ggml_backend_tensor_set(out.state.combine.residual_in, + zeros.data(), 0, + sizeof(float) * zeros.size()); + out.owner_ctx = w.ctx; + out.backend = backend; + return true; +} + +static bool ds4_try_gpu_resident_decode_ffn( + bool enabled, + ggml_backend_t backend, + ggml_backend_t cpu_backend, + const DeepSeek4Weights & w, + const MoeHybridConfig & hybrid_cfg, + const MoeLayerDesc & desc, + MoeHybridLayerStorage & storage, + ggml_tensor * ffn_post_gpu, + DeepSeek4CachedDecodeGpuFfnState & gpu_ffn_state, + const int32_t * selected_ids, + const float * selected_weights, + int n_selected, + ggml_tensor * selected_weights_gpu, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layer, + ggml_tensor ** out_gpu, + DeepSeek4StepTelemetry * step_tel) { + if (out_gpu) { + *out_gpu = nullptr; + } + if (!enabled || !ffn_post_gpu) { + return true; + } + + if (!ensure_cached_decode_gpu_ffn_state(gpu_ffn_state, backend, w)) { + return false; + } + + MoeHybridFfnTelemetry ffn_tel; + const auto ffn_t0 = Ds4TimingClock::now(); + if (!eval_moe_hybrid_ffn_gpu_resident( + backend, hybrid_cfg, desc, storage, cpu_backend, + ffn_post_gpu, + nullptr, + gpu_ffn_state.state, + selected_ids, selected_weights, n_selected, + selected_weights_gpu, + out_gpu, + expert_compute, expert_layer, + step_tel ? &ffn_tel : nullptr)) { + return false; + } + + if (step_tel) { + step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); + add_ffn_telemetry(step_tel, ffn_tel); + } + return true; +} + +static bool ds4_all_selected_hot( + const MoeHybridLayerStorage & storage, + const int32_t * selected_ids, + int n_selected) { + for (int i = 0; i < n_selected; ++i) { + const int32_t gid = selected_ids[i]; + if (gid < 0 || gid >= (int32_t) storage.hot_local_by_global.size()) { + return false; + } + if (storage.hot_local_by_global[(size_t) gid] < 0) { + return false; + } + } + return true; +} + +static Ds4MoeRouting build_moe_routing( + ggml_context * ctx, + ggml_tensor * cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, int n_tokens) { Ds4MoeRouting out; ggml_tensor * logits = ggml_mul_mat(ctx, L.ffn_gate_inp, cur); @@ -1826,7 +2396,99 @@ static void cpu_rms_norm(float * out, const float * x, int n, float eps) { for (int i = 0; i < n; i++) out[i] = x[i] * scale; } +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386) || defined(_M_IX86) +#if defined(__GNUC__) || defined(__clang__) +#define DS4_TARGET_AVX2_F16C_FMA __attribute__((target("avx2,fma,f16c"))) +#else +#define DS4_TARGET_AVX2_F16C_FMA +#endif + +static bool ds4_cpu_has_avx2_f16c_fma() { +#if defined(__GNUC__) || defined(__clang__) + static const bool ok = + __builtin_cpu_supports("avx2") && + __builtin_cpu_supports("fma") && + __builtin_cpu_supports("f16c"); + return ok; +#else + return false; +#endif +} + +DS4_TARGET_AVX2_F16C_FMA +static inline float cpu_hsum_f32x8(__m256 v) { + const __m128 lo = _mm256_castps256_ps128(v); + const __m128 hi = _mm256_extractf128_ps(v, 1); + __m128 sum = _mm_add_ps(lo, hi); + sum = _mm_hadd_ps(sum, sum); + sum = _mm_hadd_ps(sum, sum); + return _mm_cvtss_f32(sum); +} + +DS4_TARGET_AVX2_F16C_FMA +static float cpu_dot_f16_row_avx2(const uint16_t * row, const float * x, int cols) { + __m256 acc = _mm256_setzero_ps(); + int c = 0; + for (; c + 7 < cols; c += 8) { + const __m256 xv = _mm256_loadu_ps(x + c); + const __m256 rv = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(row + c))); + acc = _mm256_fmadd_ps(rv, xv, acc); + } + float sum = cpu_hsum_f32x8(acc); + for (; c < cols; ++c) { + sum += ggml_fp16_to_fp32(row[c]) * x[c]; + } + return sum; +} + +DS4_TARGET_AVX2_F16C_FMA +static void cpu_matvec_f16_rows4_avx2( + float * out, + const uint16_t * row0, + const uint16_t * row1, + const uint16_t * row2, + const uint16_t * row3, + const float * x, + int cols) { + __m256 acc0 = _mm256_setzero_ps(); + __m256 acc1 = _mm256_setzero_ps(); + __m256 acc2 = _mm256_setzero_ps(); + __m256 acc3 = _mm256_setzero_ps(); + + int c = 0; + for (; c + 7 < cols; c += 8) { + const __m256 xv = _mm256_loadu_ps(x + c); + acc0 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(row0 + c))), xv, acc0); + acc1 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(row1 + c))), xv, acc1); + acc2 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(row2 + c))), xv, acc2); + acc3 = _mm256_fmadd_ps(_mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(row3 + c))), xv, acc3); + } + + float sum0 = cpu_hsum_f32x8(acc0); + float sum1 = cpu_hsum_f32x8(acc1); + float sum2 = cpu_hsum_f32x8(acc2); + float sum3 = cpu_hsum_f32x8(acc3); + for (; c < cols; ++c) { + const float xv = x[c]; + sum0 += ggml_fp16_to_fp32(row0[c]) * xv; + sum1 += ggml_fp16_to_fp32(row1[c]) * xv; + sum2 += ggml_fp16_to_fp32(row2[c]) * xv; + sum3 += ggml_fp16_to_fp32(row3[c]) * xv; + } + + out[0] = sum0; + out[1] = sum1; + out[2] = sum2; + out[3] = sum3; +} +#endif + static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386) || defined(_M_IX86) + if (ds4_cpu_has_avx2_f16c_fma()) { + return cpu_dot_f16_row_avx2(row, x, cols); + } +#endif float acc = 0.0f; for (int c = 0; c < cols; c++) { acc += ggml_fp16_to_fp32(row[c]) * x[c]; @@ -1834,23 +2496,43 @@ static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { return acc; } -static void cpu_matvec_f16_serial(float * out, const uint16_t * mat, const float * x, int rows, int cols) { - // mat: [cols, rows] in row-major F16 (ggml layout: ne[0]=cols, ne[1]=rows) - // out[r] = dot(mat_row_r, x) for r in [0, rows) - for (int r = 0; r < rows; r++) { +static void cpu_matvec_f16_rows(float * out, const uint16_t * mat, const float * x, int r0, int r1, int cols) { + int r = r0; +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386) || defined(_M_IX86) + if (ds4_cpu_has_avx2_f16c_fma()) { + for (; r + 3 < r1; r += 4) { + cpu_matvec_f16_rows4_avx2( + out + r, + mat + (size_t)(r + 0) * cols, + mat + (size_t)(r + 1) * cols, + mat + (size_t)(r + 2) * cols, + mat + (size_t)(r + 3) * cols, + x, + cols); + } + } +#endif + for (; r < r1; ++r) { const uint16_t * row = mat + (size_t)r * cols; out[r] = cpu_dot_f16_row(row, x, cols); } } +static void cpu_matvec_f16_serial(float * out, const uint16_t * mat, const float * x, int rows, int cols) { + // mat: [cols, rows] in row-major F16 (ggml layout: ne[0]=cols, ne[1]=rows) + // out[r] = dot(mat_row_r, x) for r in [0, rows) + cpu_matvec_f16_rows(out, mat, x, 0, rows, cols); +} + static void cpu_matvec_f16(float * out, const uint16_t * mat, const float * x, int rows, int cols) { const int64_t ops = (int64_t)rows * (int64_t)cols; + if (rows <= 32 || ops < 262144) { + cpu_matvec_f16_serial(out, mat, x, rows, cols); + return; + } const int min_parallel_rows = ops >= 262144 ? 1 : 512; ds4_parallel_for_tokens(rows, min_parallel_rows, [&](int r0, int r1) { - for (int r = r0; r < r1; ++r) { - const uint16_t * row = mat + (size_t)r * cols; - out[r] = cpu_dot_f16_row(row, x, cols); - } + cpu_matvec_f16_rows(out, mat, x, r0, r1, cols); }); } @@ -2087,22 +2769,29 @@ static void hc_pre_batch(std::vector & working, post.resize((size_t)n_tokens * (size_t)n_hc); comb.resize((size_t)n_tokens * (size_t)n_hc * (size_t)n_hc); + // The single-token CUDA mix helper uses one global scratch arena guarded by + // a mutex. Reusing it across a full prefill batch serializes HC-pre token + // work and adds repeated host<->device sync/copy overhead for only a tiny + // 24-float mix result. For batch/prefill we keep the path CPU-parallel + // instead of funneling every token through that serialized helper. ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { std::vector flat(hc_dim); float mix[24]; for (int t = t0; t < t1; ++t) { - hc_pre_auto_into(working.data() + (size_t)t * n_embd, - post.data() + (size_t)t * n_hc, - comb.data() + (size_t)t * n_hc * (size_t)n_hc, - hc_state + (size_t)t * hc_dim, - weights, - fn_tensor, - n_embd, - n_hc, - sinkhorn_iters, - hc_eps, - flat.data(), - mix); + cpu_hc_pre_into(working.data() + (size_t)t * n_embd, + post.data() + (size_t)t * n_hc, + comb.data() + (size_t)t * n_hc * (size_t)n_hc, + hc_state + (size_t)t * hc_dim, + weights.fn_data.data(), + weights.scale_data.data(), + weights.base_data.data(), + n_embd, + n_hc, + sinkhorn_iters, + hc_eps, + flat.data(), + mix, + true); } }); } @@ -2209,7 +2898,11 @@ static bool deepseek4_step_hybrid( int kv_start, std::vector & out_logits, const int32_t * token_ids, - MoeHybridStreamEngine * stream_engine, + ExpertIpcClient * expert_worker, + bool worker_owns_hot_ids, + bool disable_cached_decode, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layers, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats) { const auto step_t0 = Ds4TimingClock::now(); @@ -2234,182 +2927,823 @@ static bool deepseek4_step_hybrid( static std::vector hc_layer_weights; static HcWeightsCpu hc_output_weights; static std::vector hash_routing_tables; - if (hc_layer_weights.empty()) { + static int hc_loaded_n_layer = 0; + if (hc_loaded_n_layer != w.n_layer) { hc_layer_weights.resize((size_t)w.n_layer); hash_routing_tables.resize((size_t)w.n_layer); for (int il = 0; il < w.n_layer; il++) { const DeepSeek4Layer & L = w.layers[(size_t)il]; + hc_layer_weights[il] = {}; + hash_routing_tables[(size_t)il] = {}; load_hc_weights_cpu(hc_layer_weights[il].attn, L.hc_attn_fn, L.hc_attn_scale, L.hc_attn_base); load_hc_weights_cpu(hc_layer_weights[il].ffn, L.hc_ffn_fn, L.hc_ffn_scale, L.hc_ffn_base); if (il < w.n_hash_layer && L.ffn_gate_tid2eid) { load_hash_routing_cpu(hash_routing_tables[(size_t)il], L.ffn_gate_tid2eid); } } + hc_output_weights = {}; load_hc_weights_cpu(hc_output_weights, w.output_hc_fn, w.output_hc_scale, w.output_hc_base); + hc_loaded_n_layer = w.n_layer; } - for (int il = 0; il < w.n_layer; ++il) { - const DeepSeek4Layer & L = w.layers[(size_t) il]; - DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; - const HcLayerWeightsCpu & hc_lw = hc_layer_weights[(size_t)il]; + // Multi-target expert-split currently fans out enough distinct CUDA graph + // shapes during decode to overrun 24 GiB cards on first graph instantiation. + // Keep the cached decode path for the existing single-target routes, but + // force eager decode graphs for the multi-target carrier until that cache + // policy is narrowed. + const bool use_cached_decode = + n_tokens == 1 && ds4_backend_is_gpu(backend) && !disable_cached_decode; + // Batch HC graph is opt-in: strict PR489 validation showed the eager path + // is quality-clean and slightly faster on the target HIP host. + const bool use_backend_batch_hc_graph = + n_tokens > 1 && + ds4_backend_is_gpu(backend) && + ds4_env_flag("DFLASH_DS4_ENABLE_BACKEND_BATCH_HC_GRAPH") && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_BATCH_HC_GRAPH"); + const bool disable_backend_decode_hc_direct = + ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC_DIRECT"); + const bool use_backend_decode_hc_direct = + use_cached_decode && + ds4_backend_is_hip(backend) && + !disable_backend_decode_hc_direct; + const bool use_backend_decode_hc_graph = + use_cached_decode && !use_backend_decode_hc_direct; + const bool use_backend_hc_graph = + use_backend_batch_hc_graph || use_backend_decode_hc_graph; + const bool use_backend_hc_post_graph = + use_backend_batch_hc_graph || use_cached_decode; + const bool use_gpu_resident_decode_ffn = + use_cached_decode && + ds4_env_flag("DFLASH_DS4_GPU_RESIDENT_FFN") && + expert_worker == nullptr; - // ── HC pre (attention) ────────────────────────────────────── - // For decode (n_tokens=1): compute working vector from HC state - const auto hc_pre_attn_t0 = Ds4TimingClock::now(); - std::vector cur((size_t)n_embd * (size_t)n_tokens); - HcPreResult hc_attn_result; - if (hc_lw.attn.loaded && n_tokens == 1) { - hc_attn_result = hc_pre_auto(hc_state.data(), hc_lw.attn, L.hc_attn_fn, - n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); - memcpy(cur.data(), hc_attn_result.working.data(), (size_t)n_embd * sizeof(float)); - } else { - // Fallback: use first HC stream - memcpy(cur.data(), hc_state.data(), (size_t)n_embd * (size_t)n_tokens * sizeof(float)); - } - if (telemetry) telemetry->hc_pre_attn_us += ds4_elapsed_us(hc_pre_attn_t0, Ds4TimingClock::now()); + std::vector hc_post; + std::vector hc_comb; + std::vector next_hc; + std::vector attn_out_host; + std::vector ffn_out_host_scratch; + std::vector final_embd; - // ── Build attention graph ─────────────────────────────────── - const auto attn_build_t0 = Ds4TimingClock::now(); - const size_t ctx_size = 48 * 1024 * 1024; - ggml_init_params params{}; - params.mem_size = ctx_size; - params.mem_buffer = nullptr; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - if (!ctx) { - if (hot_alloc) ggml_gallocr_free(hot_alloc); - if (cold_alloc) ggml_gallocr_free(cold_alloc); - return false; + static std::vector cached_decode_attn_hc_pre_graphs; + static std::vector cached_decode_ffn_hc_pre_graphs; + static DeepSeek4CachedDecodeHcPostGraph cached_decode_hc_post_graph; + static std::vector cached_batch_attn_hc_pre_graphs; + static std::vector cached_batch_ffn_hc_pre_graphs; + static DeepSeek4CachedDecodeHcPostGraph cached_batch_hc_post_graph; + static std::vector cached_attn_allocs; + static std::vector> cached_decode_attn_graphs; + static std::vector cached_decode_route_graphs; + static thread_local DeepSeek4CachedDecodeGpuFfnState cached_decode_gpu_ffn_state; + static DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; + static int decode_cache_n_layer = 0; + static int batch_hc_cache_n_layer = 0; + static int dynamic_attn_alloc_n_layer = 0; + static int route_cache_n_layer = 0; + if (use_cached_decode && decode_cache_n_layer != w.n_layer) { + for (auto & g : cached_decode_attn_hc_pre_graphs) { + g.free(); } - - ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); - ggml_set_input(inp); - std::vector i32_inputs; - std::vector i32_array_inputs; - std::vector i64_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); - // 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)); - if (!ggml_gallocr_alloc_graph(attn_alloc, gf)) { - ggml_gallocr_free(attn_alloc); - ggml_free(ctx); - if (hot_alloc) ggml_gallocr_free(hot_alloc); - if (cold_alloc) ggml_gallocr_free(cold_alloc); - return false; + cached_decode_attn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_ffn_hc_pre_graphs) { + g.free(); } - if (telemetry) telemetry->attn_build_us += ds4_elapsed_us(attn_build_t0, Ds4TimingClock::now()); - ggml_backend_tensor_set(inp, cur.data(), 0, sizeof(float) * cur.size()); - for (const DeepSeek4I32InputBinding & binding : i32_inputs) { - ggml_backend_tensor_set(binding.tensor, &binding.value, 0, sizeof(binding.value)); - } - for (const DeepSeek4I32ArrayBinding & binding : i32_array_inputs) { - ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, - sizeof(int32_t) * binding.values.size()); - } - for (const DeepSeek4I64ArrayBinding & binding : i64_array_inputs) { - ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, - sizeof(int64_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()); - std::vector attn_out_host((size_t)n_embd * (size_t)n_tokens); - if (ok) { - const auto attn_read_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_get(attn_out, attn_out_host.data(), 0, sizeof(float) * attn_out_host.size()); - if (telemetry) telemetry->attn_read_us += ds4_elapsed_us(attn_read_t0, Ds4TimingClock::now()); - } - ggml_gallocr_free(attn_alloc); - ggml_free(ctx); - if (!ok) { - if (hot_alloc) ggml_gallocr_free(hot_alloc); - if (cold_alloc) ggml_gallocr_free(cold_alloc); - return false; + cached_decode_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + cached_decode_hc_post_graph.free(); + for (auto & per_layer : cached_decode_attn_graphs) { + for (auto & g : per_layer) { + g.free(); + } + } + cached_decode_attn_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_route_graphs) { + g.free(); + } + cached_decode_route_graphs.assign((size_t)w.n_layer, {}); + cached_decode_output_graph.free(); + decode_cache_n_layer = w.n_layer; + route_cache_n_layer = w.n_layer; + } + if (use_backend_batch_hc_graph && batch_hc_cache_n_layer != w.n_layer) { + for (auto & g : cached_batch_attn_hc_pre_graphs) { + g.free(); + } + cached_batch_attn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_batch_ffn_hc_pre_graphs) { + g.free(); + } + cached_batch_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + cached_batch_hc_post_graph.free(); + batch_hc_cache_n_layer = w.n_layer; + } + if (dynamic_attn_alloc_n_layer != w.n_layer) { + for (auto & alloc : cached_attn_allocs) { + alloc.free(); } + cached_attn_allocs.assign((size_t)w.n_layer, {}); + dynamic_attn_alloc_n_layer = w.n_layer; + } + if (route_cache_n_layer != w.n_layer) { + for (auto & g : cached_decode_route_graphs) { + g.free(); + } + cached_decode_route_graphs.assign((size_t)w.n_layer, {}); + route_cache_n_layer = w.n_layer; + } - // ── HC post (attention) ───────────────────────────────────── - const auto hc_post_attn_t0 = Ds4TimingClock::now(); - if (hc_lw.attn.loaded && n_tokens == 1) { - std::vector new_hc((size_t)hc_dim); - cpu_hc_post(new_hc.data(), attn_out_host.data(), hc_state.data(), - hc_attn_result.post, hc_attn_result.comb, n_embd, n_hc); - memcpy(hc_state.data(), new_hc.data(), (size_t)hc_dim * sizeof(float)); - } else { - for (int i = 0; i < n_embd * n_tokens; i++) { - hc_state[(size_t)i] += attn_out_host[(size_t)i]; + ggml_tensor * hc_state_backend = nullptr; + DeepSeek4CachedDecodeHcPostGraph * hc_post_graph = nullptr; + const bool use_host_hc_post = + !use_backend_hc_post_graph || use_backend_decode_hc_direct; + if (use_backend_hc_post_graph || use_host_hc_post) { + hc_post.resize((size_t)n_hc * (size_t)n_tokens); + hc_comb.resize((size_t)n_hc * (size_t)n_hc * (size_t)n_tokens); + next_hc.resize((size_t)hc_dim * (size_t)n_tokens); + attn_out_host.resize((size_t)n_embd * (size_t)n_tokens); + } + if (use_cached_decode || use_backend_batch_hc_graph) { + ffn_out_host_scratch.resize((size_t)n_embd * (size_t)n_tokens); + final_embd.resize((size_t)n_embd * (size_t)n_tokens); + if (use_backend_hc_post_graph) { + hc_post_graph = use_backend_batch_hc_graph + ? &cached_batch_hc_post_graph + : &cached_decode_hc_post_graph; + if (!hc_post_graph->valid() || + hc_post_graph->owner_ctx != w.ctx || + hc_post_graph->backend != backend || + hc_post_graph->n_tokens != n_tokens) { + if (!build_cached_decode_hc_post_graph(*hc_post_graph, backend, w, n_tokens)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } } + ggml_backend_tensor_set(hc_post_graph->residual_hc, + hc_state.data(), 0, + sizeof(float) * hc_state.size()); + hc_state_backend = hc_post_graph->residual_hc; } - if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); + } - // ── HC pre (FFN) ──────────────────────────────────────────── - const auto hc_pre_ffn_t0 = Ds4TimingClock::now(); - std::vector ffn_working((size_t)n_embd * (size_t)n_tokens); - HcPreResult hc_ffn_result; - if (hc_lw.ffn.loaded && n_tokens == 1) { - hc_ffn_result = hc_pre_auto(hc_state.data(), hc_lw.ffn, L.hc_ffn_fn, - n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); - memcpy(ffn_working.data(), hc_ffn_result.working.data(), (size_t)n_embd * sizeof(float)); - } else { - memcpy(ffn_working.data(), hc_state.data(), (size_t)n_embd * (size_t)n_tokens * sizeof(float)); - } - if (telemetry) telemetry->hc_pre_ffn_us += ds4_elapsed_us(hc_pre_ffn_t0, Ds4TimingClock::now()); + static thread_local DeepSeek4LayerRangeScratch hybrid_scratch; + hybrid_scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, w.n_expert_used, w.n_expert); + std::vector & hybrid_ffn_normed_host = hybrid_scratch.ffn_normed_host; + std::vector & hybrid_probs_host = hybrid_scratch.probs_host; + std::vector & hybrid_weights_host = hybrid_scratch.weights_host; + std::vector & hybrid_selected_host = hybrid_scratch.selected_host; - // ── FFN ───────────────────────────────────────────────────── - std::vector ffn_out_host((size_t)n_embd * (size_t)n_tokens, 0.0f); + for (int il = 0; il < w.n_layer; ++il) { + const bool trace_decode = + ds4_env_flag("DFLASH_DS4_TRACE_DECODE") && n_tokens == 1; + if (trace_decode) { + std::fprintf(stderr, + "[deepseek4-trace] decode-step kv=%d layer=%d begin\n", + kv_start, il); + } + const DeepSeek4Layer & L = w.layers[(size_t) il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const HcLayerWeightsCpu & hc_lw = hc_layer_weights[(size_t)il]; - if (il < w.n_hash_layer && L.ffn_gate_tid2eid) { - // Hash-routed layers: selected experts come from token_id -> expert_ids, - // while weights still come from router probabilities for those experts. - if (!token_ids || !hash_routing_tables[(size_t)il].loaded) { - std::fprintf(stderr, "[deepseek4] missing token ids/hash table for layer %d\n", il); - if (hot_alloc) ggml_gallocr_free(hot_alloc); - if (cold_alloc) ggml_gallocr_free(cold_alloc); - return false; - } - ggml_init_params ffn_params{}; - const auto route_build_t0 = Ds4TimingClock::now(); - ffn_params.mem_size = 16 * 1024 * 1024; - ffn_params.mem_buffer = nullptr; - ffn_params.no_alloc = true; - ggml_context * ffn_ctx = ggml_init(ffn_params); - if (!ffn_ctx) { - if (hot_alloc) ggml_gallocr_free(hot_alloc); - if (cold_alloc) ggml_gallocr_free(cold_alloc); - return false; + // ── HC pre (attention) ────────────────────────────────────── + const auto hc_pre_attn_t0 = Ds4TimingClock::now(); + std::vector cur((size_t)n_embd * (size_t)n_tokens); + const ggml_tensor * attn_in_backend = nullptr; + const ggml_tensor * attn_post_backend = nullptr; + const ggml_tensor * attn_comb_backend = nullptr; + if (use_backend_decode_hc_direct) { + auto & cached = cached_decode_attn_hc_pre_graphs[(size_t)il]; + if (!cached.valid() || + cached.owner_ctx != w.ctx || + cached.backend != backend || + cached.layer_idx != il || + cached.ffn) { + const auto hc_pre_attn_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_hc_pre_graph(cached, backend, w, L, + hc_lw.attn.scale_data.data(), + il, false)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->hc_pre_build_us += ds4_elapsed_us(hc_pre_attn_build_t0, Ds4TimingClock::now()); } - ggml_tensor * ffn_inp = ggml_new_tensor_2d(ffn_ctx, GGML_TYPE_F32, n_embd, n_tokens); - ggml_set_input(ffn_inp); - ggml_tensor * ffn_normed = build_rms_norm(ffn_ctx, ffn_inp, L.ffn_norm, w.rms_eps); - ggml_tensor * router_logits = ggml_mul_mat(ffn_ctx, L.ffn_gate_inp, ffn_normed); - ggml_tensor * router_probs = ggml_sqrt(ffn_ctx, ggml_softplus(ffn_ctx, router_logits)); - ggml_cgraph * ffn_gf = ggml_new_graph(ffn_ctx); - ggml_build_forward_expand(ffn_gf, ffn_normed); - ggml_build_forward_expand(ffn_gf, router_probs); - ggml_gallocr_t ffn_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(ffn_alloc, ffn_gf)) { - ggml_gallocr_free(ffn_alloc); ggml_free(ffn_ctx); + const auto hc_pre_attn_compute_t0 = Ds4TimingClock::now(); + if (!ds4_try_gpu_hc_pre_device(cached.sg.hidden_states, + cached.post, + cached.comb, + backend, + il, + false, + hc_state_backend, + L.hc_attn_fn, + L.hc_attn_scale, + L.hc_attn_base, + hc_lw.attn.scale_data.data(), + hc_lw.attn.base_data.data(), + n_embd, + n_hc, + w.n_hc_sinkhorn_iter, + w.hc_eps)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->hc_pre_compute_us += ds4_elapsed_us(hc_pre_attn_compute_t0, Ds4TimingClock::now()); + attn_in_backend = cached.sg.hidden_states; + attn_post_backend = cached.post; + attn_comb_backend = cached.comb; + } else if (use_backend_hc_graph) { + auto & cached = use_backend_batch_hc_graph + ? cached_batch_attn_hc_pre_graphs[(size_t)il] + : cached_decode_attn_hc_pre_graphs[(size_t)il]; + if (!cached.valid() || + cached.owner_ctx != w.ctx || + cached.backend != backend || + cached.layer_idx != il || + cached.n_tokens != n_tokens || + cached.ffn) { + const auto hc_pre_attn_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_hc_pre_graph(cached, backend, w, L, + hc_lw.attn.scale_data.data(), + il, false, n_tokens)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->hc_pre_build_us += ds4_elapsed_us(hc_pre_attn_build_t0, Ds4TimingClock::now()); + } + const auto hc_pre_attn_input_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_copy(hc_state_backend, cached.sg.inp_embed); + if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_attn_input_t0, Ds4TimingClock::now()); + const auto hc_pre_attn_compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, cached.sg.gf) != GGML_STATUS_SUCCESS) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->hc_pre_compute_us += ds4_elapsed_us(hc_pre_attn_compute_t0, Ds4TimingClock::now()); + attn_in_backend = cached.sg.hidden_states; + attn_post_backend = cached.post; + attn_comb_backend = cached.comb; + } else if (hc_lw.attn.loaded && n_tokens == 1) { + HcPreResult hc_attn_result = hc_pre_auto(hc_state.data(), hc_lw.attn, L.hc_attn_fn, + n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); + memcpy(cur.data(), hc_attn_result.working.data(), (size_t)n_embd * sizeof(float)); + std::memcpy(hc_post.data(), hc_attn_result.post, (size_t)n_hc * sizeof(float)); + std::memcpy(hc_comb.data(), hc_attn_result.comb, (size_t)n_hc * (size_t)n_hc * sizeof(float)); + } else if (hc_lw.attn.loaded) { + hc_pre_batch(cur, hc_post, hc_comb, + hc_state.data(), hc_lw.attn, L.hc_attn_fn, + n_tokens, n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); + } else { + memcpy(cur.data(), hc_state.data(), (size_t)n_embd * (size_t)n_tokens * sizeof(float)); + } + if (telemetry) telemetry->hc_pre_attn_us += ds4_elapsed_us(hc_pre_attn_t0, Ds4TimingClock::now()); + ds4_trace_decode_marker(trace_decode, kv_start, il, "hc_pre_attn_done"); + + // ── Attention ─────────────────────────────────────────────── + ggml_tensor * attn_out = nullptr; + ggml_cgraph * gf = nullptr; + ggml_context * ctx = nullptr; + DeepSeek4CachedDecodeAttnGraph * cached_attn = nullptr; + if (use_cached_decode) { + const int token_pos = kv_start + n_tokens - 1; + const int n_raw = std::min(kv_start + 1, w.n_swa); + const int ratio = (int)w.compress_ratios[il]; + const int n_comp_attn = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + const int n_index_comp = (ratio == 4) ? ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos) : 0; + const bool flush_boundary = ratio > 0 && ((token_pos + 1) % ratio) == 0; + auto & per_layer = cached_decode_attn_graphs[(size_t)il]; + auto it = std::find_if(per_layer.begin(), per_layer.end(), + [&](const DeepSeek4CachedDecodeAttnGraph & candidate) { + return candidate.valid() && + candidate.owner_ctx == w.ctx && + candidate.backend == backend && + candidate.layer_idx == il && + candidate.n_raw == n_raw && + candidate.n_comp_attn == n_comp_attn && + candidate.n_index_comp == n_index_comp && + candidate.flush_boundary == flush_boundary; + }); + if (it == per_layer.end()) { + if (per_layer.size() >= 20) { + per_layer.front().free(); + per_layer.erase(per_layer.begin()); + } + per_layer.emplace_back(); + auto & candidate = per_layer.back(); + const auto attn_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_attn_graph(candidate, backend, w, L, lc, il, kv_start, + n_raw, n_comp_attn, n_index_comp, + flush_boundary)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->attn_build_us += ds4_elapsed_us(attn_build_t0, Ds4TimingClock::now()); + it = std::prev(per_layer.end()); + } + cached_attn = &*it; + gf = cached_attn->sg.gf; + attn_out = cached_attn->sg.hidden_states; + + const int64_t raw_row = kv_start % w.n_swa; + const int32_t rope_pos = kv_start; + const int32_t neg_pos = -kv_start; + if (attn_in_backend) { + ggml_backend_tensor_copy(attn_in_backend, cached_attn->sg.inp_embed); + } else { + ggml_backend_tensor_set(cached_attn->sg.inp_embed, cur.data(), 0, sizeof(float) * cur.size()); + } + ggml_backend_tensor_set(cached_attn->inputs.rope_pos, &rope_pos, 0, sizeof(rope_pos)); + ggml_backend_tensor_set(cached_attn->inputs.neg_pos, &neg_pos, 0, sizeof(neg_pos)); + ggml_backend_tensor_set(cached_attn->inputs.raw_kv_rows, &raw_row, 0, sizeof(raw_row)); + if (ratio > 0) { + const int pos_mod = token_pos % ratio; + const int32_t ape_row = 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; + ggml_backend_tensor_set(cached_attn->inputs.attn_ape_row, &ape_row, 0, sizeof(ape_row)); + ggml_backend_tensor_set(cached_attn->inputs.attn_state_rows, &state_row, 0, sizeof(state_row)); + if (flush_boundary) { + ggml_backend_tensor_set(cached_attn->inputs.attn_comp_rows, &comp_row, 0, sizeof(comp_row)); + ggml_backend_tensor_set(cached_attn->inputs.attn_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + } + } + if (ratio == 4) { + const int pos_mod = token_pos % ratio; + const int32_t ape_row = 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; + ggml_backend_tensor_set(cached_attn->inputs.index_ape_row, &ape_row, 0, sizeof(ape_row)); + ggml_backend_tensor_set(cached_attn->inputs.index_state_rows, &state_row, 0, sizeof(state_row)); + if (flush_boundary) { + ggml_backend_tensor_set(cached_attn->inputs.index_comp_rows, &comp_row, 0, sizeof(comp_row)); + ggml_backend_tensor_set(cached_attn->inputs.index_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + } + } + } else { + const auto attn_build_t0 = Ds4TimingClock::now(); + const size_t ctx_size = 48 * 1024 * 1024; + ggml_init_params params{}; + params.mem_size = ctx_size; + params.mem_buffer = nullptr; + params.no_alloc = true; + ctx = ggml_init(params); + if (!ctx) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + + ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp); + std::vector i32_inputs; + std::vector i32_array_inputs; + std::vector i64_array_inputs; + const size_t graph_size = n_tokens > 1 ? 32768 : 2048; + gf = ggml_new_graph_custom(ctx, graph_size, false); + + ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps); + 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); + ggml_build_forward_expand(gf, attn_out); + auto & attn_alloc = cached_attn_allocs[(size_t)il]; + if (!attn_alloc.valid() || + attn_alloc.owner_ctx != w.ctx || + attn_alloc.backend != backend) { + attn_alloc.free(); + attn_alloc.alloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + attn_alloc.owner_ctx = w.ctx; + attn_alloc.backend = backend; + } + if (!attn_alloc.alloc || !ggml_gallocr_alloc_graph(attn_alloc.alloc, gf)) { + ggml_free(ctx); + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->attn_build_us += ds4_elapsed_us(attn_build_t0, Ds4TimingClock::now()); + if (attn_in_backend) { + ggml_backend_tensor_copy(attn_in_backend, inp); + } else { + ggml_backend_tensor_set(inp, cur.data(), 0, sizeof(float) * cur.size()); + } + for (const DeepSeek4I32InputBinding & binding : i32_inputs) { + ggml_backend_tensor_set(binding.tensor, &binding.value, 0, sizeof(binding.value)); + } + for (const DeepSeek4I32ArrayBinding & binding : i32_array_inputs) { + ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, + sizeof(int32_t) * binding.values.size()); + } + for (const DeepSeek4I64ArrayBinding & binding : i64_array_inputs) { + ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, + sizeof(int64_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()); + if (!ok) { + ggml_free(ctx); + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (use_backend_batch_hc_graph) { + if (hc_state_backend != hc_post_graph->residual_hc) { + ggml_backend_tensor_copy(hc_state_backend, hc_post_graph->residual_hc); + } + ggml_backend_tensor_copy(attn_out, hc_post_graph->block_out); + ggml_backend_tensor_copy(attn_post_backend, hc_post_graph->post); + ggml_backend_tensor_copy(attn_comb_backend, hc_post_graph->comb); + const auto hc_post_attn_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, hc_post_graph->sg.gf) != GGML_STATUS_SUCCESS) { + ggml_free(ctx); + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + hc_state_backend = hc_post_graph->sg.hidden_states; + if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); + } else { + const auto attn_read_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_get(attn_out, attn_out_host.data(), 0, sizeof(float) * attn_out_host.size()); + if (telemetry) telemetry->attn_read_us += ds4_elapsed_us(attn_read_t0, Ds4TimingClock::now()); + } + ggml_free(ctx); + ctx = nullptr; + attn_out = nullptr; + } + + if (use_cached_decode) { + const auto attn_compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->attn_compute_us += ds4_elapsed_us(attn_compute_t0, Ds4TimingClock::now()); + if (use_backend_hc_post_graph) { + if (hc_state_backend != hc_post_graph->residual_hc) { + ggml_backend_tensor_copy(hc_state_backend, hc_post_graph->residual_hc); + } + ggml_backend_tensor_copy(attn_out, hc_post_graph->block_out); + ggml_backend_tensor_copy(attn_post_backend, hc_post_graph->post); + ggml_backend_tensor_copy(attn_comb_backend, hc_post_graph->comb); + const auto hc_post_attn_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, hc_post_graph->sg.gf) != GGML_STATUS_SUCCESS) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + hc_state_backend = hc_post_graph->sg.hidden_states; + if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); + } else { + const auto attn_read_t0 = Ds4TimingClock::now(); + if (use_backend_decode_hc_direct) { + const Ds4TensorReadback readbacks[] = { + {attn_out, attn_out_host.data(), + sizeof(float) * attn_out_host.size()}, + {attn_post_backend, hc_post.data(), + sizeof(float) * hc_post.size()}, + {attn_comb_backend, hc_comb.data(), + sizeof(float) * hc_comb.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } else { + ggml_backend_tensor_get(attn_out, attn_out_host.data(), 0, sizeof(float) * attn_out_host.size()); + } + if (telemetry) telemetry->attn_read_us += ds4_elapsed_us(attn_read_t0, Ds4TimingClock::now()); + const auto hc_post_attn_t0 = Ds4TimingClock::now(); + hc_post_batch(next_hc, + attn_out_host.data(), + hc_state.data(), + hc_post.data(), + hc_comb.data(), + n_tokens, + n_embd, + n_hc); + std::memcpy(hc_state.data(), next_hc.data(), + next_hc.size() * sizeof(float)); + if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); + } + } + ds4_trace_decode_marker(trace_decode, kv_start, il, "attn_done"); + + // ── HC pre (FFN) ──────────────────────────────────────────── + const auto hc_pre_ffn_t0 = Ds4TimingClock::now(); + std::vector ffn_working; + HcPreResult hc_ffn_result; + ggml_tensor * ffn_in_backend = nullptr; + ggml_tensor * ffn_post_backend = nullptr; + ggml_tensor * ffn_comb_backend = nullptr; + ggml_tensor * ffn_out_backend = nullptr; + if (use_backend_decode_hc_direct) { + auto & cached = cached_decode_ffn_hc_pre_graphs[(size_t)il]; + if (!cached.valid() || + cached.owner_ctx != w.ctx || + cached.backend != backend || + cached.layer_idx != il || + !cached.ffn) { + const auto hc_pre_ffn_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_hc_pre_graph(cached, backend, w, L, + hc_lw.ffn.scale_data.data(), + il, true)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->hc_pre_build_us += ds4_elapsed_us(hc_pre_ffn_build_t0, Ds4TimingClock::now()); + } + const auto hc_pre_ffn_compute_t0 = Ds4TimingClock::now(); + if (!ds4_try_gpu_hc_pre_device(cached.sg.hidden_states, + cached.post, + cached.comb, + backend, + il, + true, + hc_state_backend, + L.hc_ffn_fn, + L.hc_ffn_scale, + L.hc_ffn_base, + hc_lw.ffn.scale_data.data(), + hc_lw.ffn.base_data.data(), + n_embd, + n_hc, + w.n_hc_sinkhorn_iter, + w.hc_eps)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->hc_pre_compute_us += ds4_elapsed_us(hc_pre_ffn_compute_t0, Ds4TimingClock::now()); + ffn_in_backend = cached.sg.hidden_states; + ffn_post_backend = cached.post; + ffn_comb_backend = cached.comb; + } else if (use_backend_hc_graph) { + auto & cached = use_backend_batch_hc_graph + ? cached_batch_ffn_hc_pre_graphs[(size_t)il] + : cached_decode_ffn_hc_pre_graphs[(size_t)il]; + if (!cached.valid() || + cached.owner_ctx != w.ctx || + cached.backend != backend || + cached.layer_idx != il || + cached.n_tokens != n_tokens || + !cached.ffn) { + const auto hc_pre_ffn_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_hc_pre_graph(cached, backend, w, L, + hc_lw.ffn.scale_data.data(), + il, true, n_tokens)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->hc_pre_build_us += ds4_elapsed_us(hc_pre_ffn_build_t0, Ds4TimingClock::now()); + } + const auto hc_pre_ffn_input_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_copy(hc_state_backend, cached.sg.inp_embed); + if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_ffn_input_t0, Ds4TimingClock::now()); + const auto hc_pre_ffn_compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, cached.sg.gf) != GGML_STATUS_SUCCESS) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->hc_pre_compute_us += ds4_elapsed_us(hc_pre_ffn_compute_t0, Ds4TimingClock::now()); + ffn_in_backend = cached.sg.hidden_states; + ffn_post_backend = cached.post; + ffn_comb_backend = cached.comb; + } else if (hc_lw.ffn.loaded && n_tokens == 1) { + ffn_working.resize((size_t)n_embd * (size_t)n_tokens); + hc_ffn_result = hc_pre_auto(hc_state.data(), hc_lw.ffn, L.hc_ffn_fn, + n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); + memcpy(ffn_working.data(), hc_ffn_result.working.data(), (size_t)n_embd * sizeof(float)); + } else if (hc_lw.ffn.loaded) { + hc_pre_batch(ffn_working, hc_post, hc_comb, + hc_state.data(), hc_lw.ffn, L.hc_ffn_fn, + n_tokens, n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); + } else { + ffn_working.resize((size_t)n_embd * (size_t)n_tokens); + memcpy(ffn_working.data(), hc_state.data(), (size_t)n_embd * (size_t)n_tokens * sizeof(float)); + } + if (telemetry) telemetry->hc_pre_ffn_us += ds4_elapsed_us(hc_pre_ffn_t0, Ds4TimingClock::now()); + ds4_trace_decode_marker(trace_decode, kv_start, il, "hc_pre_ffn_done"); + + // ── FFN ───────────────────────────────────────────────────── + std::vector ffn_out_host_local; + std::vector & ffn_out_host = use_cached_decode ? ffn_out_host_scratch : ffn_out_host_local; + + if ((use_cached_decode || n_tokens > 1) && + il < w.n_hash_layer && L.ffn_gate_tid2eid) { + if (!token_ids || !hash_routing_tables[(size_t)il].loaded) { + std::fprintf(stderr, "[deepseek4] missing token ids/hash table for layer %d\n", il); + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + + auto & cached = cached_decode_route_graphs[(size_t)il]; + if (!cached.valid() || + cached.owner_ctx != w.ctx || + cached.backend != backend || + cached.layer_idx != il || + cached.n_tokens != n_tokens || + !cached.hash_routed) { + const auto route_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_route_graph(cached, backend, w, L, il, n_tokens, true)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->route_build_us += ds4_elapsed_us(route_build_t0, Ds4TimingClock::now()); + } + + const auto route_select_t0 = Ds4TimingClock::now(); + const auto & hash_table = hash_routing_tables[(size_t)il].ids; + for (int ti = 0; ti < n_tokens; ++ti) { + const int32_t tok = token_ids[ti]; + if (tok < 0 || tok >= w.n_vocab) { + std::fprintf(stderr, "[deepseek4] token id %d outside hash table for layer %d\n", tok, il); + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + const int32_t * row = hash_table.data() + (size_t)tok * (size_t)w.n_expert_used; + std::memcpy(hybrid_selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, + row, + sizeof(int32_t) * (size_t)w.n_expert_used); + } + if (telemetry) telemetry->route_select_us += ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); + + if (ffn_in_backend) { + ggml_backend_tensor_copy(ffn_in_backend, cached.sg.inp_embed); + } else { + ggml_backend_tensor_set(cached.sg.inp_embed, ffn_working.data(), 0, + sizeof(float) * ffn_working.size()); + } + ggml_backend_tensor_set(cached.hash_ids, hybrid_selected_host.data(), 0, + sizeof(int32_t) * hybrid_selected_host.size()); + + const auto route_compute_t0 = Ds4TimingClock::now(); + bool ok = ggml_backend_graph_compute(backend, cached.sg.gf) == GGML_STATUS_SUCCESS; + if (telemetry) telemetry->route_compute_us += ds4_elapsed_us(route_compute_t0, Ds4TimingClock::now()); + if (!ok) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + + MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); + MoeLayerDesc desc = make_ds4_moe_layer_desc(L); + auto & storage = moe_hybrid.layers[(size_t) il]; + const MoeExpertLayer * expert_layer = + expert_layers ? expert_layers + il : nullptr; + const bool all_selected_hot = ds4_all_selected_hot( + storage, hybrid_selected_host.data(), n_tokens * w.n_expert_used); + const bool can_skip_ffn_normed_host = + expert_worker == nullptr && + (all_selected_hot || + storage.cold_backend_kind == MoeHybridColdBackend::Gpu || + (expert_compute && expert_layer && + expert_compute->supports_batch_gpu_resident())); + const auto route_read_t0 = Ds4TimingClock::now(); + if (!use_gpu_resident_decode_ffn && !can_skip_ffn_normed_host) { + const Ds4TensorReadback readbacks[] = { + {cached.ffn_normed, hybrid_ffn_normed_host.data(), + sizeof(float) * hybrid_ffn_normed_host.size()}, + {cached.weights, hybrid_weights_host.data(), + sizeof(float) * hybrid_weights_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } else if (!all_selected_hot) { + const Ds4TensorReadback readbacks[] = { + {cached.weights, hybrid_weights_host.data(), + sizeof(float) * hybrid_weights_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } + if (telemetry) telemetry->route_read_us += ds4_elapsed_us(route_read_t0, Ds4TimingClock::now()); + + if (routing_stats) { + for (int ti = 0; ti < n_tokens; ++ti) { + routing_stats->observe(il, + hybrid_selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, + w.n_expert_used); + } + } + + if (!ds4_try_gpu_resident_decode_ffn( + use_gpu_resident_decode_ffn, + backend, cpu_backend, w, + hybrid_cfg, desc, storage, + cached.ffn_normed, + cached_decode_gpu_ffn_state, + hybrid_selected_host.data(), + all_selected_hot ? nullptr : hybrid_weights_host.data(), + w.n_expert_used, + cached.weights, + expert_compute, expert_layer, + &ffn_out_backend, + telemetry)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (!ffn_out_backend && + !eval_ds4_hybrid_or_worker( + backend, cpu_backend, hybrid_cfg, desc, storage, expert_worker, + il, n_embd, w.n_expert_used, + can_skip_ffn_normed_host ? nullptr : hybrid_ffn_normed_host.data(), + cached.ffn_normed, + hybrid_selected_host.data(), hybrid_weights_host.data(), + n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, + expert_compute, expert_layer, + worker_owns_hot_ids, telemetry, + use_backend_hc_post_graph ? &ffn_out_backend : nullptr)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + } else if (il < w.n_hash_layer && L.ffn_gate_tid2eid) { + // Hash-routed layers: selected experts come from token_id -> expert_ids, + // while weights still come from router probabilities for those experts. + if (!token_ids || !hash_routing_tables[(size_t)il].loaded) { + std::fprintf(stderr, "[deepseek4] missing token ids/hash table for layer %d\n", il); + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + ggml_init_params ffn_params{}; + const auto route_build_t0 = Ds4TimingClock::now(); + ffn_params.mem_size = 16 * 1024 * 1024; + ffn_params.mem_buffer = nullptr; + ffn_params.no_alloc = true; + ggml_context * ffn_ctx = ggml_init(ffn_params); + if (!ffn_ctx) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + ggml_tensor * ffn_inp = ggml_new_tensor_2d(ffn_ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(ffn_inp); + ggml_tensor * ffn_normed = build_rms_norm(ffn_ctx, ffn_inp, L.ffn_norm, w.rms_eps); + ggml_tensor * router_logits = ggml_mul_mat(ffn_ctx, L.ffn_gate_inp, ffn_normed); + ggml_tensor * router_probs = ggml_sqrt(ffn_ctx, ggml_softplus(ffn_ctx, router_logits)); + ggml_cgraph * ffn_gf = ggml_new_graph(ffn_ctx); + ggml_build_forward_expand(ffn_gf, ffn_normed); + ggml_build_forward_expand(ffn_gf, router_probs); + ggml_gallocr_t ffn_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(ffn_alloc, ffn_gf)) { + ggml_gallocr_free(ffn_alloc); ggml_free(ffn_ctx); if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; } if (telemetry) telemetry->route_build_us += ds4_elapsed_us(route_build_t0, Ds4TimingClock::now()); - ggml_backend_tensor_set(ffn_inp, ffn_working.data(), 0, sizeof(float) * ffn_working.size()); + if (ffn_in_backend) { + ggml_backend_tensor_copy(ffn_in_backend, ffn_inp); + } else { + ggml_backend_tensor_set(ffn_inp, ffn_working.data(), 0, sizeof(float) * ffn_working.size()); + } const auto route_compute_t0 = Ds4TimingClock::now(); - ok = ggml_backend_graph_compute(backend, ffn_gf) == GGML_STATUS_SUCCESS; + bool ok = ggml_backend_graph_compute(backend, ffn_gf) == GGML_STATUS_SUCCESS; if (telemetry) telemetry->route_compute_us += ds4_elapsed_us(route_compute_t0, Ds4TimingClock::now()); - std::vector ffn_normed_host((size_t)n_embd * (size_t)n_tokens); - std::vector probs_host((size_t)w.n_expert * (size_t)n_tokens); + MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); + MoeLayerDesc desc = make_ds4_moe_layer_desc(L); + auto & storage = moe_hybrid.layers[(size_t) il]; + const MoeExpertLayer * expert_layer = + expert_layers ? expert_layers + il : nullptr; + const bool can_skip_ffn_normed_host = + expert_worker == nullptr && + (storage.cold_backend_kind == MoeHybridColdBackend::Gpu || + (expert_compute && expert_layer && + expert_compute->supports_batch_gpu_resident())); if (ok) { const auto route_read_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_get(ffn_normed, ffn_normed_host.data(), 0, sizeof(float) * ffn_normed_host.size()); - ggml_backend_tensor_get(router_probs, probs_host.data(), 0, sizeof(float) * probs_host.size()); + if (can_skip_ffn_normed_host) { + const Ds4TensorReadback readbacks[] = { + {router_probs, hybrid_probs_host.data(), + sizeof(float) * hybrid_probs_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } else { + const Ds4TensorReadback readbacks[] = { + {ffn_normed, hybrid_ffn_normed_host.data(), + sizeof(float) * hybrid_ffn_normed_host.size()}, + {router_probs, hybrid_probs_host.data(), + sizeof(float) * hybrid_probs_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } if (telemetry) telemetry->route_read_us += ds4_elapsed_us(route_read_t0, Ds4TimingClock::now()); } ggml_gallocr_free(ffn_alloc); @@ -2420,8 +3754,6 @@ static bool deepseek4_step_hybrid( return false; } - std::vector selected_host((size_t)w.n_expert_used * (size_t)n_tokens); - std::vector weights_host((size_t)w.n_expert_used * (size_t)n_tokens); const auto route_select_t0 = Ds4TimingClock::now(); const auto & hash_table = hash_routing_tables[(size_t)il].ids; for (int ti = 0; ti < n_tokens; ++ti) { @@ -2432,41 +3764,186 @@ static bool deepseek4_step_hybrid( if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; } - const int32_t * row = hash_table.data() + (size_t)tok * (size_t)w.n_expert_used; - float sum = 0.0f; - for (int ei = 0; ei < w.n_expert_used; ++ei) { - const int32_t expert = row[ei]; - selected_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei] = expert; - float prob = 0.0f; - if (expert >= 0 && expert < w.n_expert) { - prob = probs_host[(size_t)ti * (size_t)w.n_expert + (size_t)expert]; + } + ds4_parallel_for_tokens(n_tokens, 16, [&](int t0, int t1) { + for (int ti = t0; ti < t1; ++ti) { + const int32_t tok = token_ids[ti]; + const int32_t * row = hash_table.data() + (size_t)tok * (size_t)w.n_expert_used; + float sum = 0.0f; + for (int ei = 0; ei < w.n_expert_used; ++ei) { + const int32_t expert = row[ei]; + hybrid_selected_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei] = expert; + float prob = 0.0f; + if (expert >= 0 && expert < w.n_expert) { + prob = hybrid_probs_host[(size_t)ti * (size_t)w.n_expert + (size_t)expert]; + } + hybrid_weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei] = prob; + sum += prob; + } + sum = std::max(sum, 6.103515625e-5f); + for (int ei = 0; ei < w.n_expert_used; ++ei) { + float & weight = hybrid_weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei]; + weight = weight / sum * w.expert_weight_scale; } - weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei] = prob; - sum += prob; } - sum = std::max(sum, 6.103515625e-5f); - for (int ei = 0; ei < w.n_expert_used; ++ei) { - float & weight = weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei]; - weight = weight / sum * w.expert_weight_scale; + }); + if (telemetry) telemetry->route_select_us += ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); + if (routing_stats) { + for (int ti = 0; ti < n_tokens; ++ti) { + routing_stats->observe(il, + hybrid_selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, + w.n_expert_used); } } - if (telemetry) telemetry->route_select_us += ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); + + if (!eval_ds4_hybrid_or_worker( + backend, cpu_backend, hybrid_cfg, desc, storage, expert_worker, + il, n_embd, w.n_expert_used, + can_skip_ffn_normed_host ? nullptr : hybrid_ffn_normed_host.data(), + ffn_normed, + hybrid_selected_host.data(), hybrid_weights_host.data(), + n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, + expert_compute, expert_layer, + worker_owns_hot_ids, telemetry, + use_backend_hc_post_graph ? &ffn_out_backend : nullptr)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + } else if (use_cached_decode || n_tokens > 1) { + auto & cached = cached_decode_route_graphs[(size_t)il]; + if (!cached.valid() || + cached.owner_ctx != w.ctx || + cached.backend != backend || + cached.layer_idx != il || + cached.n_tokens != n_tokens || + cached.hash_routed) { + const auto route_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_route_graph(cached, backend, w, L, il, n_tokens, false)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (telemetry) telemetry->route_build_us += ds4_elapsed_us(route_build_t0, Ds4TimingClock::now()); + } + + if (ffn_in_backend) { + ggml_backend_tensor_copy(ffn_in_backend, cached.sg.inp_embed); + } else { + ggml_backend_tensor_set(cached.sg.inp_embed, ffn_working.data(), 0, + sizeof(float) * ffn_working.size()); + } + + const auto route_compute_t0 = Ds4TimingClock::now(); + bool ok = ggml_backend_graph_compute(backend, cached.sg.gf) == GGML_STATUS_SUCCESS; + if (telemetry) telemetry->route_compute_us += ds4_elapsed_us(route_compute_t0, Ds4TimingClock::now()); + if (!ok) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + + MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); + MoeLayerDesc desc = make_ds4_moe_layer_desc(L); + auto & storage = moe_hybrid.layers[(size_t) il]; + const MoeExpertLayer * expert_layer = + expert_layers ? expert_layers + il : nullptr; + const bool prefer_single_hip_route_readback = + n_tokens == 1 && + ds4_backend_is_hip(backend) && + !use_gpu_resident_decode_ffn; + const auto route_read_t0 = Ds4TimingClock::now(); + if (prefer_single_hip_route_readback) { + // On HIP decode, a second D2H fence for ffn_normed is more + // expensive than eagerly reading one token's normalized input. + const Ds4TensorReadback readbacks[] = { + {cached.selected, hybrid_selected_host.data(), + sizeof(int32_t) * hybrid_selected_host.size()}, + {cached.weights, hybrid_weights_host.data(), + sizeof(float) * hybrid_weights_host.size()}, + {cached.ffn_normed, hybrid_ffn_normed_host.data(), + sizeof(float) * hybrid_ffn_normed_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } else if (!use_gpu_resident_decode_ffn) { + const Ds4TensorReadback readbacks[] = { + {cached.selected, hybrid_selected_host.data(), + sizeof(int32_t) * hybrid_selected_host.size()}, + {cached.weights, hybrid_weights_host.data(), + sizeof(float) * hybrid_weights_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } else { + // Keep cached-decode route readback to a single D2H sync. On the + // all-hot path this may fetch a small, unused weights buffer, but + // that is cheaper than paying a second per-layer sync. + const Ds4TensorReadback readbacks[] = { + {cached.selected, hybrid_selected_host.data(), + sizeof(int32_t) * hybrid_selected_host.size()}, + {cached.weights, hybrid_weights_host.data(), + sizeof(float) * hybrid_weights_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } + const bool all_selected_hot = ds4_all_selected_hot( + storage, hybrid_selected_host.data(), n_tokens * w.n_expert_used); + const bool can_skip_ffn_normed_host = + expert_worker == nullptr && + (all_selected_hot || + storage.cold_backend_kind == MoeHybridColdBackend::Gpu || + (expert_compute && expert_layer && + expert_compute->supports_batch_gpu_resident())); + if (!prefer_single_hip_route_readback && + !use_gpu_resident_decode_ffn && + !can_skip_ffn_normed_host) { + const Ds4TensorReadback readbacks[] = { + {cached.ffn_normed, hybrid_ffn_normed_host.data(), + sizeof(float) * hybrid_ffn_normed_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } + if (telemetry) telemetry->route_read_us += ds4_elapsed_us(route_read_t0, Ds4TimingClock::now()); + if (routing_stats) { for (int ti = 0; ti < n_tokens; ++ti) { routing_stats->observe(il, - selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, + hybrid_selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, w.n_expert_used); } } - MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); - MoeLayerDesc desc = make_ds4_moe_layer_desc(L); - auto & storage = moe_hybrid.layers[(size_t) il]; - if (!eval_ds4_hybrid( - backend, cpu_backend, hybrid_cfg, desc, &moe_hybrid, storage, stream_engine, + if (!ds4_try_gpu_resident_decode_ffn( + use_gpu_resident_decode_ffn, + backend, cpu_backend, w, + hybrid_cfg, desc, storage, + cached.ffn_normed, + cached_decode_gpu_ffn_state, + hybrid_selected_host.data(), + all_selected_hot ? nullptr : hybrid_weights_host.data(), + w.n_expert_used, + cached.weights, + expert_compute, expert_layer, + &ffn_out_backend, + telemetry)) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + if (!ffn_out_backend && + !eval_ds4_hybrid_or_worker( + backend, cpu_backend, hybrid_cfg, desc, storage, expert_worker, il, n_embd, w.n_expert_used, - ffn_normed_host.data(), selected_host.data(), weights_host.data(), - n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, telemetry)) { + can_skip_ffn_normed_host ? nullptr : hybrid_ffn_normed_host.data(), + cached.ffn_normed, + hybrid_selected_host.data(), hybrid_weights_host.data(), + n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, + expert_compute, expert_layer, + worker_owns_hot_ids, telemetry, + use_backend_hc_post_graph ? &ffn_out_backend : nullptr)) { if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; @@ -2500,9 +3977,13 @@ static bool deepseek4_step_hybrid( return false; } if (telemetry) telemetry->route_build_us += ds4_elapsed_us(route_build_t0, Ds4TimingClock::now()); - ggml_backend_tensor_set(ffn_inp, ffn_working.data(), 0, sizeof(float) * ffn_working.size()); + if (ffn_in_backend) { + ggml_backend_tensor_copy(ffn_in_backend, ffn_inp); + } else { + ggml_backend_tensor_set(ffn_inp, ffn_working.data(), 0, sizeof(float) * ffn_working.size()); + } const auto route_compute_t0 = Ds4TimingClock::now(); - ok = ggml_backend_graph_compute(backend, ffn_gf) == GGML_STATUS_SUCCESS; + bool ok = ggml_backend_graph_compute(backend, ffn_gf) == GGML_STATUS_SUCCESS; if (telemetry) telemetry->route_compute_us += ds4_elapsed_us(route_compute_t0, Ds4TimingClock::now()); if (!ok) { ggml_gallocr_free(ffn_alloc); ggml_free(ffn_ctx); @@ -2511,76 +3992,97 @@ static bool deepseek4_step_hybrid( return false; } - std::vector ffn_normed_host((size_t)n_embd * (size_t)n_tokens); - std::vector probs_host((size_t)w.n_expert * (size_t)n_tokens); - std::vector selected_host((size_t)w.n_expert_used * (size_t)n_tokens); - std::vector weights_host((size_t)w.n_expert_used * (size_t)n_tokens); + MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); + MoeLayerDesc desc = make_ds4_moe_layer_desc(L); + auto & storage = moe_hybrid.layers[(size_t) il]; + const MoeExpertLayer * expert_layer = + expert_layers ? expert_layers + il : nullptr; + const bool can_skip_ffn_normed_host = + expert_worker == nullptr && + (storage.cold_backend_kind == MoeHybridColdBackend::Gpu || + (expert_compute && expert_layer && + expert_compute->supports_batch_gpu_resident())); const auto route_read_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_get(ffn_normed, ffn_normed_host.data(), 0, sizeof(float) * ffn_normed_host.size()); - ggml_backend_tensor_get(router_probs, probs_host.data(), 0, sizeof(float) * probs_host.size()); + if (can_skip_ffn_normed_host) { + const Ds4TensorReadback route_readbacks[] = { + {router_probs, hybrid_probs_host.data(), + sizeof(float) * hybrid_probs_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, route_readbacks, sizeof(route_readbacks) / sizeof(route_readbacks[0])); + } else { + const Ds4TensorReadback route_readbacks[] = { + {ffn_normed, hybrid_ffn_normed_host.data(), + sizeof(float) * hybrid_ffn_normed_host.size()}, + {router_probs, hybrid_probs_host.data(), + sizeof(float) * hybrid_probs_host.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, route_readbacks, sizeof(route_readbacks) / sizeof(route_readbacks[0])); + } if (telemetry) telemetry->route_read_us += ds4_elapsed_us(route_read_t0, Ds4TimingClock::now()); ggml_gallocr_free(ffn_alloc); ggml_free(ffn_ctx); - std::vector bias_host; const auto route_select_t0 = Ds4TimingClock::now(); - if (L.ffn_exp_probs_b) { - bias_host.resize((size_t)w.n_expert); - ggml_backend_tensor_get(L.ffn_exp_probs_b, bias_host.data(), 0, - sizeof(float) * bias_host.size()); - } - for (int ti = 0; ti < n_tokens; ++ti) { - const float * probs = probs_host.data() + (size_t)ti * (size_t)w.n_expert; + const float * route_bias_host = ds4_get_or_load_route_bias_host(w, L, lc); + ds4_parallel_for_tokens(n_tokens, 16, [&](int t0, int t1) { std::vector top((size_t)w.n_expert_used, -1); - for (int expert = 0; expert < w.n_expert; ++expert) { - const float score = probs[expert] + - (!bias_host.empty() ? bias_host[(size_t)expert] : 0.0f); - for (int slot = 0; slot < w.n_expert_used; ++slot) { - const int32_t cur_expert = top[(size_t)slot]; - const float cur_score = cur_expert >= 0 - ? probs[cur_expert] + - (!bias_host.empty() ? bias_host[(size_t)cur_expert] : 0.0f) - : -INFINITY; - if (cur_expert < 0 || score > cur_score) { - for (int m = w.n_expert_used - 1; m > slot; --m) { - top[(size_t)m] = top[(size_t)m - 1]; + for (int ti = t0; ti < t1; ++ti) { + const float * probs = hybrid_probs_host.data() + (size_t)ti * (size_t)w.n_expert; + std::fill(top.begin(), top.end(), -1); + for (int expert = 0; expert < w.n_expert; ++expert) { + const float score = probs[expert] + + (route_bias_host ? route_bias_host[(size_t)expert] : 0.0f); + for (int slot = 0; slot < w.n_expert_used; ++slot) { + const int32_t cur_expert = top[(size_t)slot]; + const float cur_score = cur_expert >= 0 + ? probs[cur_expert] + + (route_bias_host ? route_bias_host[(size_t)cur_expert] : 0.0f) + : -INFINITY; + if (cur_expert < 0 || score > cur_score) { + for (int m = w.n_expert_used - 1; m > slot; --m) { + top[(size_t)m] = top[(size_t)m - 1]; + } + top[(size_t)slot] = expert; + break; } - top[(size_t)slot] = expert; - break; } } + float sum = 0.0f; + for (int slot = 0; slot < w.n_expert_used; ++slot) { + const int32_t expert = top[(size_t)slot]; + hybrid_selected_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot] = expert; + const float weight = expert >= 0 ? probs[expert] : 0.0f; + hybrid_weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot] = weight; + sum += weight; + } + sum = std::max(sum, 6.103515625e-5f); + for (int slot = 0; slot < w.n_expert_used; ++slot) { + float & weight = hybrid_weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot]; + weight = weight / sum * w.expert_weight_scale; + } } - float sum = 0.0f; - for (int slot = 0; slot < w.n_expert_used; ++slot) { - const int32_t expert = top[(size_t)slot]; - selected_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot] = expert; - const float weight = expert >= 0 ? probs[expert] : 0.0f; - weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot] = weight; - sum += weight; - } - sum = std::max(sum, 6.103515625e-5f); - for (int slot = 0; slot < w.n_expert_used; ++slot) { - float & weight = weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot]; - weight = weight / sum * w.expert_weight_scale; - } - } + }); if (telemetry) telemetry->route_select_us += ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); if (routing_stats) { for (int ti = 0; ti < n_tokens; ++ti) { routing_stats->observe(il, - selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, + hybrid_selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, w.n_expert_used); } } - MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); - MoeLayerDesc desc = make_ds4_moe_layer_desc(L); - auto & storage = moe_hybrid.layers[(size_t) il]; - if (!eval_ds4_hybrid( - backend, cpu_backend, hybrid_cfg, desc, &moe_hybrid, storage, stream_engine, + if (!eval_ds4_hybrid_or_worker( + backend, cpu_backend, hybrid_cfg, desc, storage, expert_worker, il, n_embd, w.n_expert_used, - ffn_normed_host.data(), selected_host.data(), weights_host.data(), - n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, telemetry)) { + can_skip_ffn_normed_host ? nullptr : hybrid_ffn_normed_host.data(), + ffn_normed, + hybrid_selected_host.data(), hybrid_weights_host.data(), + n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, + expert_compute, expert_layer, + worker_owns_hot_ids, telemetry, + use_backend_hc_post_graph ? &ffn_out_backend : nullptr)) { if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; @@ -2589,76 +4091,136 @@ static bool deepseek4_step_hybrid( // ── HC post (FFN) ─────────────────────────────────────────── const auto hc_post_ffn_t0 = Ds4TimingClock::now(); - if (hc_lw.ffn.loaded && n_tokens == 1) { + if (use_backend_hc_post_graph) { + if (hc_state_backend != hc_post_graph->residual_hc) { + ggml_backend_tensor_copy(hc_state_backend, hc_post_graph->residual_hc); + } + if (ffn_out_backend) { + ggml_backend_tensor_copy(ffn_out_backend, hc_post_graph->block_out); + } else { + ggml_backend_tensor_set(hc_post_graph->block_out, + ffn_out_host.data(), 0, + sizeof(float) * ffn_out_host.size()); + } + ggml_backend_tensor_copy(ffn_post_backend, hc_post_graph->post); + ggml_backend_tensor_copy(ffn_comb_backend, hc_post_graph->comb); + if (ggml_backend_graph_compute(backend, hc_post_graph->sg.gf) != GGML_STATUS_SUCCESS) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + hc_state_backend = hc_post_graph->sg.hidden_states; + } else if (hc_lw.ffn.loaded && n_tokens == 1) { std::vector new_hc((size_t)hc_dim); cpu_hc_post(new_hc.data(), ffn_out_host.data(), hc_state.data(), hc_ffn_result.post, hc_ffn_result.comb, n_embd, n_hc); memcpy(hc_state.data(), new_hc.data(), (size_t)hc_dim * sizeof(float)); + } else if (hc_lw.ffn.loaded) { + hc_post_batch(next_hc, + ffn_out_host.data(), + hc_state.data(), + hc_post.data(), + hc_comb.data(), + n_tokens, + n_embd, + n_hc); + std::memcpy(hc_state.data(), next_hc.data(), + next_hc.size() * sizeof(float)); } else { for (int i = 0; i < n_embd * n_tokens; i++) { hc_state[(size_t)i] += ffn_out_host[(size_t)i]; } } if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); + if (trace_decode) { + std::fprintf(stderr, + "[deepseek4-trace] decode-step kv=%d layer=%d end\n", + kv_start, il); + } } if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); + const bool use_output_hc = + w.output_hc_fn && w.output_hc_scale && w.output_hc_base; + + if ((!use_cached_decode || !use_output_hc) && + use_backend_hc_post_graph && hc_state_backend) { + ggml_backend_tensor_get(hc_state_backend, hc_state.data(), 0, + sizeof(float) * hc_state.size()); + } + // ── Output HC pre → norm → logits ─────────────────────────────────── const auto output_t0 = Ds4TimingClock::now(); - std::vector final_embd((size_t)n_embd * (size_t)n_tokens); - if (hc_output_weights.loaded && n_tokens == 1) { - std::vector flat((size_t)hc_dim); - cpu_rms_norm(flat.data(), hc_state.data(), hc_dim, w.hc_eps); - std::vector pre(n_hc); - cpu_matvec_f16(pre.data(), hc_output_weights.fn_data.data(), flat.data(), n_hc, hc_dim); - float hc_weights[4]; - for (int i = 0; i < n_hc; i++) { - const float z = pre[i] * hc_output_weights.scale_data[0] + hc_output_weights.base_data[i]; - hc_weights[i] = 1.0f / (1.0f + expf(-z)) + 1.0e-6f; + if (!use_cached_decode || !use_output_hc) { + if (hc_output_weights.loaded) { + hc_output_batch(final_embd, hc_state.data(), hc_output_weights, + n_tokens, n_embd, n_hc, w.hc_eps); + } else { + final_embd.resize((size_t)n_embd * (size_t)n_tokens); + memcpy(final_embd.data(), hc_state.data(), + (size_t)n_embd * (size_t)n_tokens * sizeof(float)); } - for (int d = 0; d < n_embd; d++) { - float acc = 0.0f; - for (int h = 0; h < n_hc; h++) { - acc += hc_weights[h] * hc_state[(size_t)h * n_embd + d]; + } + + bool final_ok = false; + if (use_cached_decode) { + if (!cached_decode_output_graph.valid() || + cached_decode_output_graph.owner_ctx != w.ctx || + cached_decode_output_graph.backend != backend || + cached_decode_output_graph.n_tokens != n_tokens) { + if (!build_cached_decode_output_graph(cached_decode_output_graph, backend, w, n_tokens)) { + return false; } - final_embd[d] = acc; + } + if (use_output_hc && hc_state_backend) { + ggml_backend_tensor_copy(hc_state_backend, + cached_decode_output_graph.sg.hidden_input); + } else { + ggml_backend_tensor_set(cached_decode_output_graph.sg.hidden_input, + final_embd.data(), 0, sizeof(float) * final_embd.size()); + } + final_ok = ggml_backend_graph_compute(backend, cached_decode_output_graph.sg.gf) == + GGML_STATUS_SUCCESS; + if (final_ok) { + out_logits.resize((size_t)w.n_vocab); + ggml_backend_tensor_get(cached_decode_output_graph.sg.logits, + out_logits.data(), 0, + sizeof(float) * (size_t)w.n_vocab); } } else { - memcpy(final_embd.data(), hc_state.data(), (size_t)n_embd * (size_t)n_tokens * sizeof(float)); - } - - const size_t final_ctx_size = 16 * 1024 * 1024; - ggml_init_params params2{}; - params2.mem_size = final_ctx_size; - params2.mem_buffer = nullptr; - params2.no_alloc = true; - ggml_context * ctx2 = ggml_init(params2); - if (!ctx2) return false; - - ggml_tensor * final_inp = ggml_new_tensor_2d(ctx2, GGML_TYPE_F32, n_embd, n_tokens); - ggml_set_input(final_inp); - ggml_tensor * normed_out = build_rms_norm(ctx2, final_inp, w.out_norm, w.rms_eps); - ggml_tensor * logits = ggml_mul_mat(ctx2, w.output, normed_out); - ggml_cgraph * final_gf = ggml_new_graph(ctx2); - ggml_build_forward_expand(final_gf, logits); - ggml_gallocr_t final_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(final_alloc, final_gf)) { + const size_t final_ctx_size = 16 * 1024 * 1024; + ggml_init_params params2{}; + params2.mem_size = final_ctx_size; + params2.mem_buffer = nullptr; + params2.no_alloc = true; + ggml_context * ctx2 = ggml_init(params2); + if (!ctx2) return false; + + ggml_tensor * final_inp = ggml_new_tensor_2d(ctx2, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(final_inp); + ggml_tensor * normed_out = build_rms_norm(ctx2, final_inp, w.out_norm, w.rms_eps); + ggml_tensor * logits = ggml_mul_mat(ctx2, w.output, normed_out); + ggml_cgraph * final_gf = ggml_new_graph(ctx2); + ggml_build_forward_expand(final_gf, logits); + ggml_gallocr_t final_alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!ggml_gallocr_alloc_graph(final_alloc, final_gf)) { + ggml_gallocr_free(final_alloc); + ggml_free(ctx2); + return false; + } + ggml_backend_tensor_set(final_inp, final_embd.data(), 0, sizeof(float) * final_embd.size()); + final_ok = ggml_backend_graph_compute(backend, final_gf) == GGML_STATUS_SUCCESS; + if (final_ok) { + out_logits.resize((size_t)w.n_vocab); + const size_t logits_offset = (size_t)(n_tokens - 1) * (size_t)w.n_vocab * sizeof(float); + ggml_backend_tensor_get(logits, out_logits.data(), logits_offset, + sizeof(float) * (size_t)w.n_vocab); + } ggml_gallocr_free(final_alloc); ggml_free(ctx2); - return false; } - ggml_backend_tensor_set(final_inp, final_embd.data(), 0, sizeof(float) * final_embd.size()); - bool final_ok = ggml_backend_graph_compute(backend, final_gf) == GGML_STATUS_SUCCESS; - if (final_ok) { - out_logits.resize((size_t)w.n_vocab); - const size_t logits_offset = (size_t)(n_tokens - 1) * (size_t)w.n_vocab * sizeof(float); - ggml_backend_tensor_get(logits, out_logits.data(), logits_offset, - sizeof(float) * (size_t)w.n_vocab); - } - ggml_gallocr_free(final_alloc); - ggml_free(ctx2); if (!final_ok) return false; if (telemetry) { telemetry->output_us += ds4_elapsed_us(output_t0, Ds4TimingClock::now()); @@ -2669,6 +4231,186 @@ static bool deepseek4_step_hybrid( return true; } +namespace { + +ggml_tensor * clone_snapshot_tensor(ggml_context * ctx, + const ggml_tensor * src, + const char * name) { + if (!ctx || !src) return nullptr; + ggml_tensor * dst = ggml_dup_tensor(ctx, const_cast(src)); + if (!dst) return nullptr; + if (name && *name) ggml_set_name(dst, name); + return dst; +} + +bool copy_tensor_from_backend(const ggml_tensor * src, ggml_tensor * dst) { + if (!src || !dst) return false; + const size_t bytes = ggml_nbytes(src); + if (bytes != ggml_nbytes(dst)) return false; + ggml_backend_tensor_get(src, dst->data, 0, bytes); + return true; +} + +bool copy_tensor_to_backend(const ggml_tensor * src, ggml_tensor * dst) { + if (!src || !dst) return false; + const size_t bytes = ggml_nbytes(src); + if (bytes != ggml_nbytes(dst)) return false; + ggml_backend_tensor_set(dst, src->data, 0, bytes); + return true; +} + +bool tensors_compatible(const ggml_tensor * a, const ggml_tensor * b) { + if (!!a != !!b) return false; + if (!a) return true; + if (a->type != b->type || ggml_n_dims(a) != ggml_n_dims(b)) return false; + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (a->ne[i] != b->ne[i]) return false; + } + return true; +} + +} // namespace + +bool deepseek4_snapshot_save(const DeepSeek4Cache & cache, + ggml_backend_t snapshot_backend, + DeepSeek4Snapshot & out) { + if (!snapshot_backend || !cache.ctx || !cache.buf || !cache.hc_state || + cache.layers.size() != (size_t)cache.n_layer) { + return false; + } + + free_deepseek4_snapshot(out); + + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * (size_t)(cache.n_layer * 8 + 8) + 4096; + ip.no_alloc = true; + out.ctx = ggml_init(ip); + if (!out.ctx) { + return false; + } + + out.layers.resize((size_t)cache.n_layer); + out.hc_state_snap = clone_snapshot_tensor(out.ctx, cache.hc_state, "ds4_hc_state_snap"); + if (!out.hc_state_snap) { + free_deepseek4_snapshot(out); + return false; + } + + for (int il = 0; il < cache.n_layer; ++il) { + const auto & src = cache.layers[(size_t)il]; + auto & dst = out.layers[(size_t)il]; + dst.raw_kv = clone_snapshot_tensor(out.ctx, src.raw_kv, nullptr); + dst.comp_kv = clone_snapshot_tensor(out.ctx, src.comp_kv, nullptr); + dst.index_comp_kv = clone_snapshot_tensor(out.ctx, src.index_comp_kv, nullptr); + dst.attn_compressor.state_kv = + clone_snapshot_tensor(out.ctx, src.attn_compressor.state_kv, nullptr); + dst.attn_compressor.state_score = + clone_snapshot_tensor(out.ctx, src.attn_compressor.state_score, nullptr); + dst.indexer_compressor.state_kv = + clone_snapshot_tensor(out.ctx, src.indexer_compressor.state_kv, nullptr); + dst.indexer_compressor.state_score = + clone_snapshot_tensor(out.ctx, src.indexer_compressor.state_score, nullptr); + if (!dst.raw_kv || + (src.comp_kv && !dst.comp_kv) || + (src.index_comp_kv && !dst.index_comp_kv) || + (src.attn_compressor.state_kv && !dst.attn_compressor.state_kv) || + (src.attn_compressor.state_score && !dst.attn_compressor.state_score) || + (src.indexer_compressor.state_kv && !dst.indexer_compressor.state_kv) || + (src.indexer_compressor.state_score && !dst.indexer_compressor.state_score)) { + free_deepseek4_snapshot(out); + return false; + } + } + + out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, snapshot_backend); + if (!out.buf) { + free_deepseek4_snapshot(out); + return false; + } + + if (!copy_tensor_from_backend(cache.hc_state, out.hc_state_snap)) { + free_deepseek4_snapshot(out); + return false; + } + for (int il = 0; il < cache.n_layer; ++il) { + const auto & src = cache.layers[(size_t)il]; + auto & dst = out.layers[(size_t)il]; + dst.n_comp = src.n_comp; + dst.n_index_comp = src.n_index_comp; + if (!copy_tensor_from_backend(src.raw_kv, dst.raw_kv) || + (src.comp_kv && !copy_tensor_from_backend(src.comp_kv, dst.comp_kv)) || + (src.index_comp_kv && + !copy_tensor_from_backend(src.index_comp_kv, dst.index_comp_kv)) || + (src.attn_compressor.state_kv && + !copy_tensor_from_backend(src.attn_compressor.state_kv, + dst.attn_compressor.state_kv)) || + (src.attn_compressor.state_score && + !copy_tensor_from_backend(src.attn_compressor.state_score, + dst.attn_compressor.state_score)) || + (src.indexer_compressor.state_kv && + !copy_tensor_from_backend(src.indexer_compressor.state_kv, + dst.indexer_compressor.state_kv)) || + (src.indexer_compressor.state_score && + !copy_tensor_from_backend(src.indexer_compressor.state_score, + dst.indexer_compressor.state_score))) { + free_deepseek4_snapshot(out); + return false; + } + } + + out.cur_pos = cache.cur_pos; + return true; +} + +bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, + DeepSeek4Cache & cache) { + if (!snap.ctx || !cache.ctx || !cache.buf || !snap.hc_state_snap || + snap.layers.size() != cache.layers.size()) { + return false; + } + if (!tensors_compatible(snap.hc_state_snap, cache.hc_state) || + !copy_tensor_to_backend(snap.hc_state_snap, cache.hc_state)) { + return false; + } + + for (size_t il = 0; il < cache.layers.size(); ++il) { + const auto & src = snap.layers[il]; + auto & dst = cache.layers[il]; + if (!tensors_compatible(src.raw_kv, dst.raw_kv) || + !tensors_compatible(src.comp_kv, dst.comp_kv) || + !tensors_compatible(src.index_comp_kv, dst.index_comp_kv) || + !tensors_compatible(src.attn_compressor.state_kv, dst.attn_compressor.state_kv) || + !tensors_compatible(src.attn_compressor.state_score, dst.attn_compressor.state_score) || + !tensors_compatible(src.indexer_compressor.state_kv, dst.indexer_compressor.state_kv) || + !tensors_compatible(src.indexer_compressor.state_score, dst.indexer_compressor.state_score)) { + return false; + } + if (!copy_tensor_to_backend(src.raw_kv, dst.raw_kv) || + (src.comp_kv && !copy_tensor_to_backend(src.comp_kv, dst.comp_kv)) || + (src.index_comp_kv && + !copy_tensor_to_backend(src.index_comp_kv, dst.index_comp_kv)) || + (src.attn_compressor.state_kv && + !copy_tensor_to_backend(src.attn_compressor.state_kv, + dst.attn_compressor.state_kv)) || + (src.attn_compressor.state_score && + !copy_tensor_to_backend(src.attn_compressor.state_score, + dst.attn_compressor.state_score)) || + (src.indexer_compressor.state_kv && + !copy_tensor_to_backend(src.indexer_compressor.state_kv, + dst.indexer_compressor.state_kv)) || + (src.indexer_compressor.state_score && + !copy_tensor_to_backend(src.indexer_compressor.state_score, + dst.indexer_compressor.state_score))) { + return false; + } + dst.n_comp = src.n_comp; + dst.n_index_comp = src.n_index_comp; + } + + cache.cur_pos = snap.cur_pos; + return true; +} + // ─── Full forward step ────────────────────────────────────────────────── bool deepseek4_step( @@ -2681,21 +4423,53 @@ bool deepseek4_step( std::vector & out_logits, MoeHybridStorage * moe_hybrid, const int32_t * token_ids, - MoeHybridStreamEngine * stream_engine, + ExpertIpcClient * expert_worker, + bool worker_owns_hot_ids, + bool disable_cached_decode, DeepSeek4StepTelemetry * telemetry, - MoeHybridRoutingStats * routing_stats) { + MoeHybridRoutingStats * routing_stats, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layers) { const auto step_t0 = Ds4TimingClock::now(); + const bool direct_nonhybrid = !w.moe_hybrid && moe_hybrid == nullptr; + const bool use_full_layer_range = + ds4_env_flag("DFLASH_DS4_FULL_LAYER_RANGE"); + const bool use_decode_layer_range = + n_tokens == 1 && + ds4_env_flag("DFLASH_DS4_DECODE_LAYER_RANGE"); + if (direct_nonhybrid && + (use_full_layer_range || use_decode_layer_range)) { + // This diagnostic lane is intentionally slower than the full-step graph, + // but it exposes attn/route/ffn timing that the direct pure-UMA path + // otherwise collapses into one opaque full_compute bucket. + static thread_local std::vector full_hc_state; + return deepseek4_step_layer_range( + backend, + w, + cache, + full_hc_state, + embed, + n_tokens, + kv_start, + /*layer_begin=*/0, + /*layer_end=*/w.n_layer, + &out_logits, + token_ids, + telemetry); + } if (w.moe_hybrid && moe_hybrid != nullptr) { return deepseek4_step_hybrid(backend, w, cache, *moe_hybrid, embed, n_tokens, kv_start, out_logits, - token_ids, stream_engine, telemetry, routing_stats); + token_ids, expert_worker, worker_owns_hot_ids, + disable_cached_decode, + expert_compute, expert_layers, + telemetry, routing_stats); } const int n_embd = w.n_embd; const int n_layer = w.n_layer; - // Create compute graph context — need large budget for MoE layers const size_t ctx_size = ggml_tensor_overhead() * 65536 + 16 * 1024 * 1024; const bool reuse_full_step_decode = n_tokens == 1 && @@ -2704,7 +4478,8 @@ bool deepseek4_step( static thread_local DeepSeek4LegacyFullStepCache full_step_cache; StepGraph * cached_sg = nullptr; if (reuse_full_step_decode) { - if (full_step_cache.owner_ctx != w.ctx || full_step_cache.backend != backend) { + if (full_step_cache.owner_ctx != w.ctx || + full_step_cache.backend != backend) { full_step_cache.free(); full_step_cache.owner_ctx = w.ctx; full_step_cache.backend = backend; @@ -2712,14 +4487,14 @@ bool deepseek4_step( step_graph_free(full_step_cache.sg); } cached_sg = &full_step_cache.sg; - if (full_step_cache.meta_arena.size() < ctx_size) { - full_step_cache.meta_arena.resize(ctx_size); + if (cached_sg->meta_arena.size() < ctx_size) { + cached_sg->meta_arena.resize(ctx_size); } } ggml_init_params params{}; - params.mem_size = cached_sg ? full_step_cache.meta_arena.size() : ctx_size; - params.mem_buffer = cached_sg ? full_step_cache.meta_arena.data() : nullptr; + params.mem_size = cached_sg ? cached_sg->meta_arena.size() : ctx_size; + params.mem_buffer = cached_sg ? cached_sg->meta_arena.data() : nullptr; params.no_alloc = true; const auto full_build_t0 = Ds4TimingClock::now(); ggml_context * ctx = ggml_init(params); @@ -2727,21 +4502,24 @@ bool deepseek4_step( if (cached_sg) { cached_sg->ctx = ctx; } + if (telemetry) { + telemetry->full_graph_build_us += ds4_elapsed_us(full_build_t0, Ds4TimingClock::now()); + } ggml_gallocr_t alloc = nullptr; bool owns_alloc = false; auto release_full_step = [&]() { if (cached_sg) { step_graph_free(*cached_sg); - return; - } - if (alloc && owns_alloc) { - ggml_gallocr_free(alloc); - alloc = nullptr; - } - if (ctx) { - ggml_free(ctx); - ctx = nullptr; + } else { + if (alloc && owns_alloc) { + ggml_gallocr_free(alloc); + alloc = nullptr; + } + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } } }; @@ -2754,7 +4532,8 @@ bool deepseek4_step( } ggml_tensor * cur = inp; - ggml_cgraph * gf = ggml_new_graph_custom(ctx, 32768, false); + ggml_cgraph * gf = + ggml_new_graph_custom(ctx, ds4_full_prefill_graph_size(n_tokens), false); if (cached_sg) { cached_sg->gf = gf; } @@ -2812,9 +4591,6 @@ bool deepseek4_step( // ── Build and run graph ───────────────────────────────────────────── ggml_build_forward_expand(gf, logits); - if (telemetry) { - telemetry->full_graph_build_us += ds4_elapsed_us(full_build_t0, Ds4TimingClock::now()); - } // Allocate if (cached_sg) { @@ -2898,6 +4674,28 @@ bool deepseek4_step( return true; } +bool deepseek4_step( + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const float * embed, + int n_tokens, + int kv_start, + std::vector & out_logits, + MoeHybridStorage * moe_hybrid, + const int32_t * token_ids, + ExpertIpcClient * expert_worker, + bool worker_owns_hot_ids, + DeepSeek4StepTelemetry * telemetry, + MoeHybridRoutingStats * routing_stats, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layers) { + return deepseek4_step(backend, w, cache, embed, n_tokens, kv_start, + out_logits, moe_hybrid, token_ids, expert_worker, + worker_owns_hot_ids, false, telemetry, + routing_stats, expert_compute, expert_layers); +} + bool deepseek4_step_layer_range( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -2951,6 +4749,8 @@ bool deepseek4_step_layer_range( static std::vector cached_decode_ffn_hc_pre_graphs; static DeepSeek4CachedDecodeHcPostGraph cached_decode_hc_post_graph; static std::vector> cached_decode_attn_graphs; + static std::vector cached_decode_route_graphs_range; + static std::vector cached_decode_expert_graphs_range; static std::vector cached_decode_ffn_graphs; static DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; static DeepSeek4CachedLayerAlloc cached_dynamic_output_alloc; @@ -2977,6 +4777,14 @@ bool deepseek4_step_layer_range( } } cached_decode_attn_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_route_graphs_range) { + g.free(); + } + cached_decode_route_graphs_range.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_expert_graphs_range) { + g.free(); + } + cached_decode_expert_graphs_range.assign((size_t)w.n_layer, {}); for (auto & g : cached_decode_ffn_graphs) { g.free(); } @@ -2997,7 +4805,7 @@ bool deepseek4_step_layer_range( // Per-layer execution with CPU-side HC static thread_local DeepSeek4LayerRangeScratch scratch; - scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, w.n_expert_used); + scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, w.n_expert_used, w.n_expert); std::vector & cur = scratch.cur; std::vector & ffn_working = scratch.ffn_working; std::vector & hc_post = scratch.hc_post; @@ -3008,7 +4816,12 @@ bool deepseek4_step_layer_range( std::vector & hash_expert_ids_host = scratch.hash_expert_ids; const bool reuse_decode_graphs = (n_tokens == 1); 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 disable_backend_decode_hc_direct = + ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC_DIRECT"); + const bool use_backend_decode_hc_direct = + use_backend_decode_hc && + ds4_backend_is_hip(backend) && + !disable_backend_decode_hc_direct; const bool use_backend_decode_hc_graph = use_backend_decode_hc && !use_backend_decode_hc_direct; const ggml_tensor * hc_state_backend = nullptr; @@ -3025,6 +4838,13 @@ bool deepseek4_step_layer_range( hc_state_backend = cached_decode_hc_post_graph.residual_hc; } for (int il = layer_begin; il < layer_end; ++il) { + const bool trace_decode = + ds4_env_flag("DFLASH_DS4_TRACE_DECODE") && n_tokens == 1; + if (trace_decode) { + std::fprintf(stderr, + "[deepseek4-trace] decode-step kv=%d layer=%d begin\n", + kv_start, il); + } const DeepSeek4Layer & L = w.layers[(size_t)il]; DeepSeek4LayerCache & lc = cache.layers[(size_t)il]; const HcLayerWeightsCpu & hc_lw = hc_layer_weights_range[(size_t)il]; @@ -3113,6 +4933,7 @@ bool deepseek4_step_layer_range( n_tokens, n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); } if (telemetry) telemetry->hc_pre_attn_us += ds4_elapsed_us(hc_pre_attn_t0, Ds4TimingClock::now()); + ds4_trace_decode_marker(trace_decode, kv_start, il, "hc_pre_attn_done"); // ── Build & run attention graph ───────────────────────────── { @@ -3127,6 +4948,13 @@ bool deepseek4_step_layer_range( const int n_raw = std::min(kv_start + 1, 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_index_comp = (ratio == 4) ? ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos) : 0; + const bool flush_boundary = ratio > 0 && ((token_pos + 1) % ratio) == 0; + if (trace_decode) { + std::fprintf(stderr, + "[deepseek4-trace] decode-step kv=%d layer=%d attn-shape raw=%d comp=%d index=%d flush=%d ratio=%d\n", + kv_start, il, n_raw, n_comp_attn, n_index_comp, + flush_boundary ? 1 : 0, ratio); + } auto & per_layer = cached_decode_attn_graphs[(size_t)il]; auto it = std::find_if(per_layer.begin(), per_layer.end(), [&](const DeepSeek4CachedDecodeAttnGraph & candidate) { @@ -3136,7 +4964,8 @@ bool deepseek4_step_layer_range( candidate.layer_idx == il && candidate.n_raw == n_raw && candidate.n_comp_attn == n_comp_attn && - candidate.n_index_comp == n_index_comp; + candidate.n_index_comp == n_index_comp && + candidate.flush_boundary == flush_boundary; }); if (it == per_layer.end()) { if (per_layer.size() >= 20) { @@ -3147,7 +4976,8 @@ bool deepseek4_step_layer_range( auto & candidate = per_layer.back(); const auto attn_build_t0 = Ds4TimingClock::now(); if (!build_cached_decode_attn_graph(candidate, backend, w, L, lc, il, kv_start, - n_raw, n_comp_attn, n_index_comp)) { + n_raw, n_comp_attn, n_index_comp, + flush_boundary)) { std::fprintf(stderr, "[deepseek4] cached attn graph alloc failed layer %d\n", il); return false; } @@ -3172,10 +5002,9 @@ 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; ggml_backend_tensor_set(cached_attn->inputs.attn_ape_row, &ape_row, 0, sizeof(ape_row)); ggml_backend_tensor_set(cached_attn->inputs.attn_state_rows, &state_row, 0, sizeof(state_row)); if (flush_boundary) { @@ -3186,10 +5015,9 @@ bool deepseek4_step_layer_range( if (ratio == 4) { 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; ggml_backend_tensor_set(cached_attn->inputs.index_ape_row, &ape_row, 0, sizeof(ape_row)); ggml_backend_tensor_set(cached_attn->inputs.index_state_rows, &state_row, 0, sizeof(state_row)); if (flush_boundary) { @@ -3273,10 +5101,19 @@ bool deepseek4_step_layer_range( if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); } else { const auto attn_read_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_get(attn_out, attn_out_host.data(), 0, sizeof(float) * attn_out_host.size()); if (use_backend_decode_hc_direct) { - ggml_backend_tensor_get(attn_post_backend, hc_post.data(), 0, sizeof(float) * hc_post.size()); - ggml_backend_tensor_get(attn_comb_backend, hc_comb.data(), 0, sizeof(float) * hc_comb.size()); + const Ds4TensorReadback readbacks[] = { + {attn_out, attn_out_host.data(), + sizeof(float) * attn_out_host.size()}, + {attn_post_backend, hc_post.data(), + sizeof(float) * hc_post.size()}, + {attn_comb_backend, hc_comb.data(), + sizeof(float) * hc_comb.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } else { + ggml_backend_tensor_get(attn_out, attn_out_host.data(), 0, sizeof(float) * attn_out_host.size()); } if (telemetry) telemetry->attn_read_us += ds4_elapsed_us(attn_read_t0, Ds4TimingClock::now()); } @@ -3297,6 +5134,7 @@ bool deepseek4_step_layer_range( if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); } } + ds4_trace_decode_marker(trace_decode, kv_start, il, "attn_done"); // ── HC pre (FFN) ──────────────────────────────────────────── const auto hc_pre_ffn_t0 = Ds4TimingClock::now(); @@ -3395,38 +5233,98 @@ bool deepseek4_step_layer_range( } auto & cached = cached_decode_ffn_graphs[(size_t)il]; - if (!cached.valid() || - cached.owner_ctx != w.ctx || - cached.backend != backend || - cached.layer_idx != il || - cached.n_tokens != n_tokens || - cached.hash_routed != hash_routed) { - const auto ffn_build_t0 = Ds4TimingClock::now(); - if (!build_cached_decode_ffn_graph(cached, backend, w, L, il, n_tokens, hash_routed)) { - std::fprintf(stderr, "[deepseek4] cached ffn graph alloc failed layer %d\n", il); + ggml_tensor * ffn_out = nullptr; + if (!hash_routed) { + auto & route = cached_decode_route_graphs_range[(size_t)il]; + if (!route.valid() || + route.owner_ctx != w.ctx || + route.backend != backend || + route.layer_idx != il || + route.n_tokens != n_tokens || + route.hash_routed) { + const auto route_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_route_graph(route, backend, w, L, il, n_tokens, false)) { + std::fprintf(stderr, "[deepseek4] cached route graph alloc failed layer %d\n", il); + return false; + } + if (telemetry) telemetry->route_build_us += ds4_elapsed_us(route_build_t0, Ds4TimingClock::now()); + } + + if (ffn_in_backend) { + ggml_backend_tensor_copy(ffn_in_backend, route.sg.inp_embed); + } else { + ggml_backend_tensor_set(route.sg.inp_embed, ffn_working.data(), 0, + sizeof(float) * ffn_working.size()); + } + + const auto route_compute_t0 = Ds4TimingClock::now(); + auto route_status = ggml_backend_graph_compute(backend, route.sg.gf); + if (telemetry) telemetry->route_compute_us += ds4_elapsed_us(route_compute_t0, Ds4TimingClock::now()); + if (route_status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[deepseek4] cached route compute failed layer %d\n", il); return false; } - if (telemetry) telemetry->ffn_build_us += ds4_elapsed_us(ffn_build_t0, Ds4TimingClock::now()); - } - ggml_tensor * ffn_out = cached.sg.hidden_states; - if (ffn_in_backend) { - ggml_backend_tensor_copy(ffn_in_backend, cached.sg.inp_embed); + auto & expert = cached_decode_expert_graphs_range[(size_t)il]; + if (!expert.valid() || + expert.owner_ctx != w.ctx || + expert.backend != backend || + expert.layer_idx != il || + expert.n_tokens != n_tokens) { + const auto ffn_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_expert_graph(expert, backend, w, L, il, n_tokens)) { + std::fprintf(stderr, "[deepseek4] cached expert graph alloc failed layer %d\n", il); + return false; + } + if (telemetry) telemetry->ffn_build_us += ds4_elapsed_us(ffn_build_t0, Ds4TimingClock::now()); + } + + ggml_backend_tensor_copy(route.ffn_normed, expert.sg.inp_embed); + ggml_backend_tensor_copy(route.selected, expert.selected); + ggml_backend_tensor_copy(route.weights, expert.weights); + + const auto ffn_compute_t0 = Ds4TimingClock::now(); + auto expert_status = ggml_backend_graph_compute(backend, expert.sg.gf); + if (telemetry) telemetry->ffn_compute_us += ds4_elapsed_us(ffn_compute_t0, Ds4TimingClock::now()); + if (expert_status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[deepseek4] cached expert compute failed layer %d\n", il); + return false; + } + ffn_out = expert.sg.hidden_states; } else { - ggml_backend_tensor_set(cached.sg.inp_embed, ffn_working.data(), 0, - sizeof(float) * ffn_working.size()); - } - if (cached.hash_ids) { - ggml_backend_tensor_set(cached.hash_ids, hash_expert_ids_host.data(), 0, - sizeof(int32_t) * hash_expert_ids_host.size()); - } + if (!cached.valid() || + cached.owner_ctx != w.ctx || + cached.backend != backend || + cached.layer_idx != il || + cached.n_tokens != n_tokens || + cached.hash_routed != hash_routed) { + const auto ffn_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_ffn_graph(cached, backend, w, L, il, n_tokens, hash_routed)) { + std::fprintf(stderr, "[deepseek4] cached ffn graph alloc failed layer %d\n", il); + return false; + } + if (telemetry) telemetry->ffn_build_us += ds4_elapsed_us(ffn_build_t0, Ds4TimingClock::now()); + } - const auto ffn_compute_t0 = Ds4TimingClock::now(); - auto status = ggml_backend_graph_compute(backend, cached.sg.gf); - if (telemetry) telemetry->ffn_compute_us += ds4_elapsed_us(ffn_compute_t0, Ds4TimingClock::now()); - if (status != GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "[deepseek4] cached ffn compute failed layer %d\n", il); - return false; + ffn_out = cached.sg.hidden_states; + if (ffn_in_backend) { + ggml_backend_tensor_copy(ffn_in_backend, cached.sg.inp_embed); + } else { + ggml_backend_tensor_set(cached.sg.inp_embed, ffn_working.data(), 0, + sizeof(float) * ffn_working.size()); + } + if (cached.hash_ids) { + ggml_backend_tensor_set(cached.hash_ids, hash_expert_ids_host.data(), 0, + sizeof(int32_t) * hash_expert_ids_host.size()); + } + + const auto ffn_compute_t0 = Ds4TimingClock::now(); + auto status = ggml_backend_graph_compute(backend, cached.sg.gf); + if (telemetry) telemetry->ffn_compute_us += ds4_elapsed_us(ffn_compute_t0, Ds4TimingClock::now()); + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[deepseek4] cached ffn compute failed layer %d\n", il); + return false; + } } if (use_backend_decode_hc_graph) { @@ -3445,10 +5343,19 @@ bool deepseek4_step_layer_range( if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); } else { const auto ffn_read_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_get(ffn_out, ffn_out_host.data(), 0, sizeof(float) * ffn_out_host.size()); if (use_backend_decode_hc_direct) { - ggml_backend_tensor_get(ffn_post_backend, hc_post.data(), 0, sizeof(float) * hc_post.size()); - ggml_backend_tensor_get(ffn_comb_backend, hc_comb.data(), 0, sizeof(float) * hc_comb.size()); + const Ds4TensorReadback readbacks[] = { + {ffn_out, ffn_out_host.data(), + sizeof(float) * ffn_out_host.size()}, + {ffn_post_backend, hc_post.data(), + sizeof(float) * hc_post.size()}, + {ffn_comb_backend, hc_comb.data(), + sizeof(float) * hc_comb.size()}, + }; + ds4_tensor_gets_async_and_sync( + backend, readbacks, sizeof(readbacks) / sizeof(readbacks[0])); + } else { + ggml_backend_tensor_get(ffn_out, ffn_out_host.data(), 0, sizeof(float) * ffn_out_host.size()); } if (telemetry) telemetry->ffn_read_us += ds4_elapsed_us(ffn_read_t0, Ds4TimingClock::now()); } @@ -3470,7 +5377,11 @@ bool deepseek4_step_layer_range( } } - if (use_backend_decode_hc_graph && hc_state_backend) { + const bool use_output_hc = + w.output_hc_fn && w.output_hc_scale && w.output_hc_base; + + if ((!is_last_shard || !reuse_decode_graphs || !use_output_hc) && + use_backend_decode_hc_graph && hc_state_backend) { ggml_backend_tensor_get(hc_state_backend, hc_state.data(), 0, sizeof(float) * hc_state.size()); } @@ -3479,13 +5390,15 @@ bool deepseek4_step_layer_range( // 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 (!reuse_decode_graphs || !use_output_hc) { + 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() || @@ -3496,8 +5409,13 @@ bool deepseek4_step_layer_range( return false; } } - ggml_backend_tensor_set(cached_decode_output_graph.sg.hidden_input, - final_embd.data(), 0, sizeof(float) * final_embd.size()); + if (use_output_hc && hc_state_backend) { + ggml_backend_tensor_copy(hc_state_backend, + cached_decode_output_graph.sg.hidden_input); + } else { + ggml_backend_tensor_set(cached_decode_output_graph.sg.hidden_input, + final_embd.data(), 0, sizeof(float) * final_embd.size()); + } if (ggml_backend_graph_compute(backend, cached_decode_output_graph.sg.gf) != GGML_STATUS_SUCCESS) { return false; } @@ -3658,196 +5576,16 @@ bool create_deepseek4_cache(ggml_backend_t backend, return true; } -namespace { - -ggml_tensor * clone_snapshot_tensor(ggml_context * ctx, - const ggml_tensor * src, - const char * name) { - if (!ctx || !src) return nullptr; - ggml_tensor * dst = ggml_dup_tensor(ctx, const_cast(src)); - if (!dst) return nullptr; - if (name && *name) ggml_set_name(dst, name); - return dst; -} - -bool copy_tensor_from_backend(const ggml_tensor * src, ggml_tensor * dst) { - if (!src || !dst) return false; - const size_t bytes = ggml_nbytes(src); - if (bytes != ggml_nbytes(dst)) return false; - ggml_backend_tensor_get(src, dst->data, 0, bytes); - return true; -} - -bool copy_tensor_to_backend(const ggml_tensor * src, ggml_tensor * dst) { - if (!src || !dst) return false; - const size_t bytes = ggml_nbytes(src); - if (bytes != ggml_nbytes(dst)) return false; - ggml_backend_tensor_set(dst, src->data, 0, bytes); - return true; -} - -bool tensors_compatible(const ggml_tensor * a, const ggml_tensor * b) { - if (!!a != !!b) return false; - if (!a) return true; - if (a->type != b->type || ggml_n_dims(a) != ggml_n_dims(b)) return false; - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - if (a->ne[i] != b->ne[i]) return false; - } - return true; -} - -} // namespace - -bool deepseek4_snapshot_save(const DeepSeek4Cache & cache, - ggml_backend_t snapshot_backend, - DeepSeek4Snapshot & out) { - if (!snapshot_backend || !cache.ctx || !cache.buf || !cache.hc_state || - cache.layers.size() != (size_t)cache.n_layer) { - return false; - } - - free_deepseek4_snapshot(out); - - ggml_init_params ip{}; - ip.mem_size = ggml_tensor_overhead() * (size_t)(cache.n_layer * 8 + 8) + 4096; - ip.no_alloc = true; - out.ctx = ggml_init(ip); - if (!out.ctx) { - return false; - } - - out.layers.resize((size_t)cache.n_layer); - out.hc_state_snap = clone_snapshot_tensor(out.ctx, cache.hc_state, "ds4_hc_state_snap"); - if (!out.hc_state_snap) { - free_deepseek4_snapshot(out); - return false; - } - - for (int il = 0; il < cache.n_layer; ++il) { - const auto & src = cache.layers[(size_t)il]; - auto & dst = out.layers[(size_t)il]; - dst.raw_kv = clone_snapshot_tensor(out.ctx, src.raw_kv, nullptr); - dst.comp_kv = clone_snapshot_tensor(out.ctx, src.comp_kv, nullptr); - dst.index_comp_kv = clone_snapshot_tensor(out.ctx, src.index_comp_kv, nullptr); - dst.attn_compressor.state_kv = - clone_snapshot_tensor(out.ctx, src.attn_compressor.state_kv, nullptr); - dst.attn_compressor.state_score = - clone_snapshot_tensor(out.ctx, src.attn_compressor.state_score, nullptr); - dst.indexer_compressor.state_kv = - clone_snapshot_tensor(out.ctx, src.indexer_compressor.state_kv, nullptr); - dst.indexer_compressor.state_score = - clone_snapshot_tensor(out.ctx, src.indexer_compressor.state_score, nullptr); - if (!dst.raw_kv || - (src.comp_kv && !dst.comp_kv) || - (src.index_comp_kv && !dst.index_comp_kv) || - (src.attn_compressor.state_kv && !dst.attn_compressor.state_kv) || - (src.attn_compressor.state_score && !dst.attn_compressor.state_score) || - (src.indexer_compressor.state_kv && !dst.indexer_compressor.state_kv) || - (src.indexer_compressor.state_score && !dst.indexer_compressor.state_score)) { - free_deepseek4_snapshot(out); - return false; - } - } - - out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, snapshot_backend); - if (!out.buf) { - free_deepseek4_snapshot(out); - return false; - } - - if (!copy_tensor_from_backend(cache.hc_state, out.hc_state_snap)) { - free_deepseek4_snapshot(out); - return false; - } - for (int il = 0; il < cache.n_layer; ++il) { - const auto & src = cache.layers[(size_t)il]; - auto & dst = out.layers[(size_t)il]; - dst.n_comp = src.n_comp; - dst.n_index_comp = src.n_index_comp; - if (!copy_tensor_from_backend(src.raw_kv, dst.raw_kv) || - (src.comp_kv && !copy_tensor_from_backend(src.comp_kv, dst.comp_kv)) || - (src.index_comp_kv && - !copy_tensor_from_backend(src.index_comp_kv, dst.index_comp_kv)) || - (src.attn_compressor.state_kv && - !copy_tensor_from_backend(src.attn_compressor.state_kv, - dst.attn_compressor.state_kv)) || - (src.attn_compressor.state_score && - !copy_tensor_from_backend(src.attn_compressor.state_score, - dst.attn_compressor.state_score)) || - (src.indexer_compressor.state_kv && - !copy_tensor_from_backend(src.indexer_compressor.state_kv, - dst.indexer_compressor.state_kv)) || - (src.indexer_compressor.state_score && - !copy_tensor_from_backend(src.indexer_compressor.state_score, - dst.indexer_compressor.state_score))) { - free_deepseek4_snapshot(out); - return false; - } - } - - out.cur_pos = cache.cur_pos; - return true; -} - -bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, - DeepSeek4Cache & cache) { - if (!snap.ctx || !cache.ctx || !cache.buf || !snap.hc_state_snap || - snap.layers.size() != cache.layers.size()) { - return false; - } - if (!tensors_compatible(snap.hc_state_snap, cache.hc_state) || - !copy_tensor_to_backend(snap.hc_state_snap, cache.hc_state)) { - return false; - } - - for (size_t il = 0; il < cache.layers.size(); ++il) { - const auto & src = snap.layers[il]; - auto & dst = cache.layers[il]; - if (!tensors_compatible(src.raw_kv, dst.raw_kv) || - !tensors_compatible(src.comp_kv, dst.comp_kv) || - !tensors_compatible(src.index_comp_kv, dst.index_comp_kv) || - !tensors_compatible(src.attn_compressor.state_kv, dst.attn_compressor.state_kv) || - !tensors_compatible(src.attn_compressor.state_score, dst.attn_compressor.state_score) || - !tensors_compatible(src.indexer_compressor.state_kv, dst.indexer_compressor.state_kv) || - !tensors_compatible(src.indexer_compressor.state_score, dst.indexer_compressor.state_score)) { - return false; - } - if (!copy_tensor_to_backend(src.raw_kv, dst.raw_kv) || - (src.comp_kv && !copy_tensor_to_backend(src.comp_kv, dst.comp_kv)) || - (src.index_comp_kv && - !copy_tensor_to_backend(src.index_comp_kv, dst.index_comp_kv)) || - (src.attn_compressor.state_kv && - !copy_tensor_to_backend(src.attn_compressor.state_kv, - dst.attn_compressor.state_kv)) || - (src.attn_compressor.state_score && - !copy_tensor_to_backend(src.attn_compressor.state_score, - dst.attn_compressor.state_score)) || - (src.indexer_compressor.state_kv && - !copy_tensor_to_backend(src.indexer_compressor.state_kv, - dst.indexer_compressor.state_kv)) || - (src.indexer_compressor.state_score && - !copy_tensor_to_backend(src.indexer_compressor.state_score, - dst.indexer_compressor.state_score))) { - return false; - } - dst.n_comp = src.n_comp; - dst.n_index_comp = src.n_index_comp; - } - - cache.cur_pos = snap.cur_pos; - return true; -} - void free_deepseek4_cache(DeepSeek4Cache & c) { - if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } + if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } c.layers.clear(); c.hc_state = nullptr; } void free_deepseek4_snapshot(DeepSeek4Snapshot & s) { - if (s.buf) { ggml_backend_buffer_free(s.buf); s.buf = nullptr; } if (s.ctx) { ggml_free(s.ctx); s.ctx = nullptr; } + if (s.buf) { ggml_backend_buffer_free(s.buf); s.buf = nullptr; } s.layers.clear(); s.cur_pos = 0; s.hc_state_snap = nullptr; diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index d2a12b7d3..0b248b2bd 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -23,6 +23,7 @@ #include "ggml-backend.h" #include "internal.h" +#include "common/expert_split_compute_runtime.h" #include "common/layer_split_utils.h" namespace dflash::common { @@ -30,7 +31,9 @@ namespace dflash::common { struct MoeHybridPlacement; struct MoeHybridConfig; struct MoeHybridRoutingStats; -class MoeHybridStreamEngine; +struct MoeExpertLayer; +struct MoeExpertCompute; +class ExpertIpcClient; struct DeepSeek4StepTelemetry { uint64_t total_us = 0; @@ -61,10 +64,34 @@ struct DeepSeek4StepTelemetry { uint64_t ffn_cold_us = 0; uint64_t ffn_combine_us = 0; uint64_t ffn_partition_us = 0; + uint64_t ffn_cache_promote_us = 0; uint64_t ffn_hot_graph_builds = 0; uint64_t ffn_hot_graph_hits = 0; uint64_t ffn_cold_graph_builds = 0; uint64_t ffn_cold_graph_hits = 0; + uint64_t worker_us = 0; + uint64_t worker_parent_write_us = 0; + uint64_t worker_parent_wait_us = 0; + uint64_t worker_parent_read_us = 0; + uint64_t worker_request_read_us = 0; + uint64_t worker_partition_us = 0; + uint64_t worker_resident_eval_us = 0; + uint64_t worker_miss_build_us = 0; + uint64_t worker_miss_eval_us = 0; + uint64_t worker_request_bytes = 0; + uint64_t worker_response_bytes = 0; + uint64_t worker_hot_graph_builds = 0; + uint64_t worker_hot_graph_hits = 0; + uint64_t worker_cold_graph_builds = 0; + uint64_t worker_cold_graph_hits = 0; + uint64_t worker_hot_graph_build_us = 0; + uint64_t worker_hot_input_us = 0; + uint64_t worker_hot_compute_us = 0; + uint64_t worker_hot_read_us = 0; + uint64_t worker_cold_graph_build_us = 0; + uint64_t worker_cold_input_us = 0; + uint64_t worker_cold_compute_us = 0; + uint64_t worker_cold_read_us = 0; uint64_t hc_post_ffn_us = 0; uint64_t output_us = 0; uint64_t sample_us = 0; @@ -219,7 +246,7 @@ struct DeepSeek4Weights { int32_t eos_id = -1; int32_t eos_chat_id = -1; - // MoE hybrid placement (deprecated — layer split replaces expert split) + // MoE hybrid/expert-split placement state. bool moe_hybrid = false; }; @@ -252,6 +279,10 @@ struct DeepSeek4LayerCache { // Compressor rolling state DeepSeek4CompressorState attn_compressor; DeepSeek4CompressorState indexer_compressor; + + // Optional routing bias cached on host for CPU-side top-k selection. + std::vector route_bias_host; + bool route_bias_loaded = false; }; struct DeepSeek4Cache { @@ -280,6 +311,102 @@ struct DeepSeek4BackendConfig { int max_ctx = 0; // 0 = auto from SWA + compression capacity }; +struct DeepSeek4ExpertSplitGpuMemoryInfo { + uint64_t free_bytes = 0; + uint64_t total_bytes = 0; +}; + +struct DeepSeek4ExpertSplitBudgetMemoryInfo { + uint64_t total_bytes = 0; + uint64_t free_bytes = 0; + uint64_t core_bytes = 0; + uint64_t kv_bytes = 0; + uint64_t warm_bytes = 0; + uint64_t safety_bytes = 0; + uint64_t total_expert_bytes = 0; + bool parent_is_igpu = false; + uint64_t host_available_bytes = 0; + uint64_t host_budget_cap_bytes = 0; +}; + +struct DeepSeek4ExpertSplitHotCachePlan { + uint64_t hot_bytes = 0; + int cache_slots = 0; +}; + +int deepseek4_expert_split_prefill_chunk_limit_from_memory( + int requested_chunk, + bool expert_split_enabled, + bool parent_is_gpu, + const DeepSeek4ExpertSplitGpuMemoryInfo & memory); + +bool deepseek4_expert_split_should_disable_cached_decode_from_memory( + bool expert_split_enabled, + bool parent_is_gpu, + bool parent_is_cuda, + bool parent_is_hip, + const DeepSeek4ExpertSplitGpuMemoryInfo & memory); + +bool deepseek4_expert_split_requires_parent_cuda_graph_disable( + bool expert_split_enabled, + bool parent_is_cuda, + const std::vector & targets); + +bool deepseek4_should_use_gpu_resident_decode_ffn_policy( + bool use_cached_decode, + bool env_enabled, + bool env_disabled, + bool backend_is_hip, + bool backend_is_integrated_gpu, + bool worker_active); + +uint64_t deepseek4_expert_split_budget_from_memory( + const DeepSeek4ExpertSplitBudgetMemoryInfo & memory); + +uint64_t deepseek4_expert_split_primary_capacity_for_targets( + uint64_t expert_budget_bytes, + uint64_t hot_budget_bytes, + size_t configured_targets); + +uint64_t deepseek4_expert_split_non_cpu_capacity_bytes( + const std::vector & targets); + +uint64_t deepseek4_expert_split_effective_budget_for_targets( + uint64_t requested_budget_bytes, + uint64_t total_expert_bytes, + const std::vector & targets); + +uint64_t deepseek4_expert_split_effective_hot_budget( + uint64_t expert_budget_bytes, + uint64_t hot_budget_bytes, + size_t configured_targets); + +bool deepseek4_expert_split_hash_ids_to_hotness_counts( + const int32_t * expert_ids, + int n_token_ids, + int n_expert_used, + int n_expert, + std::vector & out_counts); + +bool deepseek4_expert_split_route_bias_to_hotness_counts( + const float * route_bias, + int n_expert, + uint64_t layer_total, + std::vector & out_counts); + +int deepseek4_expert_split_effective_cache_slots( + int requested_cache_slots, + size_t configured_targets); + +DeepSeek4ExpertSplitHotCachePlan deepseek4_expert_split_hot_cache_plan( + uint64_t expert_budget_bytes, + uint64_t total_expert_bytes, + int n_expert, + int n_expert_used, + bool parent_is_igpu, + uint64_t igpu_host_cap_bytes = 0, + uint64_t igpu_free_bytes = 0); + // ─── Function declarations ────────────────────────────────────────────── bool load_deepseek4_gguf(const std::string & path, @@ -308,7 +435,8 @@ bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, // Forward: single step (prefill chunk or decode token). // embed: [n_embd, n_tokens] input embeddings (post-embedding lookup). // hc_state: [n_hc * n_embd] persistent HC residual (updated in-place). -// Returns logits for last token. +// When want_logits is false, the step updates KV / routing state but skips the +// final output head and logits readback. bool deepseek4_step( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -319,9 +447,34 @@ bool deepseek4_step( std::vector & out_logits, MoeHybridStorage * moe_hybrid = nullptr, const int32_t * token_ids = nullptr, - MoeHybridStreamEngine * stream_engine = nullptr, + ExpertIpcClient * expert_worker = nullptr, + bool worker_owns_hot_ids = false, + bool want_logits = true, + bool disable_cached_decode = false, DeepSeek4StepTelemetry * telemetry = nullptr, - MoeHybridRoutingStats * routing_stats = nullptr); + MoeHybridRoutingStats * routing_stats = nullptr, + MoeExpertCompute * expert_compute = nullptr, + const MoeExpertLayer * expert_layers = nullptr); + +// Compatibility overload for builds that still route call sites through the +// original DeepSeek4 step signature without an explicit cached-decode toggle. +bool deepseek4_step( + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const float * embed, + int n_tokens, + int kv_start, + std::vector & out_logits, + MoeHybridStorage * moe_hybrid, + const int32_t * token_ids, + ExpertIpcClient * expert_worker, + bool worker_owns_hot_ids, + bool want_logits, + DeepSeek4StepTelemetry * telemetry, + MoeHybridRoutingStats * routing_stats = nullptr, + MoeExpertCompute * expert_compute = nullptr, + const MoeExpertLayer * expert_layers = nullptr); bool deepseek4_step_layer_range( ggml_backend_t backend, @@ -344,7 +497,10 @@ bool build_deepseek4_moe_hybrid_storage_from_file( const MoeHybridPlacement & placement, const MoeHybridConfig * cfg_override, MoeHybridStorage & out, - std::string * err = nullptr); + std::string * err = nullptr, + int cache_slots = 0, + bool load_cold_tensors = true, + const std::vector> * cold_expert_order_by_layer = nullptr); bool build_deepseek4_moe_hybrid_storage_from_file( const std::string & path, @@ -354,15 +510,6 @@ bool build_deepseek4_moe_hybrid_storage_from_file( MoeHybridStorage & out, std::string * err = nullptr); -bool build_deepseek4_moe_hybrid_storage_from_file_with_mmap( - const std::string & path, - ggml_backend_t backend, - const DeepSeek4Weights & w, - const MoeHybridPlacement & placement, - const MoeHybridConfig * cfg_override, - MoeHybridStorage & out, - std::string * err = nullptr); - // Snapshot struct DeepSeek4Snapshot { int cur_pos = 0; From 6cdbcfb7cff5c27c4ec0758eb3f1491066752cd6 Mon Sep 17 00:00:00 2001 From: weicj Date: Thu, 9 Jul 2026 13:52:42 +0800 Subject: [PATCH 3/7] fix(dflash): keep DeepSeek4 CMake sources clean --- server/CMakeLists.txt | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index d952461f8..f56977292 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -258,7 +258,6 @@ add_library(dflash_common STATIC src/deepseek4/deepseek4_daemon.cpp src/deepseek4/deepseek4_layer_split_adapter.cpp src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp - src/deepseek4/deepseek4_expert_ipc_daemon.cpp src/flashprefill_q8.cpp src/kv_cache.cpp src/kv_quant.cpp @@ -276,7 +275,6 @@ add_library(dflash_common STATIC src/common/target_shard_ipc_daemon.cpp src/common/dflash_feature_ring.cpp src/common/dflash_capture.cpp - src/common/expert_ipc.cpp src/common/dflash_draft_ipc.cpp src/common/dflash_draft_ipc_daemon.cpp src/common/pflash_drafter_ipc.cpp @@ -284,13 +282,6 @@ add_library(dflash_common STATIC src/common/dflash_spec_decode.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp - src/common/expert_split_plan.cpp - src/common/expert_split_state.cpp - src/common/expert_split_runtime.cpp - src/common/expert_split_compute_runtime.cpp - src/common/expert_split_materialization.cpp - src/common/expert_split_target_config.cpp - src/common/gguf_tensor_data.cpp src/qwen35/graph_builders.cpp src/qwen35moe/qwen35moe_ffn.cpp src/qwen35moe/qwen35moe_backend.cpp @@ -305,7 +296,6 @@ add_library(dflash_common STATIC src/common/moe_hybrid_stream.cpp src/common/moe_expert_compute.cpp src/common/moe_expert_compute_cpu.cpp - src/common/moe_expert_compute_local_gpu.cpp src/common/moe_expert_compute_ipc.cpp src/common/moe_hybrid_swap_manager.cpp src/common/moe_routing_collector.cpp @@ -324,7 +314,6 @@ add_library(dflash_common STATIC src/common/daemon_loop.cpp src/common/gguf_inspect.cpp src/common/backend_factory.cpp - src/common/host_memory_mode.cpp src/placement/placement_config.cpp src/common/layer_split_utils.cpp src/common/ddtree.cpp From 948004dd499cd34163adf4363fd77e7397882022 Mon Sep 17 00:00:00 2001 From: weicj Date: Thu, 9 Jul 2026 14:53:11 +0800 Subject: [PATCH 4/7] fix(dflash): keep DeepSeek4 PR489 compatible with PR502 --- server/CMakeLists.txt | 1 + server/src/deepseek4/deepseek4_backend.cpp | 1261 +------------------- server/src/deepseek4/deepseek4_backend.h | 22 - server/src/deepseek4/deepseek4_graph.cpp | 49 +- server/src/deepseek4/deepseek4_hc_cuda.cu | 339 ++++-- server/src/deepseek4/deepseek4_hc_cuda.h | 28 + server/src/deepseek4/deepseek4_internal.h | 97 -- server/src/deepseek4/deepseek4_loader.cpp | 8 +- server/src/deepseek4/expert_ipc.h | 76 ++ 9 files changed, 414 insertions(+), 1467 deletions(-) create mode 100644 server/src/deepseek4/expert_ipc.h diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index f56977292..5f44fdeb1 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -279,6 +279,7 @@ add_library(dflash_common STATIC src/common/dflash_draft_ipc_daemon.cpp src/common/pflash_drafter_ipc.cpp src/common/dflash_draft_graph.cpp + src/common/dspark_head.cpp src/common/dflash_spec_decode.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index e6673fad9..8047429f1 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3,7 +3,6 @@ #include "deepseek4_backend.h" #include "deepseek4_internal.h" #include "common/sampler.h" -#include "../common/expert_split_target_config.h" #include "../common/moe_hybrid_types.h" #include "ggml.h" @@ -40,573 +39,6 @@ static double gib(uint64_t bytes) { return (double) bytes / 1024.0 / 1024.0 / 1024.0; } -static int env_int(const char * name, int fallback) { - const char * value = std::getenv(name); - if (!value || !value[0]) return fallback; - return std::atoi(value); -} - -static double env_double(const char * name, double fallback) { - const char * value = std::getenv(name); - if (!value || !value[0]) return fallback; - return std::atof(value); -} - -static const char * env_str(const char * name, const char * fallback) { - const char * value = std::getenv(name); - return (value && value[0]) ? value : fallback; -} - -static int env_pct(const char * name, int fallback) { - const char * value = std::getenv(name); - if (!value || !value[0]) return fallback; - char * end = nullptr; - const long parsed = std::strtol(value, &end, 10); - if (end == value || *end != '\0' || parsed <= 0 || parsed > 100) { - return fallback; - } - return (int) parsed; -} - -static uint64_t apply_pct_bytes(uint64_t total_bytes, int pct) { - if (pct >= 100) return total_bytes; - return (total_bytes / 100ULL) * (uint64_t) pct + - ((total_bytes % 100ULL) * (uint64_t) pct) / 100ULL; -} - -static bool backend_is_cuda(ggml_backend_t backend); -static bool backend_is_hip(ggml_backend_t backend); -static bool backend_is_gpu(ggml_backend_t backend); - -} // namespace - -int deepseek4_expert_split_prefill_chunk_limit_from_memory( - int requested_chunk, - bool expert_split_enabled, - bool parent_is_gpu, - const DeepSeek4ExpertSplitGpuMemoryInfo & memory) { - int chunk = std::max(1, requested_chunk); - if (!expert_split_enabled || !parent_is_gpu) { - return chunk; - } - if (memory.total_bytes == 0) { - return std::min(chunk, 64); - } - - const uint64_t gib1 = 1024ULL * 1024ULL * 1024ULL; - const uint64_t gib24 = 24ULL * gib1; - const uint64_t gib48 = 48ULL * gib1; - const uint64_t gib80 = 80ULL * gib1; - const uint64_t total = memory.total_bytes; - const uint64_t free = memory.free_bytes; - if (total <= gib24 + 512ULL * 1024ULL * 1024ULL) { - if (free >= 12ULL * gib1) { - return std::min(chunk, 256); - } - if (free >= 6ULL * gib1) { - return std::min(chunk, 128); - } - return std::min(chunk, 64); - } - if (total <= gib48 + 512ULL * 1024ULL * 1024ULL) { - if (free >= 12ULL * gib1) { - return std::min(chunk, 256); - } - return std::min(chunk, 128); - } - if (total >= gib80 && free >= 24ULL * gib1) { - return std::min(chunk, 512); - } - return std::min(chunk, 256); -} - -bool deepseek4_expert_split_should_disable_cached_decode_from_memory( - bool expert_split_enabled, - bool parent_is_gpu, - bool parent_is_cuda, - bool parent_is_hip, - const DeepSeek4ExpertSplitGpuMemoryInfo & memory) { - if (!expert_split_enabled || !parent_is_gpu) { - return false; - } - if (parent_is_hip) { - return false; - } - if (!parent_is_cuda) { - return true; - } - if (memory.total_bytes == 0) { - return true; - } - - const uint64_t gib1 = 1024ULL * 1024ULL * 1024ULL; - const uint64_t gib24 = 24ULL * gib1; - if (memory.total_bytes > gib24 + 512ULL * 1024ULL * 1024ULL) { - return false; - } - return memory.free_bytes < 3ULL * gib1; -} - -bool deepseek4_expert_split_requires_parent_cuda_graph_disable( - bool expert_split_enabled, - bool parent_is_cuda, - const std::vector & targets) { - if (!expert_split_enabled || !parent_is_cuda || targets.size() <= 1) { - return false; - } - for (size_t i = 1; i < targets.size(); ++i) { - const ExpertSplitComputeTargetRuntime & target = targets[i]; - if (target.placement.total_hot <= 0) { - continue; - } - if (target.target.backend.empty() || target.target.backend == "cpu") { - continue; - } - return true; - } - return false; -} - -uint64_t deepseek4_expert_split_budget_from_memory( - const DeepSeek4ExpertSplitBudgetMemoryInfo & memory) { - const uint64_t fixed_bytes = - memory.core_bytes + - memory.kv_bytes + - memory.warm_bytes + - memory.safety_bytes; - if (memory.total_bytes <= fixed_bytes) { - return 0; - } - - uint64_t budget_bytes = memory.total_bytes - fixed_bytes; - if (memory.parent_is_igpu && memory.host_budget_cap_bytes > 0) { - budget_bytes = std::min(budget_bytes, memory.host_budget_cap_bytes); - } - if (memory.total_expert_bytes > 0 && - budget_bytes > memory.total_expert_bytes) { - budget_bytes = memory.total_expert_bytes; - } - return budget_bytes; -} - -uint64_t deepseek4_expert_split_primary_capacity_for_targets( - uint64_t expert_budget_bytes, - uint64_t hot_budget_bytes, - size_t configured_targets) { - if (configured_targets <= 1 || hot_budget_bytes == 0) { - return expert_budget_bytes; - } - if (expert_budget_bytes == 0) { - return 0; - } - return std::min(expert_budget_bytes, hot_budget_bytes); -} - -uint64_t deepseek4_expert_split_non_cpu_capacity_bytes( - const std::vector & targets) { - uint64_t total_bytes = 0; - for (const ExpertSplitTarget & target : targets) { - if (target.backend == "cpu") { - continue; - } - if (target.unlimited) { - return std::numeric_limits::max(); - } - const uint64_t usable = target.usable_bytes(); - if (usable == 0) { - continue; - } - if (std::numeric_limits::max() - total_bytes < usable) { - return std::numeric_limits::max(); - } - total_bytes += usable; - } - return total_bytes; -} - -uint64_t deepseek4_expert_split_effective_budget_for_targets( - uint64_t requested_budget_bytes, - uint64_t total_expert_bytes, - const std::vector & targets) { - if (requested_budget_bytes == 0 || total_expert_bytes == 0 || targets.empty()) { - return requested_budget_bytes; - } - - const bool has_cpu_target = std::any_of( - targets.begin(), targets.end(), - [](const ExpertSplitTarget & target) { - return target.backend == "cpu"; - }); - if (!has_cpu_target) { - return requested_budget_bytes; - } - - const uint64_t non_cpu_capacity = - deepseek4_expert_split_non_cpu_capacity_bytes(targets); - if (non_cpu_capacity == 0) { - return requested_budget_bytes; - } - - const uint64_t promoted_budget = - std::min(total_expert_bytes, non_cpu_capacity); - return std::max(requested_budget_bytes, promoted_budget); -} - -uint64_t deepseek4_expert_split_effective_hot_budget( - uint64_t expert_budget_bytes, - uint64_t hot_budget_bytes, - size_t configured_targets) { - if (expert_budget_bytes == 0) { - return 0; - } - if (configured_targets > 1) { - return expert_budget_bytes; - } - if (hot_budget_bytes == 0) { - return expert_budget_bytes; - } - return std::min(expert_budget_bytes, hot_budget_bytes); -} - -bool deepseek4_expert_split_hash_ids_to_hotness_counts( - const int32_t * expert_ids, - int n_token_ids, - int n_expert_used, - int n_expert, - std::vector & out_counts) { - out_counts.clear(); - if (!expert_ids || - n_token_ids <= 0 || - n_expert_used <= 0 || - n_expert <= 0) { - return false; - } - - out_counts.assign((size_t) n_expert, 0); - const size_t total_ids = - (size_t) n_token_ids * (size_t) n_expert_used; - for (size_t i = 0; i < total_ids; ++i) { - const int32_t expert = expert_ids[i]; - if (expert < 0 || expert >= n_expert) { - out_counts.clear(); - return false; - } - out_counts[(size_t) expert] += 1; - } - return true; -} - -bool deepseek4_expert_split_route_bias_to_hotness_counts( - const float * route_bias, - int n_expert, - uint64_t layer_total, - std::vector & out_counts) { - out_counts.clear(); - if (!route_bias || n_expert <= 0 || layer_total == 0) { - return false; - } - - const uint64_t minimum_total = - std::max(layer_total, (uint64_t) n_expert); - float max_bias = -INFINITY; - bool saw_finite = false; - for (int expert = 0; expert < n_expert; ++expert) { - const float bias = route_bias[expert]; - if (!std::isfinite(bias)) { - continue; - } - max_bias = saw_finite ? std::max(max_bias, bias) : bias; - saw_finite = true; - } - if (!saw_finite) { - return false; - } - - out_counts.assign((size_t) n_expert, 1); - uint64_t remaining = minimum_total - (uint64_t) n_expert; - if (remaining == 0) { - return true; - } - - std::vector weights((size_t) n_expert, 0.0L); - long double weight_sum = 0.0L; - for (int expert = 0; expert < n_expert; ++expert) { - double delta = (double) route_bias[expert] - (double) max_bias; - if (!std::isfinite(delta)) { - delta = -16.0; - } - delta = std::max(-16.0, std::min(0.0, delta)); - const long double weight = std::exp(delta); - weights[(size_t) expert] = weight; - weight_sum += weight; - } - if (!(weight_sum > 0.0L)) { - out_counts.clear(); - return false; - } - - struct ShareRemainder { - int expert = 0; - long double remainder = 0.0L; - }; - std::vector remainders; - remainders.reserve((size_t) n_expert); - - uint64_t assigned = 0; - for (int expert = 0; expert < n_expert; ++expert) { - const long double raw = - ((long double) remaining * weights[(size_t) expert]) / weight_sum; - const uint64_t extra = (uint64_t) std::floor(raw); - out_counts[(size_t) expert] += extra; - assigned += extra; - remainders.push_back({expert, raw - (long double) extra}); - } - - std::stable_sort(remainders.begin(), remainders.end(), - [](const ShareRemainder & lhs, const ShareRemainder & rhs) { - if (lhs.remainder != rhs.remainder) { - return lhs.remainder > rhs.remainder; - } - return lhs.expert < rhs.expert; - }); - for (uint64_t i = assigned; i < remaining; ++i) { - out_counts[(size_t) remainders[(size_t) (i - assigned)].expert] += 1; - } - - return true; -} - -int deepseek4_expert_split_effective_cache_slots( - int requested_cache_slots, - size_t configured_targets) { - if (configured_targets > 1) { - return 0; - } - return std::max(0, requested_cache_slots); -} - -static uint64_t ceil_div_u64(uint64_t value, uint64_t denom) { - if (denom == 0) { - return 0; - } - return value / denom + ((value % denom) != 0 ? 1ULL : 0ULL); -} - -DeepSeek4ExpertSplitHotCachePlan deepseek4_expert_split_hot_cache_plan( - uint64_t expert_budget_bytes, - uint64_t total_expert_bytes, - int n_expert, - int n_expert_used, - bool parent_is_igpu, - uint64_t igpu_host_cap_bytes, - uint64_t igpu_free_bytes) { - DeepSeek4ExpertSplitHotCachePlan plan; - plan.hot_bytes = expert_budget_bytes; - if (expert_budget_bytes == 0 || total_expert_bytes == 0 || n_expert <= 0) { - return plan; - } - - const uint64_t bytes_per_slot = - total_expert_bytes / (uint64_t) n_expert; - if (bytes_per_slot == 0) { - return plan; - } - - const int min_hot_slots = - std::max(0, std::min(n_expert, n_expert_used)); - const uint64_t min_hot_bytes = - bytes_per_slot * (uint64_t) min_hot_slots; - if (expert_budget_bytes <= min_hot_bytes) { - return plan; - } - - const int max_cache_slots = (int) std::min( - (expert_budget_bytes - min_hot_bytes) / bytes_per_slot, - (uint64_t) std::max(0, n_expert - min_hot_slots)); - if (max_cache_slots <= 0) { - return plan; - } - - int requested_slots = 0; - const bool explicit_slots = - std::getenv("DFLASH_DEEPSEEK4_CACHE_SLOTS") != nullptr; - if (const char * raw = std::getenv("DFLASH_DEEPSEEK4_CACHE_SLOTS")) { - requested_slots = std::max(0, std::atoi(raw)); - } else { - auto spark = spark_budget_split( - expert_budget_bytes, - total_expert_bytes, - n_expert, - /*core_kv_safety=*/0, - /*target_bytes=*/0); - requested_slots = std::max(0, spark.cache_slots); - - if (parent_is_igpu) { - const int reserve_pct = env_pct( - "DFLASH_DEEPSEEK4_CACHE_RESERVE_PCT", 20); - uint64_t reserve_bytes = - (expert_budget_bytes / 100ULL) * (uint64_t) reserve_pct + - ((expert_budget_bytes % 100ULL) * - (uint64_t) reserve_pct) / 100ULL; - if (const char * raw = - std::getenv("DFLASH_DEEPSEEK4_CACHE_RESERVE_MB")) { - const long long reserve_mb = std::atoll(raw); - if (reserve_mb > 0) { - reserve_bytes = - (uint64_t) reserve_mb * 1024ULL * 1024ULL; - } - } else { - const int reserve_cap_mb = env_int( - "DFLASH_DEEPSEEK4_CACHE_RESERVE_CAP_MB", 16384); - if (reserve_cap_mb > 0) { - const uint64_t reserve_cap_bytes = - (uint64_t) reserve_cap_mb * 1024ULL * 1024ULL; - reserve_bytes = - std::min(reserve_bytes, reserve_cap_bytes); - } - } - const int min_cache_slots = std::min( - n_expert, - std::max(n_expert_used * 2, 12)); - if (reserve_bytes > 0) { - requested_slots = std::max( - requested_slots, - (int) (reserve_bytes / bytes_per_slot)); - } - requested_slots = std::max(requested_slots, min_cache_slots); - } - } - - if (parent_is_igpu && !explicit_slots) { - uint64_t pinned_hot_cap = expert_budget_bytes; - if (igpu_host_cap_bytes > 0) { - pinned_hot_cap = std::min(pinned_hot_cap, igpu_host_cap_bytes); - } - if (igpu_free_bytes > 0) { - pinned_hot_cap = std::min(pinned_hot_cap, igpu_free_bytes); - } - pinned_hot_cap = std::min( - pinned_hot_cap, - apply_pct_bytes( - expert_budget_bytes, - env_pct("DFLASH_DEEPSEEK4_IGPU_BOOTSTRAP_CAP_PCT", 80))); - pinned_hot_cap = std::max(pinned_hot_cap, min_hot_bytes); - if (expert_budget_bytes > pinned_hot_cap) { - const uint64_t extra_slots = ceil_div_u64( - expert_budget_bytes - pinned_hot_cap, bytes_per_slot); - requested_slots = std::max( - requested_slots, - (int) std::min(extra_slots, (uint64_t) max_cache_slots)); - } - } - - plan.cache_slots = std::clamp(requested_slots, 0, max_cache_slots); - plan.hot_bytes = - expert_budget_bytes - - (uint64_t) plan.cache_slots * bytes_per_slot; - return plan; -} - -namespace { - -static int resolve_ds4_expert_split_prefill_chunk( - int requested_chunk, - bool expert_split_enabled, - PlacementBackend device_backend, - ggml_backend_t backend, - int gpu_index) { - int chunk = std::max(1, requested_chunk); - if (!expert_split_enabled || !backend_is_gpu(backend)) { - return chunk; - } - - const int env_cap = env_int("DFLASH_EXPERT_SPLIT_PREFILL_CHUNK_MAX", 0); - if (env_cap > 0) { - return std::min(chunk, env_cap); - } - - uint64_t gpu_free = 0; - uint64_t gpu_total = 0; - if (!query_expert_split_backend_memory( - device_backend, gpu_index, gpu_free, gpu_total)) { - return std::min(chunk, 64); - } - return dflash::common::deepseek4_expert_split_prefill_chunk_limit_from_memory( - chunk, - expert_split_enabled, - /*parent_is_gpu=*/true, - DeepSeek4ExpertSplitGpuMemoryInfo{gpu_free, gpu_total}); -} - -static bool should_disable_ds4_expert_split_cached_decode( - bool expert_split_enabled, - PlacementBackend device_backend, - ggml_backend_t backend, - int gpu_index) { - if (!expert_split_enabled || !backend_is_gpu(backend)) { - return false; - } - if (env_flag_enabled("DFLASH_DS4_FORCE_MULTI_TARGET_DECODE_CACHE")) { - return false; - } - if (env_flag_enabled("DFLASH_DS4_DISABLE_MULTI_TARGET_DECODE_CACHE")) { - return true; - } - if (backend_is_hip(backend)) { - return false; - } - if (!backend_is_cuda(backend)) { - return true; - } - - uint64_t gpu_free = 0; - uint64_t gpu_total = 0; - if (!query_expert_split_backend_memory( - device_backend, gpu_index, gpu_free, gpu_total)) { - return true; - } - return dflash::common::deepseek4_expert_split_should_disable_cached_decode_from_memory( - expert_split_enabled, - /*parent_is_gpu=*/true, - /*parent_is_cuda=*/true, - /*parent_is_hip=*/false, - DeepSeek4ExpertSplitGpuMemoryInfo{gpu_free, gpu_total}); -} - -static PlacementBackend resolved_device_backend(const DevicePlacement & device) { - return device.backend == PlacementBackend::Auto - ? compiled_placement_backend() - : device.backend; -} - -static bool backend_is_cuda(ggml_backend_t backend) { - const char * name = ggml_backend_name(backend); - return name && std::strstr(name, "CUDA") != nullptr; -} - -static bool backend_is_hip(ggml_backend_t backend) { - const char * name = ggml_backend_name(backend); - return name && - (std::strstr(name, "HIP") != nullptr || - std::strstr(name, "ROCm") != nullptr); -} - -static bool backend_is_gpu(ggml_backend_t backend) { - return backend_is_cuda(backend) || backend_is_hip(backend); -} - -static bool resolve_ds4_expert_split_targets_from_env( - uint64_t primary_capacity_bytes, - std::vector & out, - std::string * err) { - return dflash::common::resolve_expert_split_targets_from_env( - "DFLASH_DEEPSEEK4_EXPERT_TARGETS", - "DFLASH_DEEPSEEK4_EXPERT_TARGET_CAPS", - primary_capacity_bytes, - out, - err); -} - static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTelemetry & src) { dst.total_us += src.total_us; dst.embed_us += src.embed_us; @@ -754,17 +186,11 @@ struct Ds4HybridBudgetInfo { Ds4ExpertMemoryInfo mem; size_t gpu_free = 0; size_t gpu_total = 0; - bool parent_is_igpu = false; - uint64_t host_available_bytes = 0; - uint64_t host_budget_cap_bytes = 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 auto_expert_budget = 0; uint64_t expert_budget = 0; - uint64_t hot_budget = 0; - int cache_slots = 0; int max_hot_per_layer = 0; }; @@ -862,23 +288,12 @@ static void fill_prefix_hot_placement(const DeepSeek4Weights & w, } static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, - PlacementBackend device_backend, int gpu, int max_ctx, Ds4HybridBudgetInfo & out, std::string * err) { out = {}; - ExpertSplitBackendMemoryInfo memory_info; - if (!query_expert_split_backend_memory_info( - device_backend, gpu, memory_info) || - memory_info.total_bytes == 0) { - if (err) *err = "could not query GPU memory"; - return false; - } - out.gpu_free = (size_t) memory_info.free_bytes; - out.gpu_total = (size_t) memory_info.total_bytes; - out.parent_is_igpu = - memory_info.device_type == GGML_BACKEND_DEVICE_TYPE_IGPU; + ggml_backend_cuda_get_device_memory(gpu, &out.gpu_free, &out.gpu_total); if (!compute_ds4_expert_memory_info(w, nullptr, out.mem, err)) { return false; @@ -887,27 +302,13 @@ static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, out.core_bytes = moe_hybrid_core_bytes_from_memory("deepseek4", out.gpu_free, out.gpu_total); out.kv_bytes = estimate_ds4_cache_bytes(w, max_ctx); - ExpertSplitIgpuHostBudgetInfo host_budget; - if (out.parent_is_igpu) { - (void) query_expert_split_igpu_host_budget_info(host_budget); - out.host_available_bytes = host_budget.available_bytes; - out.host_budget_cap_bytes = host_budget.cap_bytes; + if (out.gpu_total <= out.core_bytes + out.kv_bytes + out.warm_bytes + out.safety_bytes) { + if (err) *err = "no VRAM budget available for DS4 experts"; + return false; } - - out.auto_expert_budget = deepseek4_expert_split_budget_from_memory( - DeepSeek4ExpertSplitBudgetMemoryInfo{ - (uint64_t) out.gpu_total, - (uint64_t) out.gpu_free, - out.core_bytes, - out.kv_bytes, - out.warm_bytes, - out.safety_bytes, - out.mem.total_expert_bytes, - out.parent_is_igpu, - out.host_available_bytes, - out.host_budget_cap_bytes, - }); - out.expert_budget = out.auto_expert_budget; + out.expert_budget = + out.gpu_total - out.core_bytes - out.kv_bytes - out.warm_bytes - out.safety_bytes; + out.expert_budget = std::min(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) { @@ -919,21 +320,11 @@ static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, return false; } - const auto hot_cache_plan = deepseek4_expert_split_hot_cache_plan( - out.expert_budget, - out.mem.total_expert_bytes, - w.n_expert, - w.n_expert_used, - out.parent_is_igpu, - out.host_budget_cap_bytes, - (uint64_t) out.gpu_free); - out.hot_budget = hot_cache_plan.hot_bytes; - out.cache_slots = hot_cache_plan.cache_slots; out.max_hot_per_layer = std::min( w.n_expert, - (int) (out.hot_budget / out.mem.bytes_per_uniform_round)); + (int) (out.expert_budget / out.mem.bytes_per_uniform_round)); if (out.max_hot_per_layer <= 0) { - if (err) *err = "pinned hot budget is smaller than one uniform expert round"; + if (err) *err = "expert budget is smaller than one uniform expert round"; return false; } return true; @@ -955,22 +346,6 @@ static MoeHybridConfig make_ds4_parent_worker_cfg(const DeepSeek4Weights & w) { 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; -} - -static MoeHybridConfig make_ds4_parent_gpu_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::Gpu; - return cfg; -} - static MoeLayerDesc make_ds4_moe_layer_desc(const DeepSeek4Layer & L) { MoeLayerDesc desc; desc.ffn_gate_exps = L.ffn_gate_exps; @@ -984,168 +359,6 @@ static MoeLayerDesc make_ds4_moe_layer_desc(const DeepSeek4Layer & L) { return desc; } -static bool deepseek4_target_list_prefers_static_hotness_prior( - const std::vector & targets) { - if (targets.size() > 1) { - return true; - } - return std::any_of( - targets.begin(), targets.end(), - [](const ExpertSplitTarget & target) { - return target.backend == "cpu"; - }); -} - -static bool init_uniform_hotness( - int n_layer, - int n_expert, - int n_expert_used, - MoeHybridRoutingStats & hotness, - std::string * err) { - if (!hotness.init(n_layer, n_expert, n_expert_used)) { - if (err) *err = "failed to initialize uniform hotness"; - return false; - } - std::fill(hotness.counts.begin(), hotness.counts.end(), 1); - hotness.layer_totals.assign((size_t) n_layer, (uint64_t) n_expert); - return true; -} - -static bool build_deepseek4_static_hotness_prior( - const DeepSeek4Weights & w, - MoeHybridRoutingStats & hotness, - std::string & source) { - if (!hotness.init(w.n_layer, w.n_expert, w.n_expert_used)) { - return false; - } - - constexpr uint64_t kRouteBiasLayerTotal = 1ULL << 20; - bool used_hash_prior = false; - bool used_bias_prior = false; - std::vector layer_counts; - std::vector hash_ids; - std::vector route_bias; - - for (int layer = 0; layer < w.n_layer; ++layer) { - const DeepSeek4Layer & L = w.layers[(size_t) layer]; - bool seeded = false; - - if (layer < w.n_hash_layer && - L.ffn_gate_tid2eid && - L.ffn_gate_tid2eid->ne[0] == w.n_expert_used && - L.ffn_gate_tid2eid->ne[1] > 0) { - const int n_token_ids = (int) L.ffn_gate_tid2eid->ne[1]; - hash_ids.resize((size_t) n_token_ids * (size_t) w.n_expert_used); - ggml_backend_tensor_get( - L.ffn_gate_tid2eid, - hash_ids.data(), - 0, - sizeof(int32_t) * hash_ids.size()); - if (deepseek4_expert_split_hash_ids_to_hotness_counts( - hash_ids.data(), - n_token_ids, - w.n_expert_used, - w.n_expert, - layer_counts)) { - used_hash_prior = true; - seeded = true; - } - } - - if (!seeded && - L.ffn_exp_probs_b && - L.ffn_exp_probs_b->ne[0] == w.n_expert) { - route_bias.resize((size_t) w.n_expert); - ggml_backend_tensor_get( - L.ffn_exp_probs_b, - route_bias.data(), - 0, - sizeof(float) * route_bias.size()); - if (deepseek4_expert_split_route_bias_to_hotness_counts( - route_bias.data(), - w.n_expert, - kRouteBiasLayerTotal, - layer_counts)) { - used_bias_prior = true; - seeded = true; - } - } - - if (!seeded) { - layer_counts.assign((size_t) w.n_expert, 1); - } - - uint64_t layer_total = 0; - for (int expert = 0; expert < w.n_expert; ++expert) { - const uint64_t count = layer_counts[(size_t) expert]; - hotness.counts[(size_t) layer * (size_t) w.n_expert + (size_t) expert] = - count; - layer_total += count; - } - hotness.layer_totals[(size_t) layer] = layer_total; - } - - if (used_hash_prior && used_bias_prior) { - source = "static-prior(hash+bias)"; - } else if (used_hash_prior) { - source = "static-prior(hash)"; - } else if (used_bias_prior) { - source = "static-prior(bias)"; - } else { - source = "uniform"; - } - return used_hash_prior || used_bias_prior; -} - -static bool build_ds4_cold_owner_placement( - const MoeHybridPlacement & target_placement, - MoeHybridPlacement & out, - std::vector> & cold_order_by_layer, - std::string * err) { - if (!target_placement.matches(target_placement.n_layer, - target_placement.n_expert, - target_placement.n_expert_used)) { - if (err) *err = "invalid expert split target placement"; - return false; - } - out = target_placement; - out.total_hot = 0; - out.hot_counts.assign((size_t) target_placement.n_layer, 0); - out.hot_expert_ids.resize((size_t) target_placement.n_layer); - cold_order_by_layer.resize((size_t) target_placement.n_layer); - for (int il = 0; il < target_placement.n_layer; ++il) { - const auto & target_owned = - target_placement.hot_expert_ids[(size_t) il]; - std::vector owned((size_t) target_placement.n_expert, 0); - auto & cold_order = cold_order_by_layer[(size_t) il]; - cold_order.clear(); - cold_order.reserve(target_owned.size()); - for (int32_t expert : target_owned) { - if (expert < 0 || expert >= target_placement.n_expert) { - if (err) *err = "expert split target placement expert id out of range"; - return false; - } - if (owned[(size_t) expert]) { - if (err) *err = "expert split target placement contains duplicates"; - return false; - } - owned[(size_t) expert] = 1; - cold_order.push_back(expert); - } - auto & hot_experts = out.hot_expert_ids[(size_t) il]; - hot_experts.clear(); - hot_experts.reserve((size_t) target_placement.n_expert - target_owned.size()); - for (int expert = 0; expert < target_placement.n_expert; ++expert) { - if (!owned[(size_t) expert]) { - hot_experts.push_back(expert); - } - } - out.hot_counts[(size_t) il] = (int) hot_experts.size(); - out.total_hot += (int) hot_experts.size(); - } - return true; -} - } // namespace DeepSeek4Backend::DeepSeek4Backend(const DeepSeek4BackendConfig & cfg) @@ -1173,21 +386,14 @@ bool DeepSeek4Backend::run_step_with_runtime_path( } (void) disable_cached_decode; - MoeExpertCompute * expert_compute = - expert_split_runtime_.enabled ? expert_split_runtime_.compute_ptr() - : expert_runtime_.compute_ptr(); - const MoeExpertLayer * expert_layers = - expert_split_runtime_.enabled ? expert_split_runtime_.layer_ptr(0) - : expert_runtime_.layer_ptr(0); - return deepseek4_step(backend_, w_, cache_, embed, n_tokens, kv_start, out_logits, moe_hybrid_.get(), token_ids, nullptr, false, want_logits, telemetry, routing_stats_.get(), - expert_compute, - expert_layers); + expert_runtime_.compute_ptr(), + expert_runtime_.layer_ptr(0)); } void DeepSeek4Backend::reset_request_state() { @@ -1248,8 +454,7 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & std::string * err) const { Ds4HybridBudgetInfo budget; if (!compute_ds4_hybrid_budget_info( - w, resolved_device_backend(cfg_.device), - cfg_.device.gpu, max_ctx, budget, err)) { + w, cfg_.device.gpu, max_ctx, budget, err)) { return false; } @@ -1262,369 +467,19 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & } 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 host_cap=%.2f GiB expert_budget=%.2f GiB pinned_hot=%.2f GiB cache_slots=%d hot/layer=%d\n", + "[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.host_budget_cap_bytes), gib(budget.expert_budget), - gib(budget.hot_budget), - budget.cache_slots, hot_per_layer); log_ds4_expert_memory_info("placement", placed_mem, w.n_layer); return true; } -bool DeepSeek4Backend::build_expert_split_state_from_hotness( - const MoeHybridRoutingStats & hotness, - uint64_t expert_budget_bytes, - uint64_t hot_budget_bytes, - int max_hot_per_layer, - const std::vector & configured_targets, - const DeepSeek4Weights & w, - std::string * err) { - if ((int)layer_expert_bytes_.size() != w.n_layer) { - if (err) *err = "layer expert bytes not initialized"; - return false; - } - if (expert_budget_bytes == 0) { - if (err) *err = "expert budget must be > 0"; - return false; - } - if (!hotness.matches(w.n_layer, w.n_expert, w.n_expert_used)) { - if (err) *err = "hotness table dimensions do not match DeepSeek4 model"; - return false; - } - - ExpertSplitConfig cfg; - cfg.n_layer = w.n_layer; - cfg.n_expert = w.n_expert; - cfg.allow_implicit_cpu_fallback = true; - cfg.require_full_grid = true; - - std::vector targets = configured_targets; - if (targets.size() > 1 && !targets[0].unlimited) { - targets[0].capacity_bytes = - deepseek4_expert_split_primary_capacity_for_targets( - targets[0].capacity_bytes, - hot_budget_bytes, - targets.size()); - } - if (targets.empty()) { - const PlacementBackend backend = resolved_device_backend(cfg_.device); - targets = { - {std::string(placement_backend_name(backend)) + ":" + - std::to_string(cfg_.device.gpu), - placement_backend_name(backend), cfg_.device.gpu, - expert_budget_bytes, 0, false}, - }; - } - if (!build_capacity_weighted_expert_split_min_per_layer( - targets, std::min(w.n_expert_used, w.n_expert), - cfg.min_per_layer_by_target, err)) { - return false; - } - if (!targets.empty() && - std::none_of(cfg.min_per_layer_by_target.begin(), - cfg.min_per_layer_by_target.end(), - [](int v) { return v > 0; })) { - cfg.min_per_layer_by_target.assign(targets.size(), 0); - cfg.min_per_layer_by_target[0] = std::min(w.n_expert_used, w.n_expert); - } - - std::vector units; - units.reserve((size_t)w.n_layer * (size_t)w.n_expert); - for (int il = 0; il < w.n_layer; ++il) { - for (int ie = 0; ie < w.n_expert; ++ie) { - units.push_back(ExpertSplitUnit{ - il, - ie, - layer_expert_bytes_[(size_t)il], - (double)hotness.count(il, ie), - }); - } - } - - std::vector primary_hot_limits; - if (max_hot_per_layer > 0) { - primary_hot_limits.assign((size_t) w.n_layer, max_hot_per_layer); - } - - ExpertSplitStateComponents state; - if (!build_expert_split_state(cfg, targets, units, - w.n_expert_used, - primary_hot_limits.empty() - ? nullptr - : &primary_hot_limits, - state, err)) { - return false; - } - if (!validate_primary_expert_split_target(state.plan.targets, - resolved_device_backend(cfg_.device), - cfg_.device.gpu, err)) { - return false; - } - expert_split_state_ = std::move(state); - return true; -} - -bool DeepSeek4Backend::load_expert_split_placement( - const char * hotness_path, - const DeepSeek4Weights & w, - int max_ctx, - MoeHybridPlacement & out, - std::string * err) { - Ds4HybridBudgetInfo budget; - if (!compute_ds4_hybrid_budget_info( - w, resolved_device_backend(cfg_.device), - cfg_.device.gpu, max_ctx, budget, err)) { - return false; - } - layer_expert_bytes_ = budget.mem.layer_expert_bytes; - - std::vector configured_targets; - if (!resolve_ds4_expert_split_targets_from_env( - budget.expert_budget, configured_targets, err)) { - return false; - } - - uint64_t placement_budget_bytes = budget.expert_budget; - uint64_t placement_non_cpu_capacity_bytes = 0; - bool promoted_cpu_fallback_only = false; - if (!configured_targets.empty() && - budget.auto_expert_budget > budget.expert_budget) { - std::vector auto_capacity_targets; - if (!resolve_ds4_expert_split_targets_from_env( - budget.auto_expert_budget, auto_capacity_targets, err)) { - return false; - } - const uint64_t promoted_budget = - deepseek4_expert_split_effective_budget_for_targets( - budget.expert_budget, - budget.mem.total_expert_bytes, - auto_capacity_targets); - if (promoted_budget > placement_budget_bytes) { - placement_non_cpu_capacity_bytes = - deepseek4_expert_split_non_cpu_capacity_bytes( - auto_capacity_targets); - placement_budget_bytes = promoted_budget; - configured_targets = std::move(auto_capacity_targets); - promoted_cpu_fallback_only = true; - } - } - - MoeHybridRoutingStats hotness; - std::string hotness_source = "uniform"; - if (hotness_path && hotness_path[0]) { - if (!MoeHybridRoutingStats::load_csv(std::string(hotness_path), hotness, err)) { - return false; - } - if (!hotness.matches(w.n_layer, w.n_expert, w.n_expert_used)) { - if (err) *err = "hotness table dimensions do not match DeepSeek4 model"; - return false; - } - hotness_source = hotness_path; - } else if (deepseek4_target_list_prefers_static_hotness_prior( - configured_targets) && - build_deepseek4_static_hotness_prior( - w, hotness, hotness_source)) { - // hotness_source already describes the static prior that was used - } else if (!init_uniform_hotness( - w.n_layer, w.n_expert, w.n_expert_used, hotness, err)) { - return false; - } - - if (routing_stats_ && hotness_path && hotness_path[0] && - hotness.counts.size() == (size_t)w.n_layer * (size_t)w.n_expert) { - routing_stats_->counts = hotness.counts; - routing_stats_->layer_totals.assign((size_t)w.n_layer, 0); - for (int il = 0; il < w.n_layer; ++il) { - for (int ie = 0; ie < w.n_expert; ++ie) { - routing_stats_->layer_totals[(size_t)il] += - hotness.counts[(size_t)il * (size_t)w.n_expert + (size_t)ie]; - } - } - } - - const size_t configured_target_count = - configured_targets.empty() ? 1 : configured_targets.size(); - const uint64_t effective_hot_budget = - deepseek4_expert_split_effective_hot_budget( - placement_budget_bytes, - budget.hot_budget, - configured_target_count); - int effective_max_hot_per_layer = budget.max_hot_per_layer; - if (budget.mem.bytes_per_uniform_round > 0) { - effective_max_hot_per_layer = std::min( - w.n_expert, - (int) (effective_hot_budget / budget.mem.bytes_per_uniform_round)); - } - if (effective_max_hot_per_layer <= 0) { - if (err) { - *err = - "effective pinned hot budget is smaller than one uniform expert round"; - } - return false; - } - - expert_split_total_budget_bytes_ = budget.expert_budget; - expert_split_hot_budget_bytes_ = effective_hot_budget; - expert_split_cache_slots_ = - deepseek4_expert_split_effective_cache_slots( - budget.cache_slots, - configured_target_count); - - if (!build_expert_split_state_from_hotness( - hotness, placement_budget_bytes, effective_hot_budget, - effective_max_hot_per_layer, configured_targets, w, err)) { - return false; - } - expert_split_total_budget_bytes_ = placement_budget_bytes; - if (expert_split_cache_slots_ != budget.cache_slots) { - std::fprintf(stderr, - "[deepseek4] expert_split cache disabled for multi-target placement: requested_slots=%d targets=%zu\n", - budget.cache_slots, - configured_target_count); - } - if (promoted_cpu_fallback_only) { - std::fprintf(stderr, - "[deepseek4] expert_split promoted non-CPU placement budget: requested=%.2f GiB effective=%.2f GiB non_cpu_cap=%.2f GiB total_experts=%.2f GiB\n", - gib(budget.expert_budget), - gib(placement_budget_bytes), - gib(placement_non_cpu_capacity_bytes), - gib(budget.mem.total_expert_bytes)); - } - out = expert_split_state_.materialization.primary_placement; - - Ds4ExpertMemoryInfo placed_mem; - if (!compute_ds4_expert_memory_info(w, &out, placed_mem, err)) { - return false; - } - std::fprintf(stderr, - "[deepseek4] expert_split placement: gpu_total=%.2f GiB gpu_free=%.2f GiB core=%.2f GiB kv=%.2f GiB warm=%.2f GiB safety=%.2f GiB host_cap=%.2f GiB expert_budget=%.2f GiB placement_budget=%.2f GiB pinned_hot=%.2f GiB cache_slots=%d targets=%zu source=%s\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.host_budget_cap_bytes), - gib(budget.expert_budget), - gib(placement_budget_bytes), - gib(effective_hot_budget), - expert_split_cache_slots_, - configured_target_count, - hotness_source.c_str()); - log_ds4_expert_memory_info("expert_split", placed_mem, w.n_layer); - return true; -} - -bool DeepSeek4Backend::init_expert_split_runtime(std::string * err) { - expert_split_runtime_.reset(); - if (!moe_hybrid_ || expert_split_state_.empty()) { - return true; - } - if (!expert_split_state_.compute_runtime.matches( - w_.n_layer, w_.n_expert, w_.n_expert_used)) { - if (err) *err = "expert split compute runtime does not match DeepSeek4 model"; - return false; - } - - MoeExpertComputeRuntimeConfig runtime_cfg; - runtime_cfg.target_path = cfg_.model_path ? cfg_.model_path : ""; - runtime_cfg.n_layer = w_.n_layer; - runtime_cfg.n_expert = w_.n_expert; - runtime_cfg.n_expert_used = w_.n_expert_used; - runtime_cfg.n_embd = w_.n_embd; - runtime_cfg.n_ff_exp = w_.n_ff_exp; - runtime_cfg.enabled = true; - runtime_cfg.log_prefix = "[deepseek4-expert-split]"; - runtime_cfg.local_target_factory_key = "deepseek4-local-gpu-cold"; - runtime_cfg.local_target_compute_factory = - [this](const ExpertSplitComputeTargetRuntime & split_target, - std::string * local_err) { - return make_expert_split_local_target_compute(split_target, local_err); - }; - - std::vector layer_descs((size_t)w_.n_layer); - for (int il = 0; il < w_.n_layer; ++il) { - layer_descs[(size_t)il] = make_ds4_moe_layer_desc(w_.layers[(size_t)il]); - } - - if (!ensure_multi_target_moe_expert_compute_runtime( - expert_split_runtime_, runtime_cfg, - expert_split_state_.compute_runtime, - *moe_hybrid_, layer_descs, err)) { - return false; - } - if (expert_split_runtime_.enabled && expert_split_runtime_.compute_ptr()) { - std::fprintf(stderr, "[deepseek4] expert_split compute ready: %s\n", - expert_split_runtime_.runtime_key.c_str()); - } - return true; -} - -std::unique_ptr DeepSeek4Backend::make_expert_split_local_target_compute( - const ExpertSplitComputeTargetRuntime & split_target, - std::string * err) const { - if (!cfg_.model_path || !cfg_.model_path[0]) { - if (err) *err = "DeepSeek4 local target compute requires a model path"; - return nullptr; - } - if (split_target.target.backend == "cpu") { - return nullptr; - } - if (split_target.target.device_id < 0) { - if (err) *err = "DeepSeek4 local target compute requires an explicit device id"; - return nullptr; - } - MoeHybridPlacement cold_owner_placement; - std::vector> cold_order_by_layer; - if (!build_ds4_cold_owner_placement(split_target.placement, - cold_owner_placement, - cold_order_by_layer, - err)) { - return nullptr; - } - - ggml_backend_t target_backend = - ggml_backend_cuda_init(std::max(0, split_target.target.device_id)); - if (!target_backend) { - if (err) *err = "failed to initialize DeepSeek4 local backend"; - return nullptr; - } - - MoeHybridConfig local_cfg = make_ds4_parent_gpu_tail_cfg(w_); - auto storage = std::make_unique(); - if (!build_deepseek4_moe_hybrid_storage_from_file( - cfg_.model_path, - target_backend, - w_, - cold_owner_placement, - &local_cfg, - *storage, - err, - /*cache_slots=*/0, - /*load_cold_tensors=*/true, - &cold_order_by_layer)) { - ggml_backend_free(target_backend); - return nullptr; - } - - std::vector layer_descs((size_t) w_.n_layer); - for (int il = 0; il < w_.n_layer; ++il) { - layer_descs[(size_t) il] = make_ds4_moe_layer_desc(w_.layers[(size_t) il]); - } - return make_local_gpu_cold_moe_expert_compute( - target_backend, - std::move(storage), - local_cfg, - std::move(layer_descs)); -} - bool DeepSeek4Backend::init_single_target_expert_runtime(std::string * err) { expert_runtime_.reset(); if (!moe_hybrid_) { @@ -1658,10 +513,6 @@ bool DeepSeek4Backend::init_single_target_expert_runtime(std::string * err) { } bool DeepSeek4Backend::init_hybrid_model() { - expert_split_total_budget_bytes_ = 0; - expert_split_hot_budget_bytes_ = 0; - expert_split_cache_slots_ = 0; - TargetLoadPlan plan; plan.skip_expert_tensors = true; if (!load_deepseek4_gguf_partial(cfg_.model_path, backend_, plan, w_)) { @@ -1682,9 +533,8 @@ bool DeepSeek4Backend::init_hybrid_model() { std::string err; const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; - const char * hotness_path = std::getenv("DFLASH_DEEPSEEK4_HOTNESS"); - if (!load_expert_split_placement(hotness_path, w_, max_ctx, moe_placement_, &err)) { - std::fprintf(stderr, "[deepseek4] failed to compute expert_split placement: %s\n", err.c_str()); + 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()); return false; } @@ -1699,15 +549,10 @@ bool DeepSeek4Backend::init_hybrid_model() { } auto hybrid = std::make_shared(); - const std::vector> * cold_order_by_layer = - expert_split_state_.materialization.ordered_cold_union - ? &expert_split_state_.materialization.cold_expert_ids_by_layer - : nullptr; + const MoeHybridConfig hybrid_cfg = make_ds4_parent_worker_cfg(w_); if (!build_deepseek4_moe_hybrid_storage_from_file( - cfg_.model_path, backend_, w_, moe_placement_, nullptr, - *hybrid, &err, expert_split_cache_slots_, - /*load_cold_tensors=*/true, - cold_order_by_layer)) { + 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()); return false; } @@ -1718,45 +563,15 @@ bool DeepSeek4Backend::init_hybrid_model() { err.c_str()); return false; } - if (!init_expert_split_runtime(&err)) { - std::fprintf(stderr, "[deepseek4] failed to initialize expert_split runtime: %s\n", - err.c_str()); - return false; - } - if (deepseek4_expert_split_requires_parent_cuda_graph_disable( - expert_split_runtime_.enabled, - backend_is_cuda(backend_), - expert_split_state_.compute_runtime.targets)) { - setenv("GGML_CUDA_DISABLE_GRAPHS", "1", 1); - std::fprintf(stderr, - "[deepseek4] expert_split multi-target on CUDA: disabling CUDA graphs for parent runtime\n"); - } - if (expert_split_runtime_.enabled) { - const PlacementBackend device_backend = resolved_device_backend(cfg_.device); - const bool disable_cached_decode = - should_disable_ds4_expert_split_cached_decode( - /*expert_split_enabled=*/true, - device_backend, - backend_, - cfg_.device.gpu); - std::fprintf(stderr, - "[deepseek4] expert_split decode cache: %s parent_backend=%s device=%d\n", - disable_cached_decode ? "eager" : "cached", - placement_backend_name(device_backend), - cfg_.device.gpu); - } 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 cache_slots=%d hot_budget=%.2f GiB expert_budget=%.2f GiB\n", + "[deepseek4] hybrid experts ready: hot=%d cold=%d cold_backend=%s\n", moe_placement_.total_hot, total_cold, - cold_backend, - expert_split_cache_slots_, - gib(expert_split_hot_budget_bytes_), - gib(expert_split_total_budget_bytes_)); + cold_backend); return true; } @@ -1782,27 +597,8 @@ bool DeepSeek4Backend::unpark(const std::string & what) { int DeepSeek4Backend::do_prefill(const std::vector & tokens, const DaemonIO & io, int kv_offset) { - const bool use_hybrid_batch_prefill = - moe_hybrid_ && - backend_is_gpu(backend_) && - !env_flag_enabled("DFLASH_DS4_DISABLE_HYBRID_BATCH_PREFILL"); const int requested_chunk = cfg_.chunk > 0 ? cfg_.chunk : 512; - const PlacementBackend device_backend = resolved_device_backend(cfg_.device); - const int chunk = use_hybrid_batch_prefill - ? resolve_ds4_expert_split_prefill_chunk( - requested_chunk, - expert_split_runtime_.enabled, - device_backend, - backend_, - cfg_.device.gpu) - : (moe_hybrid_ ? 1 : requested_chunk); - if (use_hybrid_batch_prefill && - expert_split_runtime_.enabled && - chunk < requested_chunk) { - std::fprintf(stderr, - "[deepseek4] expert_split prefill chunk capped: requested=%d effective=%d\n", - requested_chunk, chunk); - } + const int chunk = moe_hybrid_ ? 1 : requested_chunk; const int n_total = (int)tokens.size(); int pos = kv_offset; last_logits_.clear(); @@ -1817,12 +613,11 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, if (timing && moe_hybrid_) { const char * backend_name = backend_ ? ggml_backend_name(backend_) : "(null)"; std::fprintf(stderr, - "[deepseek4] prefill plan: tokens=%d requested_chunk=%d effective_chunk=%d hybrid_batch=%s expert_split=%s intermediate_logits=%s backend=%s\n", + "[deepseek4] prefill plan: tokens=%d requested_chunk=%d effective_chunk=%d hybrid=%s intermediate_logits=%s backend=%s\n", n_total, requested_chunk, chunk, - use_hybrid_batch_prefill ? "on" : "off", - expert_split_runtime_.enabled ? "on" : "off", + moe_hybrid_ ? "on" : "off", force_intermediate_logits ? "on" : "off", backend_name ? backend_name : "(null)"); } @@ -1881,13 +676,7 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, int steps = 0; std::vector embed((size_t)w_.n_embd); std::vector logits; - const PlacementBackend device_backend = resolved_device_backend(cfg_.device); - const bool disable_cached_decode = - should_disable_ds4_expert_split_cached_decode( - expert_split_runtime_.enabled, - device_backend, - backend_, - cfg_.device.gpu); + const bool disable_cached_decode = false; for (int generated = 0; generated < n_gen; generated++) { if (io.cancelled) break; @@ -2085,8 +874,6 @@ void DeepSeek4Backend::shutdown() { free_deepseek4_cache(cache_); moe_hybrid_.reset(); expert_runtime_.reset(); - expert_split_runtime_.reset(); - expert_split_state_ = {}; layer_expert_bytes_.clear(); hc_state_.clear(); last_logits_.clear(); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index f83ff4f9c..5c26d21b5 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -11,7 +11,6 @@ #include "../common/moe_hybrid_placement.h" #include "../common/moe_hybrid_routing_stats.h" #include "../common/moe_hybrid_storage.h" -#include "../common/expert_split_state.h" #include "../common/moe_expert_compute.h" #include "deepseek4_internal.h" @@ -103,34 +102,13 @@ class DeepSeek4Backend : public ModelBackend { int max_ctx, MoeHybridPlacement & out, std::string * err) const; - bool load_expert_split_placement(const char * hotness_path, - const DeepSeek4Weights & w, - int max_ctx, - MoeHybridPlacement & out, - std::string * err); - bool build_expert_split_state_from_hotness(const MoeHybridRoutingStats & hotness, - uint64_t expert_budget_bytes, - uint64_t hot_budget_bytes, - int max_hot_per_layer, - const std::vector & configured_targets, - const DeepSeek4Weights & w, - std::string * err); - std::unique_ptr make_expert_split_local_target_compute( - const ExpertSplitComputeTargetRuntime & split_target, - std::string * err) const; bool init_single_target_expert_runtime(std::string * err); - bool init_expert_split_runtime(std::string * err); void maybe_save_routing_stats(); std::shared_ptr moe_hybrid_; MoeHybridPlacement moe_placement_; - ExpertSplitStateComponents expert_split_state_; MoeExpertComputeRuntime expert_runtime_; - MoeMultiTargetExpertRuntime expert_split_runtime_; std::vector layer_expert_bytes_; - uint64_t expert_split_total_budget_bytes_ = 0; - uint64_t expert_split_hot_budget_bytes_ = 0; - int expert_split_cache_slots_ = 0; std::shared_ptr routing_stats_; std::string routing_stats_out_path_; }; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index ddabe28dc..890810578 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -130,7 +130,6 @@ static void add_ffn_telemetry(DeepSeek4StepTelemetry * dst, dst->ffn_cold_us += src.cold_us; dst->ffn_combine_us += src.combine_us; dst->ffn_partition_us += src.partition_us; - dst->ffn_cache_promote_us += src.cache_promote_us; dst->ffn_hot_graph_builds += src.hot_graph_builds; dst->ffn_hot_graph_hits += src.hot_graph_hits; dst->ffn_cold_graph_builds += src.cold_graph_builds; @@ -2004,8 +2003,7 @@ static bool eval_ds4_hybrid_or_worker( ffn_normed_host, selected_host, weights_host, n_tokens, ffn_out_host, nullptr, hot_alloc, cold_alloc, expert_compute, expert_layer, - step_tel ? &ffn_tel : nullptr, - ffn_normed_backend); + step_tel ? &ffn_tel : nullptr); if (ffn_ok) { if (trace_decode) { std::fprintf(stderr, @@ -2220,19 +2218,30 @@ static bool ds4_try_gpu_resident_decode_ffn( MoeHybridFfnTelemetry ffn_tel; const auto ffn_t0 = Ds4TimingClock::now(); + std::vector selected_weights_storage; + if (!selected_weights && selected_weights_gpu) { + selected_weights_storage.resize((size_t)n_selected); + ggml_backend_tensor_get(selected_weights_gpu, + selected_weights_storage.data(), 0, + sizeof(float) * selected_weights_storage.size()); + selected_weights = selected_weights_storage.data(); + } + if (!selected_weights) { + return false; + } if (!eval_moe_hybrid_ffn_gpu_resident( backend, hybrid_cfg, desc, storage, cpu_backend, ffn_post_gpu, nullptr, gpu_ffn_state.state, selected_ids, selected_weights, n_selected, - selected_weights_gpu, - out_gpu, - expert_compute, expert_layer, - step_tel ? &ffn_tel : nullptr)) { + expert_compute, expert_layer)) { return false; } + if (out_gpu) { + *out_gpu = gpu_ffn_state.state.act_cur; + } if (step_tel) { step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); add_ffn_telemetry(step_tel, ffn_tel); @@ -3605,12 +3614,7 @@ static bool deepseek4_step_hybrid( expert_layers ? expert_layers + il : nullptr; const bool all_selected_hot = ds4_all_selected_hot( storage, hybrid_selected_host.data(), n_tokens * w.n_expert_used); - const bool can_skip_ffn_normed_host = - expert_worker == nullptr && - (all_selected_hot || - storage.cold_backend_kind == MoeHybridColdBackend::Gpu || - (expert_compute && expert_layer && - expert_compute->supports_batch_gpu_resident())); + const bool can_skip_ffn_normed_host = false; const auto route_read_t0 = Ds4TimingClock::now(); if (!use_gpu_resident_decode_ffn && !can_skip_ffn_normed_host) { const Ds4TensorReadback readbacks[] = { @@ -3720,11 +3724,7 @@ static bool deepseek4_step_hybrid( auto & storage = moe_hybrid.layers[(size_t) il]; const MoeExpertLayer * expert_layer = expert_layers ? expert_layers + il : nullptr; - const bool can_skip_ffn_normed_host = - expert_worker == nullptr && - (storage.cold_backend_kind == MoeHybridColdBackend::Gpu || - (expert_compute && expert_layer && - expert_compute->supports_batch_gpu_resident())); + const bool can_skip_ffn_normed_host = false; if (ok) { const auto route_read_t0 = Ds4TimingClock::now(); if (can_skip_ffn_normed_host) { @@ -3890,12 +3890,7 @@ static bool deepseek4_step_hybrid( } const bool all_selected_hot = ds4_all_selected_hot( storage, hybrid_selected_host.data(), n_tokens * w.n_expert_used); - const bool can_skip_ffn_normed_host = - expert_worker == nullptr && - (all_selected_hot || - storage.cold_backend_kind == MoeHybridColdBackend::Gpu || - (expert_compute && expert_layer && - expert_compute->supports_batch_gpu_resident())); + const bool can_skip_ffn_normed_host = false; if (!prefer_single_hip_route_readback && !use_gpu_resident_decode_ffn && !can_skip_ffn_normed_host) { @@ -3997,11 +3992,7 @@ static bool deepseek4_step_hybrid( auto & storage = moe_hybrid.layers[(size_t) il]; const MoeExpertLayer * expert_layer = expert_layers ? expert_layers + il : nullptr; - const bool can_skip_ffn_normed_host = - expert_worker == nullptr && - (storage.cold_backend_kind == MoeHybridColdBackend::Gpu || - (expert_compute && expert_layer && - expert_compute->supports_batch_gpu_resident())); + const bool can_skip_ffn_normed_host = false; const auto route_read_t0 = Ds4TimingClock::now(); if (can_skip_ffn_normed_host) { const Ds4TensorReadback route_readbacks[] = { diff --git a/server/src/deepseek4/deepseek4_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index 8f784d1f2..d8ea9c6ad 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -1,13 +1,17 @@ #include "deepseek4_hc_cuda.h" #include "common/gpu_runtime_compat.h" +#include "ggml-backend-impl.h" +#include "ggml-cuda/common.cuh" #include #include #include #include +#include #include +#include #include namespace dflash::common { @@ -63,7 +67,17 @@ struct HcCudaScratch { }; std::mutex g_mu; -HcCudaScratch g_scratch; +std::unordered_map> g_scratch_by_stream; + +HcCudaScratch & hc_scratch_for_stream(cudaStream_t stream) { + void * key = reinterpret_cast(stream); + auto it = g_scratch_by_stream.find(key); + if (it == g_scratch_by_stream.end()) { + auto inserted = g_scratch_by_stream.emplace(key, std::make_unique()); + it = inserted.first; + } + return *it->second; +} void hc_log_cuda_error(const char * label, cudaError_t err) { if (err != cudaSuccess) { @@ -111,6 +125,57 @@ __global__ void hc_mix_kernel(const float * x, if (tid == 0) mix[row] = smem[0]; } +__global__ void hc_finalize_inv_rms_kernel(const float * sums, + int n_sums, + int cols, + float eps, + float * inv_rms_out) { + __shared__ float smem[kThreads]; + const int tid = threadIdx.x; + float acc = 0.0f; + for (int i = tid; i < n_sums; i += blockDim.x) { + acc += sums[i]; + } + smem[tid] = acc; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + smem[tid] += smem[tid + stride]; + } + __syncthreads(); + } + if (tid == 0) { + inv_rms_out[0] = rsqrtf(smem[0] / (float) cols + eps); + } +} + +__global__ void hc_mix_kernel_device_inv(const float * x, + const __half * fn, + int cols, + const float * inv_rms_ptr, + float * mix) { + __shared__ float smem[kThreads]; + const int row = blockIdx.x; + const int tid = threadIdx.x; + const float inv_rms = inv_rms_ptr[0]; + float acc = 0.0f; + const __half * w = fn + (size_t) row * (size_t) cols; + for (int c = tid; c < cols; c += blockDim.x) { + acc += __half2float(w[c]) * (x[c] * inv_rms); + } + smem[tid] = acc; + __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]; + } +} + __device__ float hc_sigmoid(float x) { return 1.0f / (1.0f + expf(-x)); } @@ -257,6 +322,9 @@ bool hc_pre_device_locked(const void * hc_state_device, void * working_device, void * post_device, void * comb_device, + HcCudaScratch & scratch, + cudaStream_t stream, + bool sync_device, bool log_errors) { const int hc_dim = n_embd * n_hc; const int mix_dim = 2 * n_hc + n_hc * n_hc; @@ -266,80 +334,78 @@ bool hc_pre_device_locked(const void * hc_state_device, } return false; }; + const float * hc_state_ptr = static_cast(hc_state_device); - if (hc_state_device != g_scratch.d_state) { - cudaError_t err = cudaMemcpy(g_scratch.d_state, hc_state_device, - sizeof(float) * (size_t) hc_dim, - cudaMemcpyDeviceToDevice); - if (err != cudaSuccess) { - return fail("copy state d2d", err); - } + hc_sumsq_kernel<<>>( + hc_state_ptr, + hc_dim, + scratch.d_sums); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + return fail("sumsq kernel", err); } - hc_mix_norm_kernel<<>>( - g_scratch.d_state, - static_cast(fn_device), + hc_finalize_inv_rms_kernel<<<1, kThreads, 0, stream>>>( + scratch.d_sums, + kSums, hc_dim, - mix_dim, eps, - g_scratch.d_mix); - cudaError_t err = cudaGetLastError(); + scratch.d_sums); + err = cudaGetLastError(); + if (err != cudaSuccess) { + return fail("inv_rms kernel", err); + } + + hc_mix_kernel_device_inv<<>>( + hc_state_ptr, + static_cast(fn_device), + hc_dim, + scratch.d_sums, + scratch.d_mix); + err = cudaGetLastError(); if (err != cudaSuccess) { return fail("mix kernel", err); } - hc_finish_kernel<<<1, kThreads>>>( - g_scratch.d_state, - g_scratch.d_mix, + hc_finish_kernel<<<1, kThreads, 0, stream>>>( + hc_state_ptr, + scratch.d_mix, static_cast(scale_device), static_cast(base_device), n_embd, n_hc, sinkhorn_iters, - g_scratch.d_working, - g_scratch.d_post, - g_scratch.d_comb); + static_cast(working_device), + static_cast(post_device), + static_cast(comb_device)); err = cudaGetLastError(); if (err != cudaSuccess) { return fail("finish kernel", err); } -#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) - err = cudaDeviceSynchronize(); - if (err != cudaSuccess) { - return fail("device sync", err); - } -#endif - - if (working_device != g_scratch.d_working) { - err = cudaMemcpy(working_device, g_scratch.d_working, - sizeof(float) * (size_t) n_embd, - cudaMemcpyDeviceToDevice); + if (sync_device) { + err = cudaDeviceSynchronize(); if (err != cudaSuccess) { - return fail("copy working d2d", err); + return fail("finish sync", err); } } - if (post_device != g_scratch.d_post) { - err = cudaMemcpy(post_device, g_scratch.d_post, - sizeof(float) * (size_t) n_hc, - cudaMemcpyDeviceToDevice); - if (err != cudaSuccess) { - return fail("copy post d2d", err); - } + + return true; +} + +static cudaStream_t hc_backend_stream_or_default(ggml_backend_t backend) { + if (!backend || !backend->context) { + return nullptr; } - if (comb_device != g_scratch.d_comb) { - err = cudaMemcpy(comb_device, g_scratch.d_comb, - sizeof(float) * (size_t) n_hc * (size_t) n_hc, - cudaMemcpyDeviceToDevice); - if (err != cudaSuccess) { - return fail("copy comb d2d", err); - } + if (!backend->device) { + return nullptr; } - err = cudaDeviceSynchronize(); - if (err != cudaSuccess) { - return fail("post-copy sync", err); + const auto dev_type = ggml_backend_dev_type(backend->device); + if (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && + dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + return nullptr; } - - return true; + auto * cuda_ctx = static_cast(backend->context); + return cuda_ctx ? cuda_ctx->stream() : nullptr; } } // namespace @@ -355,30 +421,31 @@ bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, } const int hc_dim = n_embd * n_hc; std::lock_guard lock(g_mu); - if (!g_scratch.ensure((size_t)hc_dim, (size_t)n_embd, (size_t)n_hc)) { + HcCudaScratch & scratch = hc_scratch_for_stream(nullptr); + if (!scratch.ensure((size_t)hc_dim, (size_t)n_embd, (size_t)n_hc)) { return false; } - if (cudaMemcpy(g_scratch.d_state, hc_state_host, sizeof(float) * (size_t)hc_dim, + if (cudaMemcpy(scratch.d_state, hc_state_host, sizeof(float) * (size_t)hc_dim, cudaMemcpyHostToDevice) != cudaSuccess) { return false; } - hc_sumsq_kernel<<>>(g_scratch.d_state, hc_dim, g_scratch.d_sums); + hc_sumsq_kernel<<>>(scratch.d_state, hc_dim, scratch.d_sums); if (cudaGetLastError() != cudaSuccess) return false; std::vector sums(kSums); - if (cudaMemcpy(sums.data(), g_scratch.d_sums, sizeof(float) * sums.size(), + if (cudaMemcpy(sums.data(), scratch.d_sums, sizeof(float) * sums.size(), cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } float ss = 0.0f; for (float v : sums) ss += v; const float inv_rms = 1.0f / std::sqrt(ss / (float)hc_dim + eps); - hc_mix_kernel<<>>(g_scratch.d_state, + hc_mix_kernel<<>>(scratch.d_state, static_cast(fn_device), hc_dim, inv_rms, - g_scratch.d_mix); + scratch.d_mix); if (cudaGetLastError() != cudaSuccess) return false; - if (cudaMemcpy(mix_host, g_scratch.d_mix, sizeof(float) * kMixDim, + if (cudaMemcpy(mix_host, scratch.d_mix, sizeof(float) * kMixDim, cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } @@ -408,46 +475,50 @@ bool deepseek4_cuda_hc_pre(const float * hc_state_host, } std::lock_guard lock(g_mu); - if (!g_scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + HcCudaScratch & scratch = hc_scratch_for_stream(nullptr); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { return false; } - if (cudaMemcpy(g_scratch.d_state, hc_state_host, sizeof(float) * (size_t) hc_dim, + if (cudaMemcpy(scratch.d_state, hc_state_host, sizeof(float) * (size_t) hc_dim, cudaMemcpyHostToDevice) != cudaSuccess) { return false; } - if (cudaMemcpy(g_scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, + if (cudaMemcpy(scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, cudaMemcpyHostToDevice) != cudaSuccess) { return false; } - if (cudaMemcpy(g_scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, + if (cudaMemcpy(scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, cudaMemcpyHostToDevice) != cudaSuccess) { return false; } if (!hc_pre_device_locked( - g_scratch.d_state, + scratch.d_state, fn_device, - g_scratch.d_scale, - g_scratch.d_base, + scratch.d_scale, + scratch.d_base, n_embd, n_hc, sinkhorn_iters, eps, - g_scratch.d_working, - g_scratch.d_post, - g_scratch.d_comb, + scratch.d_working, + scratch.d_post, + scratch.d_comb, + scratch, + nullptr, + true, false)) { return false; } - if (cudaMemcpy(working_host, g_scratch.d_working, sizeof(float) * (size_t) n_embd, + if (cudaMemcpy(working_host, scratch.d_working, sizeof(float) * (size_t) n_embd, cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } - if (cudaMemcpy(post_host, g_scratch.d_post, sizeof(float) * (size_t) n_hc, + if (cudaMemcpy(post_host, scratch.d_post, sizeof(float) * (size_t) n_hc, cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } - if (cudaMemcpy(comb_host, g_scratch.d_comb, sizeof(float) * (size_t) n_hc * (size_t) n_hc, + if (cudaMemcpy(comb_host, scratch.d_comb, sizeof(float) * (size_t) n_hc * (size_t) n_hc, cudaMemcpyDeviceToHost) != cudaSuccess) { return false; } @@ -478,7 +549,54 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, } std::lock_guard lock(g_mu); - if (!g_scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + HcCudaScratch & scratch = hc_scratch_for_stream(nullptr); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + return false; + } + return hc_pre_device_locked(hc_state_device, + fn_device, + scale_device, + base_device, + n_embd, + n_hc, + sinkhorn_iters, + eps, + working_device, + post_device, + comb_device, + scratch, + nullptr, + true, + false); +} + +bool deepseek4_cuda_hc_pre_device_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const void * scale_device, + const void * base_device, + 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_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { + return false; + } + const int hc_dim = n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + if (mix_dim > kMaxMixDim) { + return false; + } + + cudaStream_t stream = hc_backend_stream_or_default(backend); + std::lock_guard lock(g_mu); + HcCudaScratch & scratch = hc_scratch_for_stream(stream); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { return false; } return hc_pre_device_locked(hc_state_device, @@ -492,6 +610,9 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, working_device, post_device, comb_device, + scratch, + stream, + false, false); } @@ -518,29 +639,82 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, } std::lock_guard lock(g_mu); - if (!g_scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + HcCudaScratch & scratch = hc_scratch_for_stream(nullptr); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { std::fprintf(stderr, "[deepseek4-hc-direct] ensure failed\n"); return false; } - if (cudaMemcpy(g_scratch.d_state, hc_state_device, sizeof(float) * (size_t) hc_dim, - cudaMemcpyDeviceToDevice) != cudaSuccess) { - hc_log_cuda_error("copy state d2d", cudaGetLastError()); + if (cudaMemcpy(scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, + cudaMemcpyHostToDevice) != cudaSuccess) { + hc_log_cuda_error("copy scale", cudaGetLastError()); return false; } - if (cudaMemcpy(g_scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, + if (cudaMemcpy(scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, + cudaMemcpyHostToDevice) != cudaSuccess) { + hc_log_cuda_error("copy base", cudaGetLastError()); + return false; + } + return hc_pre_device_locked(hc_state_device, + fn_device, + scratch.d_scale, + scratch.d_base, + n_embd, + n_hc, + sinkhorn_iters, + eps, + working_device, + post_device, + comb_device, + scratch, + nullptr, + true, + true); +} + +bool deepseek4_cuda_hc_pre_device_params_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const float * scale_host, + const float * base_host, + 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_host || !base_host || + !working_device || !post_device || !comb_device || + n_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { + return false; + } + const int hc_dim = n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + if (mix_dim > kMaxMixDim) { + return false; + } + + cudaStream_t stream = hc_backend_stream_or_default(backend); + std::lock_guard lock(g_mu); + HcCudaScratch & scratch = hc_scratch_for_stream(stream); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + std::fprintf(stderr, "[deepseek4-hc-direct] ensure failed\n"); + return false; + } + if (cudaMemcpy(scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, cudaMemcpyHostToDevice) != cudaSuccess) { hc_log_cuda_error("copy scale", cudaGetLastError()); return false; } - if (cudaMemcpy(g_scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, + if (cudaMemcpy(scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, cudaMemcpyHostToDevice) != cudaSuccess) { hc_log_cuda_error("copy base", cudaGetLastError()); return false; } - return hc_pre_device_locked(g_scratch.d_state, + return hc_pre_device_locked(hc_state_device, fn_device, - g_scratch.d_scale, - g_scratch.d_base, + scratch.d_scale, + scratch.d_base, n_embd, n_hc, sinkhorn_iters, @@ -548,6 +722,9 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, working_device, post_device, comb_device, + scratch, + stream, + false, true); } diff --git a/server/src/deepseek4/deepseek4_hc_cuda.h b/server/src/deepseek4/deepseek4_hc_cuda.h index d2268ac12..b5b0eddee 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.h +++ b/server/src/deepseek4/deepseek4_hc_cuda.h @@ -2,6 +2,8 @@ #include +#include "ggml-backend.h" + namespace dflash::common { bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, @@ -35,6 +37,19 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, void * post_device, void * comb_device); +bool deepseek4_cuda_hc_pre_device_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const void * scale_device, + const void * base_device, + int n_embd, + int n_hc, + int sinkhorn_iters, + float eps, + void * working_device, + void * post_device, + void * comb_device); + bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, const void * fn_device, const float * scale_host, @@ -47,4 +62,17 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, void * post_device, void * comb_device); +bool deepseek4_cuda_hc_pre_device_params_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const float * scale_host, + const float * base_host, + int n_embd, + int n_hc, + int sinkhorn_iters, + float eps, + void * working_device, + void * post_device, + void * comb_device); + } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 0b248b2bd..a36a735de 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -23,7 +23,6 @@ #include "ggml-backend.h" #include "internal.h" -#include "common/expert_split_compute_runtime.h" #include "common/layer_split_utils.h" namespace dflash::common { @@ -311,102 +310,6 @@ struct DeepSeek4BackendConfig { int max_ctx = 0; // 0 = auto from SWA + compression capacity }; -struct DeepSeek4ExpertSplitGpuMemoryInfo { - uint64_t free_bytes = 0; - uint64_t total_bytes = 0; -}; - -struct DeepSeek4ExpertSplitBudgetMemoryInfo { - uint64_t total_bytes = 0; - uint64_t free_bytes = 0; - uint64_t core_bytes = 0; - uint64_t kv_bytes = 0; - uint64_t warm_bytes = 0; - uint64_t safety_bytes = 0; - uint64_t total_expert_bytes = 0; - bool parent_is_igpu = false; - uint64_t host_available_bytes = 0; - uint64_t host_budget_cap_bytes = 0; -}; - -struct DeepSeek4ExpertSplitHotCachePlan { - uint64_t hot_bytes = 0; - int cache_slots = 0; -}; - -int deepseek4_expert_split_prefill_chunk_limit_from_memory( - int requested_chunk, - bool expert_split_enabled, - bool parent_is_gpu, - const DeepSeek4ExpertSplitGpuMemoryInfo & memory); - -bool deepseek4_expert_split_should_disable_cached_decode_from_memory( - bool expert_split_enabled, - bool parent_is_gpu, - bool parent_is_cuda, - bool parent_is_hip, - const DeepSeek4ExpertSplitGpuMemoryInfo & memory); - -bool deepseek4_expert_split_requires_parent_cuda_graph_disable( - bool expert_split_enabled, - bool parent_is_cuda, - const std::vector & targets); - -bool deepseek4_should_use_gpu_resident_decode_ffn_policy( - bool use_cached_decode, - bool env_enabled, - bool env_disabled, - bool backend_is_hip, - bool backend_is_integrated_gpu, - bool worker_active); - -uint64_t deepseek4_expert_split_budget_from_memory( - const DeepSeek4ExpertSplitBudgetMemoryInfo & memory); - -uint64_t deepseek4_expert_split_primary_capacity_for_targets( - uint64_t expert_budget_bytes, - uint64_t hot_budget_bytes, - size_t configured_targets); - -uint64_t deepseek4_expert_split_non_cpu_capacity_bytes( - const std::vector & targets); - -uint64_t deepseek4_expert_split_effective_budget_for_targets( - uint64_t requested_budget_bytes, - uint64_t total_expert_bytes, - const std::vector & targets); - -uint64_t deepseek4_expert_split_effective_hot_budget( - uint64_t expert_budget_bytes, - uint64_t hot_budget_bytes, - size_t configured_targets); - -bool deepseek4_expert_split_hash_ids_to_hotness_counts( - const int32_t * expert_ids, - int n_token_ids, - int n_expert_used, - int n_expert, - std::vector & out_counts); - -bool deepseek4_expert_split_route_bias_to_hotness_counts( - const float * route_bias, - int n_expert, - uint64_t layer_total, - std::vector & out_counts); - -int deepseek4_expert_split_effective_cache_slots( - int requested_cache_slots, - size_t configured_targets); - -DeepSeek4ExpertSplitHotCachePlan deepseek4_expert_split_hot_cache_plan( - uint64_t expert_budget_bytes, - uint64_t total_expert_bytes, - int n_expert, - int n_expert_used, - bool parent_is_igpu, - uint64_t igpu_host_cap_bytes = 0, - uint64_t igpu_free_bytes = 0); - // ─── Function declarations ────────────────────────────────────────────── bool load_deepseek4_gguf(const std::string & path, diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index ef3aa9fbf..07096f00f 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -711,7 +711,13 @@ bool build_deepseek4_moe_hybrid_storage_from_file( const MoeHybridPlacement & placement, const MoeHybridConfig * cfg_override, MoeHybridStorage & out, - std::string * err) { + std::string * err, + int cache_slots, + bool load_cold_tensors, + const std::vector> * cold_expert_order_by_layer) { + (void) cache_slots; + (void) load_cold_tensors; + (void) cold_expert_order_by_layer; ggml_context * expert_meta = nullptr; gguf_init_params gip{}; gip.no_alloc = true; diff --git a/server/src/deepseek4/expert_ipc.h b/server/src/deepseek4/expert_ipc.h new file mode 100644 index 000000000..6a4bb206f --- /dev/null +++ b/server/src/deepseek4/expert_ipc.h @@ -0,0 +1,76 @@ +// Compatibility shim for PR489's DeepSeek4 graph worker hook. +// +// The online PR489 path never creates an ExpertIpcClient; the hook is kept in +// the graph signature only so this patch can stay source-compatible with local +// experimental expert IPC trees without requiring those files in a clean PR489 +// checkout. +#pragma once + +#include +#include + +namespace dflash::common { + +struct ExpertIpcTiming { + uint64_t parent_write_us = 0; + uint64_t parent_wait_us = 0; + uint64_t parent_read_us = 0; + uint64_t worker_request_read_us = 0; + uint64_t worker_partition_us = 0; + uint64_t worker_resident_eval_us = 0; + uint64_t worker_miss_build_us = 0; + uint64_t worker_miss_eval_us = 0; + uint64_t request_bytes = 0; + uint64_t response_bytes = 0; + uint64_t worker_hot_graph_builds = 0; + uint64_t worker_hot_graph_hits = 0; + uint64_t worker_cold_graph_builds = 0; + uint64_t worker_cold_graph_hits = 0; + uint64_t worker_hot_graph_build_us = 0; + uint64_t worker_hot_input_us = 0; + uint64_t worker_hot_compute_us = 0; + uint64_t worker_hot_read_us = 0; + uint64_t worker_cold_graph_build_us = 0; + uint64_t worker_cold_input_us = 0; + uint64_t worker_cold_compute_us = 0; + uint64_t worker_cold_read_us = 0; +}; + +class ExpertIpcClient { +public: + struct PendingEval {}; + + bool active() const { return false; } + + bool eval_begin(int, + int, + int, + int, + const float *, + const int32_t *, + const float *, + PendingEval &, + ExpertIpcTiming * = nullptr) { + return false; + } + + bool eval_end(PendingEval &, + std::vector &, + ExpertIpcTiming * = nullptr) { + return false; + } + + bool eval(int, + int, + int, + int, + const float *, + const int32_t *, + const float *, + std::vector &, + ExpertIpcTiming * = nullptr) { + return false; + } +}; + +} // namespace dflash::common From 16eba371ff25ec90e531948346b6d6c724fbfbc1 Mon Sep 17 00:00:00 2001 From: weicj Date: Thu, 9 Jul 2026 15:37:33 +0800 Subject: [PATCH 5/7] fix(dflash): address DeepSeek4 ROCmFPx review feedback --- server/CMakeLists.txt | 5 - server/src/deepseek4/deepseek4_hc_cuda.cu | 121 ---------------------- server/src/deepseek4/deepseek4_hc_cuda.h | 28 ----- server/src/deepseek4/deepseek4_internal.h | 5 +- server/src/deepseek4/deepseek4_loader.cpp | 8 +- 5 files changed, 2 insertions(+), 165 deletions(-) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 5f44fdeb1..2e54ce514 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -591,11 +591,6 @@ find_package(OpenMP) if(OpenMP_CXX_FOUND) target_link_libraries(dflash_common PRIVATE OpenMP::OpenMP_CXX) endif() -if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") - target_link_libraries(dflash_common PUBLIC CUDA::cudart) -elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") - target_link_libraries(dflash_common PUBLIC hip::host) -endif() if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang") target_compile_options(dflash_common PRIVATE diff --git a/server/src/deepseek4/deepseek4_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index d8ea9c6ad..f2c435241 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -1,8 +1,6 @@ #include "deepseek4_hc_cuda.h" #include "common/gpu_runtime_compat.h" -#include "ggml-backend-impl.h" -#include "ggml-cuda/common.cuh" #include @@ -392,22 +390,6 @@ bool hc_pre_device_locked(const void * hc_state_device, return true; } -static cudaStream_t hc_backend_stream_or_default(ggml_backend_t backend) { - if (!backend || !backend->context) { - return nullptr; - } - if (!backend->device) { - return nullptr; - } - const auto dev_type = ggml_backend_dev_type(backend->device); - if (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && - dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { - return nullptr; - } - auto * cuda_ctx = static_cast(backend->context); - return cuda_ctx ? cuda_ctx->stream() : nullptr; -} - } // namespace bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, @@ -570,52 +552,6 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, false); } -bool deepseek4_cuda_hc_pre_device_on_backend(ggml_backend_t backend, - const void * hc_state_device, - const void * fn_device, - const void * scale_device, - const void * base_device, - 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_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { - return false; - } - const int hc_dim = n_embd * n_hc; - const int mix_dim = 2 * n_hc + n_hc * n_hc; - if (mix_dim > kMaxMixDim) { - return false; - } - - cudaStream_t stream = hc_backend_stream_or_default(backend); - std::lock_guard lock(g_mu); - HcCudaScratch & scratch = hc_scratch_for_stream(stream); - if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { - return false; - } - return hc_pre_device_locked(hc_state_device, - fn_device, - scale_device, - base_device, - n_embd, - n_hc, - sinkhorn_iters, - eps, - working_device, - post_device, - comb_device, - scratch, - stream, - false, - false); -} - bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, const void * fn_device, const float * scale_host, @@ -671,61 +607,4 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, true); } -bool deepseek4_cuda_hc_pre_device_params_on_backend(ggml_backend_t backend, - const void * hc_state_device, - const void * fn_device, - const float * scale_host, - const float * base_host, - 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_host || !base_host || - !working_device || !post_device || !comb_device || - n_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { - return false; - } - const int hc_dim = n_embd * n_hc; - const int mix_dim = 2 * n_hc + n_hc * n_hc; - if (mix_dim > kMaxMixDim) { - return false; - } - - cudaStream_t stream = hc_backend_stream_or_default(backend); - std::lock_guard lock(g_mu); - HcCudaScratch & scratch = hc_scratch_for_stream(stream); - if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { - std::fprintf(stderr, "[deepseek4-hc-direct] ensure failed\n"); - return false; - } - if (cudaMemcpy(scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, - cudaMemcpyHostToDevice) != cudaSuccess) { - hc_log_cuda_error("copy scale", cudaGetLastError()); - return false; - } - if (cudaMemcpy(scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, - cudaMemcpyHostToDevice) != cudaSuccess) { - hc_log_cuda_error("copy base", cudaGetLastError()); - return false; - } - return hc_pre_device_locked(hc_state_device, - fn_device, - scratch.d_scale, - scratch.d_base, - n_embd, - n_hc, - sinkhorn_iters, - eps, - working_device, - post_device, - comb_device, - scratch, - stream, - false, - true); -} - } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_hc_cuda.h b/server/src/deepseek4/deepseek4_hc_cuda.h index b5b0eddee..d2268ac12 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.h +++ b/server/src/deepseek4/deepseek4_hc_cuda.h @@ -2,8 +2,6 @@ #include -#include "ggml-backend.h" - namespace dflash::common { bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, @@ -37,19 +35,6 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, void * post_device, void * comb_device); -bool deepseek4_cuda_hc_pre_device_on_backend(ggml_backend_t backend, - const void * hc_state_device, - const void * fn_device, - const void * scale_device, - const void * base_device, - int n_embd, - int n_hc, - int sinkhorn_iters, - float eps, - void * working_device, - void * post_device, - void * comb_device); - bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, const void * fn_device, const float * scale_host, @@ -62,17 +47,4 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, void * post_device, void * comb_device); -bool deepseek4_cuda_hc_pre_device_params_on_backend(ggml_backend_t backend, - const void * hc_state_device, - const void * fn_device, - const float * scale_host, - const float * base_host, - int n_embd, - int n_hc, - int sinkhorn_iters, - float eps, - void * working_device, - void * post_device, - void * comb_device); - } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index a36a735de..f13612255 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -400,10 +400,7 @@ bool build_deepseek4_moe_hybrid_storage_from_file( const MoeHybridPlacement & placement, const MoeHybridConfig * cfg_override, MoeHybridStorage & out, - std::string * err = nullptr, - int cache_slots = 0, - bool load_cold_tensors = true, - const std::vector> * cold_expert_order_by_layer = nullptr); + std::string * err = nullptr); bool build_deepseek4_moe_hybrid_storage_from_file( const std::string & path, diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index 07096f00f..ef3aa9fbf 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -711,13 +711,7 @@ bool build_deepseek4_moe_hybrid_storage_from_file( const MoeHybridPlacement & placement, const MoeHybridConfig * cfg_override, MoeHybridStorage & out, - std::string * err, - int cache_slots, - bool load_cold_tensors, - const std::vector> * cold_expert_order_by_layer) { - (void) cache_slots; - (void) load_cold_tensors; - (void) cold_expert_order_by_layer; + std::string * err) { ggml_context * expert_meta = nullptr; gguf_init_params gip{}; gip.no_alloc = true; From 2555e41c697ddd38aa6ecd20a2c09858f16c1420 Mon Sep 17 00:00:00 2001 From: weicj Date: Fri, 10 Jul 2026 07:50:43 +0800 Subject: [PATCH 6/7] fix(deepseek4): stabilize HC decode output --- server/src/deepseek4/deepseek4_graph.cpp | 1447 ++++++++++++++++++--- server/src/deepseek4/deepseek4_hc_cuda.cu | 677 +++++++++- server/src/deepseek4/deepseek4_hc_cuda.h | 67 + server/src/deepseek4/deepseek4_loader.cpp | 17 +- server/src/server/http_server.cpp | 191 ++- 5 files changed, 2094 insertions(+), 305 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 890810578..e6ab64734 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -38,17 +38,35 @@ namespace dflash::common { namespace { using Ds4TimingClock = std::chrono::steady_clock; +static std::atomic ds4_request_graph_epoch{1}; static uint64_t ds4_elapsed_us(Ds4TimingClock::time_point start, Ds4TimingClock::time_point end) { return (uint64_t)std::chrono::duration_cast(end - start).count(); } +static uint64_t ds4_current_request_graph_epoch() { + return ds4_request_graph_epoch.load(std::memory_order_relaxed); +} + static bool ds4_env_flag(const char * name) { const char * value = std::getenv(name); return value && value[0] && std::strcmp(value, "0") != 0; } +static int64_t ds4_env_i64(const char * name, int64_t fallback) { + const char * value = std::getenv(name); + if (!value || !value[0]) { + return fallback; + } + char * end = nullptr; + const long long parsed = std::strtoll(value, &end, 10); + if (end == value) { + return fallback; + } + return (int64_t) parsed; +} + static size_t ds4_full_prefill_graph_size(int n_tokens) { if (n_tokens <= 1) { // The non-hybrid all-hot decode path still builds one full-layer graph @@ -187,6 +205,10 @@ static const float * ds4_get_or_load_route_bias_host( } // namespace +void deepseek4_invalidate_request_graph_caches() { + ds4_request_graph_epoch.fetch_add(1, std::memory_order_relaxed); +} + struct DeepSeek4I32InputBinding { ggml_tensor * tensor = nullptr; int32_t value = 0; @@ -225,6 +247,9 @@ struct DeepSeek4CachedDecodeFfnGraph { int layer_idx = -1; int n_tokens = 0; bool hash_routed = false; + size_t estimated_bytes = 0; + uint64_t last_used_epoch = 0; + uint32_t reuse_count = 0; StepGraph sg; ggml_tensor * hash_ids = nullptr; @@ -242,6 +267,9 @@ struct DeepSeek4CachedDecodeFfnGraph { layer_idx = -1; n_tokens = 0; hash_routed = false; + estimated_bytes = 0; + last_used_epoch = 0; + reuse_count = 0; } }; @@ -364,6 +392,7 @@ struct DeepSeek4CachedDecodeAttnGraph { bool flush_boundary = false; bool compressed = false; bool indexed = false; + bool uses_shared_inputs = false; StepGraph sg; DeepSeek4AttentionGraphInputs inputs; @@ -391,6 +420,7 @@ struct DeepSeek4CachedDecodeAttnGraph { flush_boundary = false; compressed = false; indexed = false; + uses_shared_inputs = false; } }; @@ -478,6 +508,122 @@ struct DeepSeek4LayerRangeScratch { } }; +static size_t ds4_estimate_cached_ffn_graph_bytes( + const DeepSeek4Weights & w, + int n_tokens, + bool hash_routed) { + if (n_tokens <= 0) { + return 0; + } + const size_t tokens = (size_t)n_tokens; + const size_t embd = (size_t)std::max(w.n_embd, 1); + const size_t experts = (size_t)std::max(w.n_expert, 1); + const size_t used = (size_t)std::max(w.n_expert_used, 1); + const size_t ff = (size_t)std::max(w.n_ff_exp, 1); + + // Conservative working-set estimate for cache pressure decisions. + size_t f32_elems_per_token = 0; + f32_elems_per_token += embd * 6; + f32_elems_per_token += experts * 2; + f32_elems_per_token += ff * used * 4; + if (!hash_routed) { + f32_elems_per_token += experts; + } + size_t bytes = tokens * f32_elems_per_token * sizeof(float); + bytes += tokens * used * sizeof(int32_t); + bytes += 4ull * 1024ull * 1024ull; + return bytes; +} + +static size_t ds4_prefill_ffn_cache_budget_bytes() { + const int64_t budget_mb = + ds4_env_i64("DFLASH_DS4_PREFILL_FFN_CACHE_BUDGET_MB", 4096); + if (budget_mb < 0) { + return (size_t)-1; + } + return (size_t)budget_mb * 1024ull * 1024ull; +} + +static size_t ds4_prefill_ffn_cache_retain_target_bytes(size_t budget) { + if (budget == (size_t)-1) { + return budget; + } + int64_t pct = + ds4_env_i64("DFLASH_DS4_PREFILL_FFN_CACHE_RETAIN_PCT", 50); + pct = std::max(0, std::min(100, pct)); + return budget * (size_t)pct / 100ull; +} + +static bool ds4_maybe_evict_prefill_ffn_cache_for_decode( + std::vector & graphs, + const DeepSeek4Weights & w) { + const size_t budget = ds4_prefill_ffn_cache_budget_bytes(); + if (budget == (size_t)-1) { + return false; + } + + size_t total = 0; + size_t live = 0; + for (auto & g : graphs) { + if (!g.valid()) { + continue; + } + if (g.estimated_bytes == 0) { + g.estimated_bytes = + ds4_estimate_cached_ffn_graph_bytes(w, g.n_tokens, g.hash_routed); + } + total += g.estimated_bytes; + ++live; + } + if (total <= budget) { + return false; + } + + const size_t target = ds4_prefill_ffn_cache_retain_target_bytes(budget); + size_t freed = 0; + size_t evicted = 0; + while (total > target) { + size_t victim = graphs.size(); + size_t victim_bytes = 0; + for (size_t i = 0; i < graphs.size(); ++i) { + const auto & g = graphs[i]; + if (!g.valid()) { + continue; + } + const size_t bytes = g.estimated_bytes; + if (victim == graphs.size() || + bytes > victim_bytes || + (bytes == victim_bytes && + (g.reuse_count < graphs[victim].reuse_count || + (g.reuse_count == graphs[victim].reuse_count && + g.last_used_epoch < graphs[victim].last_used_epoch)))) { + victim = i; + victim_bytes = bytes; + } + } + if (victim == graphs.size()) { + break; + } + graphs[victim].free(); + freed += victim_bytes; + total = total > victim_bytes ? total - victim_bytes : 0; + ++evicted; + } + + if (ds4_env_flag("DFLASH_DS4_PREFILL_FFN_CACHE_TRACE")) { + std::fprintf(stderr, + "[deepseek4] prefill-ffn-cache evict-for-decode live=%zu evicted=%zu freed=%.1fMB retained=%.1fMB budget=%.1fMB target=%.1fMB\n", + live, + evicted, + (double)freed / (1024.0 * 1024.0), + (double)total / (1024.0 * 1024.0), + (double)budget / (1024.0 * 1024.0), + (double)target / (1024.0 * 1024.0)); + std::fflush(stderr); + } + return evicted > 0; +} + static bool build_cached_decode_ffn_graph( DeepSeek4CachedDecodeFfnGraph & out, ggml_backend_t backend, @@ -565,6 +711,9 @@ static bool build_cached_decode_ffn_graph( out.layer_idx = layer_idx; out.n_tokens = n_tokens; out.hash_routed = hash_routed; + out.estimated_bytes = ds4_estimate_cached_ffn_graph_bytes(w, n_tokens, hash_routed); + out.last_used_epoch = 0; + out.reuse_count = 0; return true; } @@ -1434,6 +1583,7 @@ struct DeepSeek4CachedDecodeHcPreGraph { int n_tokens = 0; bool ffn = false; StepGraph sg; + ggml_tensor * pre = nullptr; ggml_tensor * post = nullptr; ggml_tensor * comb = nullptr; @@ -1441,7 +1591,7 @@ struct DeepSeek4CachedDecodeHcPreGraph { return owner_ctx && backend && layer_idx >= 0 && n_tokens > 0 && sg.ctx && sg.gf && sg.alloc && sg.inp_embed && sg.hidden_states && - post && comb; + pre && post && comb; } void free() { @@ -1451,6 +1601,7 @@ struct DeepSeek4CachedDecodeHcPreGraph { layer_idx = -1; n_tokens = 0; ffn = false; + pre = nullptr; post = nullptr; comb = nullptr; } @@ -1485,6 +1636,136 @@ struct DeepSeek4CachedDecodeHcPostGraph { } }; +// Per-step decode scalar inputs shared by all cached per-layer attention graphs. +// Values depend only on (kv_start, ratio), so all layers with the same ratio can +// reuse one i32/i64 bundle instead of issuing many tiny tensor_set calls. +struct Ds4DecodeSharedInputs { + static constexpr int MAX_RATIOS = 4; + + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + int ratios[MAX_RATIOS] = {0}; + int n_ratios = 0; + ggml_tensor * i32_bundle = nullptr; // [6 * n_ratios] + ggml_tensor * i64_bundle = nullptr; // [5 * n_ratios] + + ggml_tensor * v_rope_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_neg_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_ape_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_comp_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_ape[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_cpos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_raw_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_state_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_comp_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_state[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_comp[MAX_RATIOS] = {nullptr}; + + void free() { + if (buf) { + ggml_backend_buffer_free(buf); + buf = nullptr; + } + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } + owner_ctx = nullptr; + backend = nullptr; + n_ratios = 0; + i32_bundle = nullptr; + i64_bundle = nullptr; + } + + int slot(int ratio) const { + for (int i = 0; i < n_ratios; ++i) { + if (ratios[i] == ratio) { + return i; + } + } + return -1; + } + + bool ensure(const DeepSeek4Weights & w, ggml_backend_t bk) { + if (ctx && owner_ctx == w.ctx && backend == bk) { + return true; + } + free(); + for (int il = 0; il < w.n_layer; ++il) { + const int r = (int) w.compress_ratios[il]; + if (slot(r) >= 0) { + continue; + } + if (n_ratios >= MAX_RATIOS) { + return false; + } + ratios[n_ratios++] = r; + } + + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * (size_t)(2 + 11 * MAX_RATIOS) + 4096; + params.mem_buffer = nullptr; + params.no_alloc = true; + ctx = ggml_init(params); + if (!ctx) { + return false; + } + i32_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 6 * (int64_t)n_ratios); + i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 5 * (int64_t)n_ratios); + for (int s = 0; s < n_ratios; ++s) { + v_rope_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t)s * 6 + 0) * sizeof(int32_t)); + v_neg_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t)s * 6 + 1) * sizeof(int32_t)); + v_ape_row[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t)s * 6 + 2) * sizeof(int32_t)); + v_comp_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t)s * 6 + 3) * sizeof(int32_t)); + v_index_ape[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t)s * 6 + 4) * sizeof(int32_t)); + v_index_cpos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t)s * 6 + 5) * sizeof(int32_t)); + v_raw_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t)s * 5 + 0) * sizeof(int64_t)); + v_state_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t)s * 5 + 1) * sizeof(int64_t)); + v_comp_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t)s * 5 + 2) * sizeof(int64_t)); + v_index_state[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t)s * 5 + 3) * sizeof(int64_t)); + v_index_comp[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t)s * 5 + 4) * sizeof(int64_t)); + } + buf = ggml_backend_alloc_ctx_tensors(ctx, bk); + if (!buf) { + free(); + return false; + } + owner_ctx = w.ctx; + backend = bk; + return true; + } + + void set_step(const DeepSeek4Weights & w, int kv_start) { + const int token_pos = kv_start; + int32_t i32v[6 * MAX_RATIOS] = {0}; + int64_t i64v[5 * MAX_RATIOS] = {0}; + for (int s = 0; s < n_ratios; ++s) { + const int ratio = ratios[s]; + i32v[s * 6 + 0] = kv_start; + i32v[s * 6 + 1] = -kv_start; + i64v[s * 5 + 0] = kv_start % w.n_swa; + if (ratio > 0) { + const int pos_mod = token_pos % ratio; + i32v[s * 6 + 2] = pos_mod; + i32v[s * 6 + 3] = token_pos + 1 - ratio; + i64v[s * 5 + 1] = (ratio == 4) ? (int64_t)(ratio + pos_mod) : (int64_t)pos_mod; + i64v[s * 5 + 2] = token_pos / ratio; + } + if (ratio == 4) { + const int pos_mod = token_pos % ratio; + i32v[s * 6 + 4] = pos_mod; + i32v[s * 6 + 5] = token_pos + 1 - ratio; + i64v[s * 5 + 3] = ratio + pos_mod; + i64v[s * 5 + 4] = token_pos / ratio; + } + } + ggml_backend_tensor_set(i32_bundle, i32v, 0, sizeof(int32_t) * 6 * (size_t)n_ratios); + ggml_backend_tensor_set(i64_bundle, i64v, 0, sizeof(int64_t) * 5 * (size_t)n_ratios); + } +}; + static bool build_cached_decode_attn_graph( DeepSeek4CachedDecodeAttnGraph & out, ggml_backend_t backend, @@ -1496,7 +1777,8 @@ static bool build_cached_decode_attn_graph( int raw_attn_count, int comp_attn_count, int index_comp_count, - bool flush_boundary) { + bool flush_boundary, + const Ds4DecodeSharedInputs * shared_inputs = nullptr) { out.free(); const size_t ctx_size = 48 * 1024 * 1024; @@ -1523,32 +1805,52 @@ static bool build_cached_decode_attn_graph( out.sg.gf = ggml_new_graph_custom( out.sg.ctx, DS4_CACHED_DECODE_ATTN_GRAPH_SIZE, false); - out.inputs.rope_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.neg_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - ggml_set_input(out.inputs.rope_pos); - ggml_set_input(out.inputs.neg_pos); - - out.inputs.raw_kv_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - ggml_set_input(out.inputs.raw_kv_rows); - if (ratio > 0) { - out.inputs.attn_ape_row = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.attn_comp_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.attn_state_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - out.inputs.attn_comp_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - ggml_set_input(out.inputs.attn_ape_row); - ggml_set_input(out.inputs.attn_comp_pos); - ggml_set_input(out.inputs.attn_state_rows); - ggml_set_input(out.inputs.attn_comp_rows); - } - if (ratio == 4) { - out.inputs.index_ape_row = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.index_comp_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.index_state_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - out.inputs.index_comp_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - ggml_set_input(out.inputs.index_ape_row); - ggml_set_input(out.inputs.index_comp_pos); - ggml_set_input(out.inputs.index_state_rows); - ggml_set_input(out.inputs.index_comp_rows); + const int shared_slot = shared_inputs ? shared_inputs->slot(ratio) : -1; + if (shared_slot >= 0) { + out.inputs.rope_pos = shared_inputs->v_rope_pos[shared_slot]; + out.inputs.neg_pos = shared_inputs->v_neg_pos[shared_slot]; + out.inputs.raw_kv_rows = shared_inputs->v_raw_row[shared_slot]; + if (ratio > 0) { + out.inputs.attn_ape_row = shared_inputs->v_ape_row[shared_slot]; + out.inputs.attn_comp_pos = shared_inputs->v_comp_pos[shared_slot]; + out.inputs.attn_state_rows = shared_inputs->v_state_row[shared_slot]; + out.inputs.attn_comp_rows = shared_inputs->v_comp_row[shared_slot]; + } + if (ratio == 4) { + out.inputs.index_ape_row = shared_inputs->v_index_ape[shared_slot]; + out.inputs.index_comp_pos = shared_inputs->v_index_cpos[shared_slot]; + out.inputs.index_state_rows = shared_inputs->v_index_state[shared_slot]; + out.inputs.index_comp_rows = shared_inputs->v_index_comp[shared_slot]; + } + out.uses_shared_inputs = true; + } else { + out.inputs.rope_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.neg_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + ggml_set_input(out.inputs.rope_pos); + ggml_set_input(out.inputs.neg_pos); + + out.inputs.raw_kv_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + ggml_set_input(out.inputs.raw_kv_rows); + if (ratio > 0) { + out.inputs.attn_ape_row = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.attn_comp_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.attn_state_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + out.inputs.attn_comp_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + ggml_set_input(out.inputs.attn_ape_row); + ggml_set_input(out.inputs.attn_comp_pos); + ggml_set_input(out.inputs.attn_state_rows); + ggml_set_input(out.inputs.attn_comp_rows); + } + if (ratio == 4) { + out.inputs.index_ape_row = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.index_comp_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.index_state_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + out.inputs.index_comp_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + ggml_set_input(out.inputs.index_ape_row); + ggml_set_input(out.inputs.index_comp_pos); + ggml_set_input(out.inputs.index_state_rows); + ggml_set_input(out.inputs.index_comp_rows); + } } std::vector i32_inputs; @@ -1675,29 +1977,31 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, !ggml_backend_buffer_is_host(scale_tensor->buffer) && !ggml_backend_buffer_is_host(base_tensor->buffer); if (can_use_device_params) { - return deepseek4_cuda_hc_pre_device(hc_state->data, - fn_tensor->data, - scale_tensor->data, - base_tensor->data, - n_embd, - n_hc, - sinkhorn_iters, - hc_eps, - working->data, - post->data, - comb->data); - } - return deepseek4_cuda_hc_pre_device_params(hc_state->data, - fn_tensor->data, - scale_data, - base_data, - n_embd, - n_hc, - sinkhorn_iters, - hc_eps, - working->data, - post->data, - comb->data); + return deepseek4_cuda_hc_pre_device_on_backend(backend, + hc_state->data, + fn_tensor->data, + scale_tensor->data, + base_tensor->data, + n_embd, + n_hc, + sinkhorn_iters, + hc_eps, + working->data, + post->data, + comb->data); + } + return deepseek4_cuda_hc_pre_device_params_on_backend(backend, + hc_state->data, + fn_tensor->data, + scale_data, + base_data, + n_embd, + n_hc, + sinkhorn_iters, + hc_eps, + working->data, + post->data, + comb->data); #else (void) working; (void) post; @@ -1719,6 +2023,121 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, #endif } +static bool ds4_try_gpu_hc_post_device(ggml_backend_t backend, + ggml_tensor * residual_hc, + const ggml_tensor * block_out, + const ggml_tensor * post, + const ggml_tensor * comb, + int n_embd, + int n_hc, + ggml_tensor * out_hc) { +#if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + if (!residual_hc || !block_out || !post || !comb || !out_hc || + !residual_hc->data || !block_out->data || !post->data || + !comb->data || !out_hc->data) { + return false; + } + return deepseek4_cuda_hc_post_device_on_backend(backend, + residual_hc->data, + block_out->data, + post->data, + comb->data, + n_embd, + n_hc, + out_hc->data); +#else + (void) backend; + (void) residual_hc; + (void) block_out; + (void) post; + (void) comb; + (void) n_embd; + (void) n_hc; + (void) out_hc; + return false; +#endif +} + +static bool ds4_try_gpu_hc_pre_batch_device(ggml_tensor * working, + ggml_tensor * post, + ggml_tensor * comb, + ggml_tensor * pre, + ggml_backend_t backend, + ggml_tensor * hc_state, + ggml_tensor * fn_tensor, + ggml_tensor * scale_tensor, + ggml_tensor * base_tensor, + const float * scale_data, + const float * base_data, + int n_tokens, + int n_embd, + int n_hc, + int sinkhorn_iters, + float hc_eps) { +#if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + if (!working || !post || !comb || !hc_state || !fn_tensor || !scale_data || !base_data || + !working->data || !post->data || !comb->data || !hc_state->data || !fn_tensor->data || + n_tokens <= 1) { + return false; + } + const bool can_use_device_params = + ds4_backend_is_gpu(backend) && + scale_tensor && base_tensor && + scale_tensor->data && base_tensor->data && + scale_tensor->buffer && base_tensor->buffer && + !ggml_backend_buffer_is_host(scale_tensor->buffer) && + !ggml_backend_buffer_is_host(base_tensor->buffer); + if (can_use_device_params) { + return deepseek4_cuda_hc_pre_batch_device_on_backend(backend, + hc_state->data, + fn_tensor->data, + scale_tensor->data, + base_tensor->data, + n_tokens, + n_embd, + n_hc, + sinkhorn_iters, + hc_eps, + working->data, + post->data, + comb->data, + pre ? pre->data : nullptr); + } + return deepseek4_cuda_hc_pre_batch_device_params_on_backend(backend, + hc_state->data, + fn_tensor->data, + scale_data, + base_data, + n_tokens, + n_embd, + n_hc, + sinkhorn_iters, + hc_eps, + working->data, + post->data, + comb->data, + pre ? pre->data : nullptr); +#else + (void) working; + (void) post; + (void) comb; + (void) pre; + (void) backend; + (void) hc_state; + (void) fn_tensor; + (void) scale_tensor; + (void) base_tensor; + (void) scale_data; + (void) base_data; + (void) n_tokens; + (void) n_embd; + (void) n_hc; + (void) sinkhorn_iters; + (void) hc_eps; + return false; +#endif +} + static bool build_cached_decode_hc_pre_graph( DeepSeek4CachedDecodeHcPreGraph & out, ggml_backend_t backend, @@ -1791,19 +2210,53 @@ static bool build_cached_decode_hc_pre_graph( comb = ds4_hc_col_normalize(out.sg.ctx, comb); } - ggml_tensor * hc_state_3d = ggml_reshape_3d(out.sg.ctx, out.sg.inp_embed, w.n_embd, w.n_hc, n_tokens); - ggml_tensor * hc_state_t = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, hc_state_3d, 1, 0, 2, 3)); - ggml_tensor * pre_3d = ggml_reshape_3d(out.sg.ctx, pre, w.n_hc, 1, n_tokens); - ggml_tensor * working = ggml_mul_mat(out.sg.ctx, hc_state_t, pre_3d); - working = ggml_reshape_2d(out.sg.ctx, working, w.n_embd, n_tokens); + ggml_tensor * working = nullptr; + if (n_tokens > 1) { + // Keep the batch weighted sum in flat 2D token-major views. HIP batch + // 3D matmul/reduction here corrupts the HC stream accumulation. + for (int h = 0; h < w.n_hc; ++h) { + ggml_tensor * hc_block = ggml_cont( + out.sg.ctx, + ggml_view_2d( + out.sg.ctx, + out.sg.inp_embed, + w.n_embd, + n_tokens, + out.sg.inp_embed->nb[1], + (size_t) h * (size_t) w.n_embd * out.sg.inp_embed->nb[0])); + ggml_tensor * pre_row = ggml_cont( + out.sg.ctx, + ggml_view_2d( + out.sg.ctx, + pre, + 1, + n_tokens, + pre->nb[1], + (size_t) h * pre->nb[0])); + ggml_tensor * contrib = ggml_mul( + out.sg.ctx, + hc_block, + ggml_repeat(out.sg.ctx, pre_row, hc_block)); + working = working ? ggml_add(out.sg.ctx, working, contrib) : contrib; + } + } else { + ggml_tensor * hc_state_3d = ggml_reshape_3d(out.sg.ctx, out.sg.inp_embed, w.n_embd, w.n_hc, n_tokens); + ggml_tensor * hc_state_t = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, hc_state_3d, 1, 0, 2, 3)); + ggml_tensor * pre_3d = ggml_reshape_3d(out.sg.ctx, pre, w.n_hc, 1, n_tokens); + working = ggml_mul_mat(out.sg.ctx, hc_state_t, pre_3d); + working = ggml_reshape_2d(out.sg.ctx, working, w.n_embd, n_tokens); + } out.sg.hidden_states = working; + out.pre = pre; out.post = post; out.comb = comb; ggml_set_output(out.sg.hidden_states); + ggml_set_output(out.pre); ggml_set_output(out.post); ggml_set_output(out.comb); ggml_build_forward_expand(out.sg.gf, out.sg.hidden_states); + ggml_build_forward_expand(out.sg.gf, out.pre); ggml_build_forward_expand(out.sg.gf, out.post); ggml_build_forward_expand(out.sg.gf, out.comb); @@ -2208,6 +2661,7 @@ static bool ds4_try_gpu_resident_decode_ffn( if (out_gpu) { *out_gpu = nullptr; } + (void) selected_weights_gpu; if (!enabled || !ffn_post_gpu) { return true; } @@ -2218,17 +2672,6 @@ static bool ds4_try_gpu_resident_decode_ffn( MoeHybridFfnTelemetry ffn_tel; const auto ffn_t0 = Ds4TimingClock::now(); - std::vector selected_weights_storage; - if (!selected_weights && selected_weights_gpu) { - selected_weights_storage.resize((size_t)n_selected); - ggml_backend_tensor_get(selected_weights_gpu, - selected_weights_storage.data(), 0, - sizeof(float) * selected_weights_storage.size()); - selected_weights = selected_weights_storage.data(); - } - if (!selected_weights) { - return false; - } if (!eval_moe_hybrid_ffn_gpu_resident( backend, hybrid_cfg, desc, storage, cpu_backend, ffn_post_gpu, @@ -2238,10 +2681,10 @@ static bool ds4_try_gpu_resident_decode_ffn( expert_compute, expert_layer)) { return false; } - if (out_gpu) { *out_gpu = gpu_ffn_state.state.act_cur; } + if (step_tel) { step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); add_ffn_telemetry(step_tel, ffn_tel); @@ -2395,6 +2838,8 @@ struct HcLayerWeightsCpu { struct HashRoutingTableCpu { std::vector ids; // [n_vocab, n_expert_used] + int width = 0; + int n_vocab = 0; bool loaded = false; }; @@ -2805,6 +3250,124 @@ static void hc_pre_batch(std::vector & working, }); } +static float ds4_max_abs_diff(const float * a, const float * b, size_t n) { + float max_diff = 0.0f; + for (size_t i = 0; i < n; ++i) { + max_diff = std::max(max_diff, std::fabs(a[i] - b[i])); + } + return max_diff; +} + +static void ds4_verify_backend_batch_hc_pre_graph( + const char * phase, + ggml_backend_t backend, + int kv_start, + int layer_idx, + int n_tokens, + ggml_tensor * hc_state_backend, + const std::vector & hc_state_host, + const HcWeightsCpu & weights, + ggml_tensor * fn_tensor, + const DeepSeek4CachedDecodeHcPreGraph & cached, + int n_embd, + int n_hc, + int sinkhorn_iters, + float hc_eps) { + if (!ds4_env_flag("DFLASH_DS4_VERIFY_BACKEND_BATCH_HC_PRE_GRAPH") || + n_tokens <= 1) { + return; + } + + const size_t hc_state_elems = + (size_t) n_tokens * (size_t) n_embd * (size_t) n_hc; + std::vector verify_hc_state(hc_state_backend ? hc_state_elems : 0); + const float * ref_hc_state = hc_state_host.data(); + if (hc_state_backend) { + ggml_backend_tensor_get(hc_state_backend, + verify_hc_state.data(), + 0, + sizeof(float) * hc_state_elems); + ref_hc_state = verify_hc_state.data(); + } + + const size_t hc_dim = (size_t) n_embd * (size_t) n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + std::vector ref_working((size_t) n_tokens * (size_t) n_embd); + std::vector ref_pre((size_t) n_tokens * (size_t) n_hc); + std::vector ref_post((size_t) n_tokens * (size_t) n_hc); + std::vector ref_comb((size_t) n_tokens * (size_t) n_hc * (size_t) n_hc); + ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + std::vector flat(hc_dim); + std::vector mix((size_t) mix_dim); + std::vector split((size_t) mix_dim); + for (int t = t0; t < t1; ++t) { + const float * token_hc = ref_hc_state + (size_t) t * hc_dim; + cpu_rms_norm(flat.data(), token_hc, (int) hc_dim, hc_eps); + cpu_matvec_f16_serial(mix.data(), + weights.fn_data.data(), + flat.data(), + mix_dim, + (int) hc_dim); + cpu_hc_sinkhorn(split.data(), + mix.data(), + weights.scale_data.data(), + weights.base_data.data(), + n_hc, + sinkhorn_iters, + 1.0e-6f); + std::memcpy(ref_pre.data() + (size_t) t * (size_t) n_hc, + split.data(), + (size_t) n_hc * sizeof(float)); + std::memcpy(ref_post.data() + (size_t) t * (size_t) n_hc, + split.data() + n_hc, + (size_t) n_hc * sizeof(float)); + std::memcpy(ref_comb.data() + (size_t) t * (size_t) n_hc * (size_t) n_hc, + split.data() + 2 * n_hc, + (size_t) n_hc * (size_t) n_hc * sizeof(float)); + for (int d = 0; d < n_embd; ++d) { + float acc = 0.0f; + for (int h = 0; h < n_hc; ++h) { + acc += split[(size_t) h] * token_hc[(size_t) h * (size_t) n_embd + (size_t) d]; + } + ref_working[(size_t) t * (size_t) n_embd + (size_t) d] = acc; + } + } + }); + + std::vector graph_working(ref_working.size()); + std::vector graph_pre(ref_pre.size()); + std::vector graph_post(ref_post.size()); + std::vector graph_comb(ref_comb.size()); + ggml_backend_tensor_get(cached.sg.hidden_states, + graph_working.data(), + 0, + sizeof(float) * graph_working.size()); + ggml_backend_tensor_get(cached.pre, + graph_pre.data(), + 0, + sizeof(float) * graph_pre.size()); + ggml_backend_tensor_get(cached.post, + graph_post.data(), + 0, + sizeof(float) * graph_post.size()); + ggml_backend_tensor_get(cached.comb, + graph_comb.data(), + 0, + sizeof(float) * graph_comb.size()); + + std::fprintf(stderr, + "[deepseek4-hc-verify] kv=%d layer=%d phase=%s tokens=%d diff_pre=%.6g diff_working=%.6g diff_post=%.6g diff_comb=%.6g\n", + kv_start, + layer_idx, + phase ? phase : "?", + n_tokens, + ds4_max_abs_diff(graph_pre.data(), ref_pre.data(), ref_pre.size()), + ds4_max_abs_diff(graph_working.data(), ref_working.data(), ref_working.size()), + ds4_max_abs_diff(graph_post.data(), ref_post.data(), ref_post.size()), + ds4_max_abs_diff(graph_comb.data(), ref_comb.data(), ref_comb.size())); + std::fflush(stderr); +} + static void cpu_hc_post(float * out_hc, const float * block_out, const float * residual_hc, const float * post, const float * comb, int n_embd, int n_hc) { @@ -2891,6 +3454,11 @@ static void load_hc_weights_cpu(HcWeightsCpu & dst, ggml_tensor * fn, static bool load_hash_routing_cpu(HashRoutingTableCpu & dst, ggml_tensor * table) { if (dst.loaded) return true; if (!table) return false; + dst.width = (int)table->ne[0]; + dst.n_vocab = table->ne[1] > 0 ? (int)table->ne[1] : 0; + if (dst.width <= 0 || dst.n_vocab <= 0) { + return false; + } dst.ids.resize(ggml_nelements(table)); ggml_backend_tensor_get(table, dst.ids.data(), 0, ggml_nbytes(table)); dst.loaded = true; @@ -2962,25 +3530,27 @@ static bool deepseek4_step_hybrid( // policy is narrowed. const bool use_cached_decode = n_tokens == 1 && ds4_backend_is_gpu(backend) && !disable_cached_decode; - // Batch HC graph is opt-in: strict PR489 validation showed the eager path - // is quality-clean and slightly faster on the target HIP host. const bool use_backend_batch_hc_graph = n_tokens > 1 && ds4_backend_is_gpu(backend) && - ds4_env_flag("DFLASH_DS4_ENABLE_BACKEND_BATCH_HC_GRAPH") && !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_BATCH_HC_GRAPH"); + const bool use_backend_decode_hc = + use_cached_decode && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC") && + (!ds4_backend_is_hip(backend) || + ds4_env_flag("DFLASH_DS4_ENABLE_BACKEND_DECODE_HC")); const bool disable_backend_decode_hc_direct = ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC_DIRECT"); const bool use_backend_decode_hc_direct = - use_cached_decode && + use_backend_decode_hc && ds4_backend_is_hip(backend) && !disable_backend_decode_hc_direct; const bool use_backend_decode_hc_graph = - use_cached_decode && !use_backend_decode_hc_direct; + use_backend_decode_hc && !use_backend_decode_hc_direct; const bool use_backend_hc_graph = use_backend_batch_hc_graph || use_backend_decode_hc_graph; const bool use_backend_hc_post_graph = - use_backend_batch_hc_graph || use_cached_decode; + use_backend_batch_hc_graph || use_backend_decode_hc; const bool use_gpu_resident_decode_ffn = use_cached_decode && ds4_env_flag("DFLASH_DS4_GPU_RESIDENT_FFN") && @@ -3004,6 +3574,7 @@ static bool deepseek4_step_hybrid( static std::vector cached_decode_route_graphs; static thread_local DeepSeek4CachedDecodeGpuFfnState cached_decode_gpu_ffn_state; static DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; + static Ds4DecodeSharedInputs decode_shared_inputs; static int decode_cache_n_layer = 0; static int batch_hc_cache_n_layer = 0; static int dynamic_attn_alloc_n_layer = 0; @@ -3029,6 +3600,7 @@ static bool deepseek4_step_hybrid( } cached_decode_route_graphs.assign((size_t)w.n_layer, {}); cached_decode_output_graph.free(); + decode_shared_inputs.free(); decode_cache_n_layer = w.n_layer; route_cache_n_layer = w.n_layer; } @@ -3058,6 +3630,15 @@ static bool deepseek4_step_hybrid( cached_decode_route_graphs.assign((size_t)w.n_layer, {}); route_cache_n_layer = w.n_layer; } + const bool use_decode_shared_inputs = + use_cached_decode && + !ds4_env_flag("DFLASH_DS4_DISABLE_DECODE_SHARED_INPUTS") && + decode_shared_inputs.ensure(w, backend); + if (use_decode_shared_inputs) { + decode_shared_inputs.set_step(w, kv_start); + } + const Ds4DecodeSharedInputs * decode_shared_inputs_ptr = + use_decode_shared_inputs ? &decode_shared_inputs : nullptr; ggml_tensor * hc_state_backend = nullptr; DeepSeek4CachedDecodeHcPostGraph * hc_post_graph = nullptr; @@ -3231,7 +3812,8 @@ static bool deepseek4_step_hybrid( candidate.n_raw == n_raw && candidate.n_comp_attn == n_comp_attn && candidate.n_index_comp == n_index_comp && - candidate.flush_boundary == flush_boundary; + candidate.flush_boundary == flush_boundary && + candidate.uses_shared_inputs == (decode_shared_inputs_ptr != nullptr); }); if (it == per_layer.end()) { if (per_layer.size() >= 20) { @@ -3243,7 +3825,8 @@ static bool deepseek4_step_hybrid( const auto attn_build_t0 = Ds4TimingClock::now(); if (!build_cached_decode_attn_graph(candidate, backend, w, L, lc, il, kv_start, n_raw, n_comp_attn, n_index_comp, - flush_boundary)) { + flush_boundary, + decode_shared_inputs_ptr)) { if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; @@ -3255,41 +3838,43 @@ static bool deepseek4_step_hybrid( gf = cached_attn->sg.gf; attn_out = cached_attn->sg.hidden_states; - const int64_t raw_row = kv_start % w.n_swa; - const int32_t rope_pos = kv_start; - const int32_t neg_pos = -kv_start; if (attn_in_backend) { ggml_backend_tensor_copy(attn_in_backend, cached_attn->sg.inp_embed); } else { ggml_backend_tensor_set(cached_attn->sg.inp_embed, cur.data(), 0, sizeof(float) * cur.size()); } - ggml_backend_tensor_set(cached_attn->inputs.rope_pos, &rope_pos, 0, sizeof(rope_pos)); - ggml_backend_tensor_set(cached_attn->inputs.neg_pos, &neg_pos, 0, sizeof(neg_pos)); - ggml_backend_tensor_set(cached_attn->inputs.raw_kv_rows, &raw_row, 0, sizeof(raw_row)); - if (ratio > 0) { - const int pos_mod = token_pos % ratio; - const int32_t ape_row = 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; - ggml_backend_tensor_set(cached_attn->inputs.attn_ape_row, &ape_row, 0, sizeof(ape_row)); - ggml_backend_tensor_set(cached_attn->inputs.attn_state_rows, &state_row, 0, sizeof(state_row)); - if (flush_boundary) { - ggml_backend_tensor_set(cached_attn->inputs.attn_comp_rows, &comp_row, 0, sizeof(comp_row)); - ggml_backend_tensor_set(cached_attn->inputs.attn_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + if (!cached_attn->uses_shared_inputs) { + const int64_t raw_row = kv_start % w.n_swa; + const int32_t rope_pos = kv_start; + const int32_t neg_pos = -kv_start; + ggml_backend_tensor_set(cached_attn->inputs.rope_pos, &rope_pos, 0, sizeof(rope_pos)); + ggml_backend_tensor_set(cached_attn->inputs.neg_pos, &neg_pos, 0, sizeof(neg_pos)); + ggml_backend_tensor_set(cached_attn->inputs.raw_kv_rows, &raw_row, 0, sizeof(raw_row)); + if (ratio > 0) { + const int pos_mod = token_pos % ratio; + const int32_t ape_row = 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; + ggml_backend_tensor_set(cached_attn->inputs.attn_ape_row, &ape_row, 0, sizeof(ape_row)); + ggml_backend_tensor_set(cached_attn->inputs.attn_state_rows, &state_row, 0, sizeof(state_row)); + if (flush_boundary) { + ggml_backend_tensor_set(cached_attn->inputs.attn_comp_rows, &comp_row, 0, sizeof(comp_row)); + ggml_backend_tensor_set(cached_attn->inputs.attn_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + } } - } - if (ratio == 4) { - const int pos_mod = token_pos % ratio; - const int32_t ape_row = 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; - ggml_backend_tensor_set(cached_attn->inputs.index_ape_row, &ape_row, 0, sizeof(ape_row)); - ggml_backend_tensor_set(cached_attn->inputs.index_state_rows, &state_row, 0, sizeof(state_row)); - if (flush_boundary) { - ggml_backend_tensor_set(cached_attn->inputs.index_comp_rows, &comp_row, 0, sizeof(comp_row)); - ggml_backend_tensor_set(cached_attn->inputs.index_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + if (ratio == 4) { + const int pos_mod = token_pos % ratio; + const int32_t ape_row = 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; + ggml_backend_tensor_set(cached_attn->inputs.index_ape_row, &ape_row, 0, sizeof(ape_row)); + ggml_backend_tensor_set(cached_attn->inputs.index_state_rows, &state_row, 0, sizeof(state_row)); + if (flush_boundary) { + ggml_backend_tensor_set(cached_attn->inputs.index_comp_rows, &comp_row, 0, sizeof(comp_row)); + ggml_backend_tensor_set(cached_attn->inputs.index_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + } } } } else { @@ -3573,16 +4158,18 @@ static bool deepseek4_step_hybrid( } const auto route_select_t0 = Ds4TimingClock::now(); - const auto & hash_table = hash_routing_tables[(size_t)il].ids; + const auto & hash_table = hash_routing_tables[(size_t)il]; for (int ti = 0; ti < n_tokens; ++ti) { const int32_t tok = token_ids[ti]; - if (tok < 0 || tok >= w.n_vocab) { + if (tok < 0 || tok >= hash_table.n_vocab || + hash_table.width < w.n_expert_used) { std::fprintf(stderr, "[deepseek4] token id %d outside hash table for layer %d\n", tok, il); if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; } - const int32_t * row = hash_table.data() + (size_t)tok * (size_t)w.n_expert_used; + const int32_t * row = + hash_table.ids.data() + (size_t)tok * (size_t)hash_table.width; std::memcpy(hybrid_selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, row, sizeof(int32_t) * (size_t)w.n_expert_used); @@ -3614,7 +4201,10 @@ static bool deepseek4_step_hybrid( expert_layers ? expert_layers + il : nullptr; const bool all_selected_hot = ds4_all_selected_hot( storage, hybrid_selected_host.data(), n_tokens * w.n_expert_used); - const bool can_skip_ffn_normed_host = false; + const bool can_skip_ffn_normed_host = + expert_worker == nullptr && + (all_selected_hot || + storage.cold_backend_kind == MoeHybridColdBackend::Gpu); const auto route_read_t0 = Ds4TimingClock::now(); if (!use_gpu_resident_decode_ffn && !can_skip_ffn_normed_host) { const Ds4TensorReadback readbacks[] = { @@ -3724,7 +4314,9 @@ static bool deepseek4_step_hybrid( auto & storage = moe_hybrid.layers[(size_t) il]; const MoeExpertLayer * expert_layer = expert_layers ? expert_layers + il : nullptr; - const bool can_skip_ffn_normed_host = false; + const bool can_skip_ffn_normed_host = + expert_worker == nullptr && + storage.cold_backend_kind == MoeHybridColdBackend::Gpu; if (ok) { const auto route_read_t0 = Ds4TimingClock::now(); if (can_skip_ffn_normed_host) { @@ -3755,10 +4347,11 @@ static bool deepseek4_step_hybrid( } const auto route_select_t0 = Ds4TimingClock::now(); - const auto & hash_table = hash_routing_tables[(size_t)il].ids; + const auto & hash_table = hash_routing_tables[(size_t)il]; for (int ti = 0; ti < n_tokens; ++ti) { const int32_t tok = token_ids[ti]; - if (tok < 0 || tok >= w.n_vocab) { + if (tok < 0 || tok >= hash_table.n_vocab || + hash_table.width < w.n_expert_used) { std::fprintf(stderr, "[deepseek4] token id %d outside hash table for layer %d\n", tok, il); if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); @@ -3768,7 +4361,8 @@ static bool deepseek4_step_hybrid( ds4_parallel_for_tokens(n_tokens, 16, [&](int t0, int t1) { for (int ti = t0; ti < t1; ++ti) { const int32_t tok = token_ids[ti]; - const int32_t * row = hash_table.data() + (size_t)tok * (size_t)w.n_expert_used; + const int32_t * row = + hash_table.ids.data() + (size_t)tok * (size_t)hash_table.width; float sum = 0.0f; for (int ei = 0; ei < w.n_expert_used; ++ei) { const int32_t expert = row[ei]; @@ -3890,7 +4484,10 @@ static bool deepseek4_step_hybrid( } const bool all_selected_hot = ds4_all_selected_hot( storage, hybrid_selected_host.data(), n_tokens * w.n_expert_used); - const bool can_skip_ffn_normed_host = false; + const bool can_skip_ffn_normed_host = + expert_worker == nullptr && + (all_selected_hot || + storage.cold_backend_kind == MoeHybridColdBackend::Gpu); if (!prefer_single_hip_route_readback && !use_gpu_resident_decode_ffn && !can_skip_ffn_normed_host) { @@ -3992,7 +4589,9 @@ static bool deepseek4_step_hybrid( auto & storage = moe_hybrid.layers[(size_t) il]; const MoeExpertLayer * expert_layer = expert_layers ? expert_layers + il : nullptr; - const bool can_skip_ffn_normed_host = false; + const bool can_skip_ffn_normed_host = + expert_worker == nullptr && + storage.cold_backend_kind == MoeHybridColdBackend::Gpu; const auto route_read_t0 = Ds4TimingClock::now(); if (can_skip_ffn_normed_host) { const Ds4TensorReadback route_readbacks[] = { @@ -4136,7 +4735,7 @@ static bool deepseek4_step_hybrid( const bool use_output_hc = w.output_hc_fn && w.output_hc_scale && w.output_hc_base; - if ((!use_cached_decode || !use_output_hc) && + if ((!use_cached_decode || !use_output_hc || !hc_state_backend) && use_backend_hc_post_graph && hc_state_backend) { ggml_backend_tensor_get(hc_state_backend, hc_state.data(), 0, sizeof(float) * hc_state.size()); @@ -4144,7 +4743,7 @@ static bool deepseek4_step_hybrid( // ── Output HC pre → norm → logits ─────────────────────────────────── const auto output_t0 = Ds4TimingClock::now(); - if (!use_cached_decode || !use_output_hc) { + if (!use_cached_decode || !use_output_hc || !hc_state_backend) { if (hc_output_weights.loaded) { hc_output_batch(final_embd, hc_state.data(), hc_output_weights, n_tokens, n_embd, n_hc, w.hc_eps); @@ -4687,6 +5286,31 @@ bool deepseek4_step( routing_stats, expert_compute, expert_layers); } +bool deepseek4_step( + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const float * embed, + int n_tokens, + int kv_start, + std::vector & out_logits, + MoeHybridStorage * moe_hybrid, + const int32_t * token_ids, + ExpertIpcClient * expert_worker, + bool worker_owns_hot_ids, + bool want_logits, + bool disable_cached_decode, + DeepSeek4StepTelemetry * telemetry, + MoeHybridRoutingStats * routing_stats, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layers) { + (void) want_logits; + return deepseek4_step(backend, w, cache, embed, n_tokens, kv_start, + out_logits, moe_hybrid, token_ids, expert_worker, + worker_owns_hot_ids, disable_cached_decode, telemetry, + routing_stats, expert_compute, expert_layers); +} + bool deepseek4_step_layer_range( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -4739,13 +5363,78 @@ bool deepseek4_step_layer_range( static std::vector cached_decode_attn_hc_pre_graphs; static std::vector cached_decode_ffn_hc_pre_graphs; static DeepSeek4CachedDecodeHcPostGraph cached_decode_hc_post_graph; + static std::vector cached_batch_attn_hc_pre_graphs; + static std::vector cached_batch_ffn_hc_pre_graphs; + static DeepSeek4CachedDecodeHcPostGraph cached_batch_hc_post_graph; static std::vector> cached_decode_attn_graphs; static std::vector cached_decode_route_graphs_range; static std::vector cached_decode_expert_graphs_range; static std::vector cached_decode_ffn_graphs; + static std::vector cached_prefill_ffn_graphs; static DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; static DeepSeek4CachedLayerAlloc cached_dynamic_output_alloc; + static Ds4DecodeSharedInputs decode_shared_inputs; static int hc_loaded_n_layer = 0; + static int batch_hc_cache_n_layer = 0; + static uint64_t prefill_ffn_cache_epoch = 0; + static bool prefill_ffn_cache_decode_evict_done = true; + static bool prefill_ffn_cache_decode_disable_hc_post_direct_ffn = false; + static thread_local uint64_t cached_request_epoch = 0; + const uint64_t request_epoch = ds4_current_request_graph_epoch(); + if (cached_request_epoch != request_epoch && hc_loaded_n_layer == w.n_layer) { + for (auto & alloc : cached_attn_allocs) { + alloc.free(); + } + cached_attn_allocs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_attn_hc_pre_graphs) { + g.free(); + } + cached_decode_attn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_ffn_hc_pre_graphs) { + g.free(); + } + cached_decode_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + cached_decode_hc_post_graph.free(); + for (auto & g : cached_batch_attn_hc_pre_graphs) { + g.free(); + } + cached_batch_attn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_batch_ffn_hc_pre_graphs) { + g.free(); + } + cached_batch_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + cached_batch_hc_post_graph.free(); + for (auto & per_layer : cached_decode_attn_graphs) { + for (auto & g : per_layer) { + g.free(); + } + } + cached_decode_attn_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_route_graphs_range) { + g.free(); + } + cached_decode_route_graphs_range.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_expert_graphs_range) { + g.free(); + } + cached_decode_expert_graphs_range.assign((size_t)w.n_layer, {}); + for (auto & g : cached_decode_ffn_graphs) { + g.free(); + } + cached_decode_ffn_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_prefill_ffn_graphs) { + g.free(); + } + cached_prefill_ffn_graphs.assign((size_t)w.n_layer, {}); + cached_decode_output_graph.free(); + cached_dynamic_output_alloc.free(); + decode_shared_inputs.free(); + batch_hc_cache_n_layer = 0; + prefill_ffn_cache_epoch = 0; + prefill_ffn_cache_decode_evict_done = true; + prefill_ffn_cache_decode_disable_hc_post_direct_ffn = false; + cached_request_epoch = request_epoch; + } if (hc_loaded_n_layer != w.n_layer) { hc_layer_weights_range.resize((size_t)w.n_layer); hash_routing_tables_range.resize((size_t)w.n_layer); @@ -4762,6 +5451,16 @@ bool deepseek4_step_layer_range( } cached_decode_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); cached_decode_hc_post_graph.free(); + for (auto & g : cached_batch_attn_hc_pre_graphs) { + g.free(); + } + cached_batch_attn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_batch_ffn_hc_pre_graphs) { + g.free(); + } + cached_batch_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + cached_batch_hc_post_graph.free(); + batch_hc_cache_n_layer = w.n_layer; for (auto & per_layer : cached_decode_attn_graphs) { for (auto & g : per_layer) { g.free(); @@ -4780,8 +5479,16 @@ bool deepseek4_step_layer_range( g.free(); } cached_decode_ffn_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_prefill_ffn_graphs) { + g.free(); + } + cached_prefill_ffn_graphs.assign((size_t)w.n_layer, {}); cached_decode_output_graph.free(); cached_dynamic_output_alloc.free(); + decode_shared_inputs.free(); + prefill_ffn_cache_epoch = 0; + prefill_ffn_cache_decode_evict_done = true; + prefill_ffn_cache_decode_disable_hc_post_direct_ffn = false; for (int il = 0; il < w.n_layer; il++) { const DeepSeek4Layer & L = w.layers[(size_t)il]; load_hc_weights_cpu(hc_layer_weights_range[il].attn, L.hc_attn_fn, L.hc_attn_scale, L.hc_attn_base); @@ -4806,28 +5513,141 @@ 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_backend_decode_hc = reuse_decode_graphs && ds4_backend_is_gpu(backend); + const bool use_decode_shared_inputs = + reuse_decode_graphs && + ds4_backend_is_gpu(backend) && + !ds4_env_flag("DFLASH_DS4_DISABLE_DECODE_SHARED_INPUTS") && + decode_shared_inputs.ensure(w, backend); + if (use_decode_shared_inputs) { + decode_shared_inputs.set_step(w, kv_start); + } + const Ds4DecodeSharedInputs * decode_shared_inputs_ptr = + use_decode_shared_inputs ? &decode_shared_inputs : nullptr; + const bool split_layer_range_route_expert = + ds4_env_flag("DFLASH_DS4_SPLIT_LAYER_RANGE_ROUTE_EXPERT"); + const bool cache_layer_range_prefill_ffn = + n_tokens > 1 && + !ds4_env_flag("DFLASH_DS4_DISABLE_LAYER_RANGE_PREFILL_FFN_CACHE"); + if (cache_layer_range_prefill_ffn) { + prefill_ffn_cache_decode_evict_done = false; + prefill_ffn_cache_decode_disable_hc_post_direct_ffn = false; + } else if (reuse_decode_graphs && !prefill_ffn_cache_decode_evict_done) { + const bool cache_pressure = + ds4_maybe_evict_prefill_ffn_cache_for_decode( + cached_prefill_ffn_graphs, w); + prefill_ffn_cache_decode_disable_hc_post_direct_ffn = + cache_pressure && + !ds4_env_flag("DFLASH_DS4_DISABLE_ADAPTIVE_HC_POST_DIRECT_GUARD"); + if (prefill_ffn_cache_decode_disable_hc_post_direct_ffn && + ds4_env_flag("DFLASH_DS4_PREFILL_FFN_CACHE_TRACE")) { + std::fprintf(stderr, + "[deepseek4] prefill-ffn-cache adaptive guard disabled ffn hc-post-direct for decode\n"); + std::fflush(stderr); + } + prefill_ffn_cache_decode_evict_done = true; + } + const bool enable_backend_batch_hc_graph = + n_tokens > 1 && + ds4_backend_is_gpu(backend) && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_BATCH_HC_GRAPH"); + const bool use_backend_batch_hc_pre_graph = + enable_backend_batch_hc_graph && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_BATCH_HC_PRE_GRAPH"); + const bool use_backend_batch_hc_pre_direct = + use_backend_batch_hc_pre_graph && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_BATCH_HC_PRE_DIRECT"); + const bool use_backend_batch_hc_post_graph = + enable_backend_batch_hc_graph && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_BATCH_HC_POST_GRAPH"); + const bool use_backend_batch_hc_state_graph = + use_backend_batch_hc_pre_graph || use_backend_batch_hc_post_graph; + const bool use_backend_decode_hc = + reuse_decode_graphs && + ds4_backend_is_gpu(backend) && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC"); const bool disable_backend_decode_hc_direct = ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC_DIRECT"); const bool use_backend_decode_hc_direct = use_backend_decode_hc && ds4_backend_is_hip(backend) && !disable_backend_decode_hc_direct; + const bool disable_backend_decode_hc_post_direct = + ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC_POST_DIRECT"); + const bool use_backend_decode_hc_post_direct_attn = + use_backend_decode_hc_direct && + !disable_backend_decode_hc_post_direct && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC_POST_DIRECT_ATTN"); + const bool use_backend_decode_hc_post_direct_ffn = + use_backend_decode_hc_direct && + !disable_backend_decode_hc_post_direct && + !prefill_ffn_cache_decode_disable_hc_post_direct_ffn && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC_POST_DIRECT_FFN"); + const bool use_backend_decode_hc_post_direct = + use_backend_decode_hc_post_direct_attn || + use_backend_decode_hc_post_direct_ffn; + const bool sync_backend_decode_hc_post_direct = + use_backend_decode_hc_post_direct && + !ds4_env_flag("DFLASH_DS4_DISABLE_BACKEND_DECODE_HC_POST_DIRECT_SYNC"); const bool use_backend_decode_hc_graph = use_backend_decode_hc && !use_backend_decode_hc_direct; - 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() || - cached_decode_hc_post_graph.owner_ctx != w.ctx || - cached_decode_hc_post_graph.backend != backend) { - if (!build_cached_decode_hc_post_graph(cached_decode_hc_post_graph, backend, w)) { + const bool use_backend_decode_hc_pre_graph = use_backend_decode_hc_graph; + const bool use_backend_decode_hc_post_graph = use_backend_decode_hc_graph; + const bool use_backend_hc_pre_graph = + use_backend_batch_hc_pre_graph || use_backend_decode_hc_pre_graph; + const bool use_backend_hc_post_graph = + use_backend_batch_hc_post_graph || use_backend_decode_hc_post_graph; + if (use_backend_batch_hc_state_graph && batch_hc_cache_n_layer != w.n_layer) { + for (auto & g : cached_batch_attn_hc_pre_graphs) { + g.free(); + } + cached_batch_attn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + for (auto & g : cached_batch_ffn_hc_pre_graphs) { + g.free(); + } + cached_batch_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); + cached_batch_hc_post_graph.free(); + batch_hc_cache_n_layer = w.n_layer; + } + DeepSeek4CachedDecodeHcPostGraph * hc_post_graph = nullptr; + ggml_tensor * hc_state_backend = nullptr; + bool hc_state_backend_current = false; + bool hc_state_host_current = true; + const bool need_backend_hc_state = + use_backend_hc_pre_graph || use_backend_hc_post_graph || use_backend_decode_hc_direct; + if (need_backend_hc_state) { + hc_post_graph = use_backend_batch_hc_state_graph + ? &cached_batch_hc_post_graph + : &cached_decode_hc_post_graph; + if (!hc_post_graph->valid() || + hc_post_graph->owner_ctx != w.ctx || + hc_post_graph->backend != backend || + hc_post_graph->n_tokens != n_tokens) { + if (!build_cached_decode_hc_post_graph(*hc_post_graph, backend, w, n_tokens)) { return false; } } - ggml_backend_tensor_set(cached_decode_hc_post_graph.residual_hc, + ggml_backend_tensor_set(hc_post_graph->residual_hc, hc_state.data(), 0, sizeof(float) * hc_state.size()); - hc_state_backend = cached_decode_hc_post_graph.residual_hc; + hc_state_backend = hc_post_graph->residual_hc; + hc_state_backend_current = true; } + auto sync_hc_state_from_backend_for_host_pre = [&]() { + if (!hc_state_host_current && hc_state_backend_current && hc_state_backend) { + ggml_backend_tensor_get(hc_state_backend, + hc_state.data(), 0, + sizeof(float) * hc_state.size()); + hc_state_host_current = true; + } + }; + auto sync_hc_state_to_backend_for_graph_pre = [&]() { + if (!use_backend_hc_post_graph && use_backend_hc_pre_graph && hc_post_graph) { + ggml_backend_tensor_set(hc_post_graph->residual_hc, + hc_state.data(), 0, + sizeof(float) * hc_state.size()); + hc_state_backend = hc_post_graph->residual_hc; + hc_state_backend_current = true; + } + }; for (int il = layer_begin; il < layer_end; ++il) { const bool trace_decode = ds4_env_flag("DFLASH_DS4_TRACE_DECODE") && n_tokens == 1; @@ -4852,8 +5672,15 @@ bool deepseek4_step_layer_range( const auto hc_pre_attn_t0 = Ds4TimingClock::now(); 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()); + ggml_tensor * hc_pre_input = hc_post_graph->residual_hc; + if (hc_state_backend_current && hc_state_backend) { + hc_pre_input = hc_state_backend; + } else { + ggml_backend_tensor_set(hc_post_graph->residual_hc, + hc_state.data(), 0, sizeof(float) * hc_state.size()); + hc_state_backend = hc_post_graph->residual_hc; + hc_state_backend_current = true; + } if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_attn_input_t0, Ds4TimingClock::now()); auto & cached = cached_decode_attn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || @@ -4875,7 +5702,7 @@ bool deepseek4_step_layer_range( backend, il, false, - cached_decode_hc_post_graph.residual_hc, + hc_pre_input, L.hc_attn_fn, L.hc_attn_scale, L.hc_attn_base, @@ -4892,33 +5719,76 @@ bool deepseek4_step_layer_range( attn_in_backend = cached.sg.hidden_states; attn_post_backend = cached.post; attn_comb_backend = cached.comb; - } else if (use_backend_decode_hc_graph) { - auto & cached = cached_decode_attn_hc_pre_graphs[(size_t)il]; + } else if (use_backend_hc_pre_graph) { + auto & cached = use_backend_batch_hc_pre_graph + ? cached_batch_attn_hc_pre_graphs[(size_t)il] + : cached_decode_attn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || cached.owner_ctx != w.ctx || cached.backend != backend || cached.layer_idx != il || + cached.n_tokens != n_tokens || cached.ffn) { const auto hc_pre_attn_build_t0 = Ds4TimingClock::now(); - if (!build_cached_decode_hc_pre_graph(cached, backend, w, L, hc_lw.attn.scale_data.data(), il, false)) { + if (!build_cached_decode_hc_pre_graph(cached, backend, w, L, hc_lw.attn.scale_data.data(), il, false, n_tokens)) { std::fprintf(stderr, "[deepseek4] cached hc-pre graph alloc failed layer %d attn\n", il); return false; } if (telemetry) telemetry->hc_pre_build_us += ds4_elapsed_us(hc_pre_attn_build_t0, Ds4TimingClock::now()); } const auto hc_pre_attn_input_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_copy(hc_state_backend, cached.sg.inp_embed); + sync_hc_state_to_backend_for_graph_pre(); + if (!use_backend_batch_hc_pre_direct) { + ggml_backend_tensor_copy(hc_state_backend, cached.sg.inp_embed); + } if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_attn_input_t0, Ds4TimingClock::now()); const auto hc_pre_attn_compute_t0 = Ds4TimingClock::now(); - if (ggml_backend_graph_compute(backend, cached.sg.gf) != GGML_STATUS_SUCCESS) { + if (use_backend_batch_hc_pre_direct) { + if (!ds4_try_gpu_hc_pre_batch_device(cached.sg.hidden_states, + cached.post, + cached.comb, + cached.pre, + backend, + hc_state_backend, + L.hc_attn_fn, + L.hc_attn_scale, + L.hc_attn_base, + hc_lw.attn.scale_data.data(), + hc_lw.attn.base_data.data(), + n_tokens, + n_embd, + n_hc, + w.n_hc_sinkhorn_iter, + w.hc_eps)) { + std::fprintf(stderr, "[deepseek4] direct batch hc-pre compute failed layer %d attn\n", il); + return false; + } + } else if (ggml_backend_graph_compute(backend, cached.sg.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] cached hc-pre compute 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()); + if (use_backend_batch_hc_pre_graph) { + ds4_verify_backend_batch_hc_pre_graph("attn", + backend, + kv_start, + il, + n_tokens, + hc_state_backend, + hc_state, + hc_lw.attn, + L.hc_attn_fn, + cached, + n_embd, + n_hc, + w.n_hc_sinkhorn_iter, + w.hc_eps); + } attn_in_backend = cached.sg.hidden_states; attn_post_backend = cached.post; attn_comb_backend = cached.comb; } else { + sync_hc_state_from_backend_for_host_pre(); hc_pre_batch(cur, hc_post, hc_comb, hc_state.data(), hc_lw.attn, L.hc_attn_fn, n_tokens, n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); @@ -4956,7 +5826,8 @@ bool deepseek4_step_layer_range( candidate.n_raw == n_raw && candidate.n_comp_attn == n_comp_attn && candidate.n_index_comp == n_index_comp && - candidate.flush_boundary == flush_boundary; + candidate.flush_boundary == flush_boundary && + candidate.uses_shared_inputs == (decode_shared_inputs_ptr != nullptr); }); if (it == per_layer.end()) { if (per_layer.size() >= 20) { @@ -4968,7 +5839,8 @@ bool deepseek4_step_layer_range( const auto attn_build_t0 = Ds4TimingClock::now(); if (!build_cached_decode_attn_graph(candidate, backend, w, L, lc, il, kv_start, n_raw, n_comp_attn, n_index_comp, - flush_boundary)) { + flush_boundary, + decode_shared_inputs_ptr)) { std::fprintf(stderr, "[deepseek4] cached attn graph alloc failed layer %d\n", il); return false; } @@ -4979,41 +5851,43 @@ bool deepseek4_step_layer_range( gf = cached_attn->sg.gf; attn_out = cached_attn->sg.hidden_states; - const int64_t raw_row = kv_start % w.n_swa; - const int32_t rope_pos = kv_start; - const int32_t neg_pos = -kv_start; if (attn_in_backend) { ggml_backend_tensor_copy(attn_in_backend, cached_attn->sg.inp_embed); } else { ggml_backend_tensor_set(cached_attn->sg.inp_embed, cur.data(), 0, sizeof(float) * cur.size()); } - ggml_backend_tensor_set(cached_attn->inputs.rope_pos, &rope_pos, 0, sizeof(rope_pos)); - ggml_backend_tensor_set(cached_attn->inputs.neg_pos, &neg_pos, 0, sizeof(neg_pos)); - ggml_backend_tensor_set(cached_attn->inputs.raw_kv_rows, &raw_row, 0, sizeof(raw_row)); - if (ratio > 0) { - const int pos_mod = token_pos % ratio; - const int32_t ape_row = 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; - ggml_backend_tensor_set(cached_attn->inputs.attn_ape_row, &ape_row, 0, sizeof(ape_row)); - ggml_backend_tensor_set(cached_attn->inputs.attn_state_rows, &state_row, 0, sizeof(state_row)); - if (flush_boundary) { - ggml_backend_tensor_set(cached_attn->inputs.attn_comp_rows, &comp_row, 0, sizeof(comp_row)); - ggml_backend_tensor_set(cached_attn->inputs.attn_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + if (!cached_attn->uses_shared_inputs) { + const int64_t raw_row = kv_start % w.n_swa; + const int32_t rope_pos = kv_start; + const int32_t neg_pos = -kv_start; + ggml_backend_tensor_set(cached_attn->inputs.rope_pos, &rope_pos, 0, sizeof(rope_pos)); + ggml_backend_tensor_set(cached_attn->inputs.neg_pos, &neg_pos, 0, sizeof(neg_pos)); + ggml_backend_tensor_set(cached_attn->inputs.raw_kv_rows, &raw_row, 0, sizeof(raw_row)); + if (ratio > 0) { + const int pos_mod = token_pos % ratio; + const int32_t ape_row = 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; + ggml_backend_tensor_set(cached_attn->inputs.attn_ape_row, &ape_row, 0, sizeof(ape_row)); + ggml_backend_tensor_set(cached_attn->inputs.attn_state_rows, &state_row, 0, sizeof(state_row)); + if (flush_boundary) { + ggml_backend_tensor_set(cached_attn->inputs.attn_comp_rows, &comp_row, 0, sizeof(comp_row)); + ggml_backend_tensor_set(cached_attn->inputs.attn_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + } } - } - if (ratio == 4) { - const int pos_mod = token_pos % ratio; - const int32_t ape_row = 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; - ggml_backend_tensor_set(cached_attn->inputs.index_ape_row, &ape_row, 0, sizeof(ape_row)); - ggml_backend_tensor_set(cached_attn->inputs.index_state_rows, &state_row, 0, sizeof(state_row)); - if (flush_boundary) { - ggml_backend_tensor_set(cached_attn->inputs.index_comp_rows, &comp_row, 0, sizeof(comp_row)); - ggml_backend_tensor_set(cached_attn->inputs.index_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + if (ratio == 4) { + const int pos_mod = token_pos % ratio; + const int32_t ape_row = 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; + ggml_backend_tensor_set(cached_attn->inputs.index_ape_row, &ape_row, 0, sizeof(ape_row)); + ggml_backend_tensor_set(cached_attn->inputs.index_state_rows, &state_row, 0, sizeof(state_row)); + if (flush_boundary) { + ggml_backend_tensor_set(cached_attn->inputs.index_comp_rows, &comp_row, 0, sizeof(comp_row)); + ggml_backend_tensor_set(cached_attn->inputs.index_comp_pos, &comp_pos, 0, sizeof(comp_pos)); + } } } } else { @@ -5075,24 +5949,60 @@ 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 (hc_state_backend != cached_decode_hc_post_graph.residual_hc) { - ggml_backend_tensor_copy(hc_state_backend, cached_decode_hc_post_graph.residual_hc); + if (use_backend_hc_post_graph) { + if (hc_state_backend != hc_post_graph->residual_hc) { + ggml_backend_tensor_copy(hc_state_backend, hc_post_graph->residual_hc); + } + ggml_backend_tensor_copy(attn_out, hc_post_graph->block_out); + if (attn_post_backend && attn_comb_backend) { + ggml_backend_tensor_copy(attn_post_backend, hc_post_graph->post); + ggml_backend_tensor_copy(attn_comb_backend, hc_post_graph->comb); + } else { + ggml_backend_tensor_set(hc_post_graph->post, + hc_post.data(), 0, + sizeof(float) * hc_post.size()); + ggml_backend_tensor_set(hc_post_graph->comb, + hc_comb.data(), 0, + sizeof(float) * hc_comb.size()); } - ggml_backend_tensor_copy(attn_out, cached_decode_hc_post_graph.block_out); - ggml_backend_tensor_copy(attn_post_backend, cached_decode_hc_post_graph.post); - ggml_backend_tensor_copy(attn_comb_backend, cached_decode_hc_post_graph.comb); const auto hc_post_attn_t0 = Ds4TimingClock::now(); - if (ggml_backend_graph_compute(backend, cached_decode_hc_post_graph.sg.gf) != GGML_STATUS_SUCCESS) { + if (ggml_backend_graph_compute(backend, hc_post_graph->sg.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] cached hc-post compute failed layer %d attn\n", il); if (ctx) ggml_free(ctx); return false; } - hc_state_backend = cached_decode_hc_post_graph.sg.hidden_states; + hc_state_backend = hc_post_graph->sg.hidden_states; + hc_state_backend_current = true; + hc_state_host_current = false; + if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); + } else if (use_backend_decode_hc_post_direct_attn) { + const auto hc_post_attn_t0 = Ds4TimingClock::now(); + ggml_tensor * hc_next_backend = + hc_state_backend == hc_post_graph->residual_hc + ? hc_post_graph->sg.hidden_states + : hc_post_graph->residual_hc; + if (!ds4_try_gpu_hc_post_device(backend, + hc_state_backend, + attn_out, + attn_post_backend, + attn_comb_backend, + n_embd, + n_hc, + hc_next_backend)) { + std::fprintf(stderr, "[deepseek4] direct hc-post compute failed layer %d attn\n", il); + if (ctx) ggml_free(ctx); + return false; + } + if (sync_backend_decode_hc_post_direct) { + ggml_backend_synchronize(backend); + } + hc_state_backend = hc_next_backend; + hc_state_backend_current = true; + hc_state_host_current = false; if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); } else { const auto attn_read_t0 = Ds4TimingClock::now(); - if (use_backend_decode_hc_direct) { + if (attn_post_backend && attn_comb_backend) { const Ds4TensorReadback readbacks[] = { {attn_out, attn_out_host.data(), sizeof(float) * attn_out_host.size()}, @@ -5111,8 +6021,9 @@ bool deepseek4_step_layer_range( if (ctx) ggml_free(ctx); // ── HC post (attention) ───────────────────────────────── - if (!use_backend_decode_hc_graph) { + if (!use_backend_hc_post_graph && !use_backend_decode_hc_post_direct_attn) { const auto hc_post_attn_t0 = Ds4TimingClock::now(); + sync_hc_state_from_backend_for_host_pre(); hc_post_batch(next_hc, attn_out_host.data(), hc_state.data(), @@ -5122,6 +6033,8 @@ bool deepseek4_step_layer_range( n_embd, n_hc); std::memcpy(hc_state.data(), next_hc.data(), next_hc.size() * sizeof(float)); + hc_state_backend_current = false; + hc_state_host_current = true; if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us(hc_post_attn_t0, Ds4TimingClock::now()); } } @@ -5131,8 +6044,15 @@ bool deepseek4_step_layer_range( const auto hc_pre_ffn_t0 = Ds4TimingClock::now(); 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()); + ggml_tensor * hc_pre_input = hc_post_graph->residual_hc; + if (hc_state_backend_current && hc_state_backend) { + hc_pre_input = hc_state_backend; + } else { + ggml_backend_tensor_set(hc_post_graph->residual_hc, + hc_state.data(), 0, sizeof(float) * hc_state.size()); + hc_state_backend = hc_post_graph->residual_hc; + hc_state_backend_current = true; + } if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_ffn_input_t0, Ds4TimingClock::now()); auto & cached = cached_decode_ffn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || @@ -5154,7 +6074,7 @@ bool deepseek4_step_layer_range( backend, il, true, - cached_decode_hc_post_graph.residual_hc, + hc_pre_input, L.hc_ffn_fn, L.hc_ffn_scale, L.hc_ffn_base, @@ -5171,33 +6091,76 @@ bool deepseek4_step_layer_range( ffn_in_backend = cached.sg.hidden_states; ffn_post_backend = cached.post; ffn_comb_backend = cached.comb; - } else if (use_backend_decode_hc_graph) { - auto & cached = cached_decode_ffn_hc_pre_graphs[(size_t)il]; + } else if (use_backend_hc_pre_graph) { + auto & cached = use_backend_batch_hc_pre_graph + ? cached_batch_ffn_hc_pre_graphs[(size_t)il] + : cached_decode_ffn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || cached.owner_ctx != w.ctx || cached.backend != backend || cached.layer_idx != il || + cached.n_tokens != n_tokens || !cached.ffn) { const auto hc_pre_ffn_build_t0 = Ds4TimingClock::now(); - if (!build_cached_decode_hc_pre_graph(cached, backend, w, L, hc_lw.ffn.scale_data.data(), il, true)) { + if (!build_cached_decode_hc_pre_graph(cached, backend, w, L, hc_lw.ffn.scale_data.data(), il, true, n_tokens)) { std::fprintf(stderr, "[deepseek4] cached hc-pre graph alloc failed layer %d ffn\n", il); return false; } if (telemetry) telemetry->hc_pre_build_us += ds4_elapsed_us(hc_pre_ffn_build_t0, Ds4TimingClock::now()); } const auto hc_pre_ffn_input_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_copy(hc_state_backend, cached.sg.inp_embed); + sync_hc_state_to_backend_for_graph_pre(); + if (!use_backend_batch_hc_pre_direct) { + ggml_backend_tensor_copy(hc_state_backend, cached.sg.inp_embed); + } if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_ffn_input_t0, Ds4TimingClock::now()); const auto hc_pre_ffn_compute_t0 = Ds4TimingClock::now(); - if (ggml_backend_graph_compute(backend, cached.sg.gf) != GGML_STATUS_SUCCESS) { + if (use_backend_batch_hc_pre_direct) { + if (!ds4_try_gpu_hc_pre_batch_device(cached.sg.hidden_states, + cached.post, + cached.comb, + cached.pre, + backend, + hc_state_backend, + L.hc_ffn_fn, + L.hc_ffn_scale, + L.hc_ffn_base, + hc_lw.ffn.scale_data.data(), + hc_lw.ffn.base_data.data(), + n_tokens, + n_embd, + n_hc, + w.n_hc_sinkhorn_iter, + w.hc_eps)) { + std::fprintf(stderr, "[deepseek4] direct batch hc-pre compute failed layer %d ffn\n", il); + return false; + } + } else if (ggml_backend_graph_compute(backend, cached.sg.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] cached hc-pre compute 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()); + if (use_backend_batch_hc_pre_graph) { + ds4_verify_backend_batch_hc_pre_graph("ffn", + backend, + kv_start, + il, + n_tokens, + hc_state_backend, + hc_state, + hc_lw.ffn, + L.hc_ffn_fn, + cached, + n_embd, + n_hc, + w.n_hc_sinkhorn_iter, + w.hc_eps); + } ffn_in_backend = cached.sg.hidden_states; ffn_post_backend = cached.post; ffn_comb_backend = cached.comb; } else { + sync_hc_state_from_backend_for_host_pre(); hc_pre_batch(ffn_working, hc_post, hc_comb, hc_state.data(), hc_lw.ffn, L.hc_ffn_fn, n_tokens, n_embd, n_hc, w.n_hc_sinkhorn_iter, w.hc_eps); @@ -5212,20 +6175,32 @@ bool deepseek4_step_layer_range( il < w.n_hash_layer && L.ffn_gate_tid2eid && token_ids && hash_routing_tables_range[(size_t)il].loaded; if (hash_routed) { - const auto & ht = hash_routing_tables_range[(size_t)il].ids; + const auto & ht = hash_routing_tables_range[(size_t)il]; const int n_used = w.n_expert_used; + if (ht.width < n_used) { + std::fprintf(stderr, "[deepseek4] hash table width %d < topk %d for layer %d\n", + ht.width, n_used, il); + return false; + } hash_expert_ids_host.resize((size_t)n_used * (size_t)n_tokens); for (int ti = 0; ti < n_tokens; ti++) { const int32_t tok = token_ids[ti]; - const int32_t * row = ht.data() + (size_t)tok * (size_t)n_used; + if (tok < 0 || tok >= ht.n_vocab) { + std::fprintf(stderr, "[deepseek4] token id %d outside hash table for layer %d\n", + tok, il); + return false; + } + const int32_t * row = ht.ids.data() + (size_t)tok * (size_t)ht.width; memcpy(hash_expert_ids_host.data() + (size_t)ti * n_used, row, (size_t)n_used * sizeof(int32_t)); } } - auto & cached = cached_decode_ffn_graphs[(size_t)il]; + auto & cached = cache_layer_range_prefill_ffn + ? cached_prefill_ffn_graphs[(size_t)il] + : cached_decode_ffn_graphs[(size_t)il]; ggml_tensor * ffn_out = nullptr; - if (!hash_routed) { + if (!hash_routed && split_layer_range_route_expert) { auto & route = cached_decode_route_graphs_range[(size_t)il]; if (!route.valid() || route.owner_ctx != w.ctx || @@ -5283,18 +6258,29 @@ bool deepseek4_step_layer_range( } ffn_out = expert.sg.hidden_states; } else { - if (!cached.valid() || - cached.owner_ctx != w.ctx || - cached.backend != backend || - cached.layer_idx != il || - cached.n_tokens != n_tokens || - cached.hash_routed != hash_routed) { + const bool cached_ffn_hit = + cached.valid() && + cached.owner_ctx == w.ctx && + cached.backend == backend && + cached.layer_idx == il && + cached.n_tokens == n_tokens && + cached.hash_routed == hash_routed; + if (!cached_ffn_hit) { const auto ffn_build_t0 = Ds4TimingClock::now(); if (!build_cached_decode_ffn_graph(cached, backend, w, L, il, n_tokens, hash_routed)) { std::fprintf(stderr, "[deepseek4] cached ffn graph alloc failed layer %d\n", il); return false; } if (telemetry) telemetry->ffn_build_us += ds4_elapsed_us(ffn_build_t0, Ds4TimingClock::now()); + } else if (cache_layer_range_prefill_ffn) { + ++cached.reuse_count; + } + if (cache_layer_range_prefill_ffn) { + cached.last_used_epoch = ++prefill_ffn_cache_epoch; + if (cached.estimated_bytes == 0) { + cached.estimated_bytes = + ds4_estimate_cached_ffn_graph_bytes(w, n_tokens, hash_routed); + } } ffn_out = cached.sg.hidden_states; @@ -5318,23 +6304,58 @@ bool deepseek4_step_layer_range( } } - 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); + if (use_backend_hc_post_graph) { + if (hc_state_backend != hc_post_graph->residual_hc) { + ggml_backend_tensor_copy(hc_state_backend, hc_post_graph->residual_hc); + } + ggml_backend_tensor_copy(ffn_out, hc_post_graph->block_out); + if (ffn_post_backend && ffn_comb_backend) { + ggml_backend_tensor_copy(ffn_post_backend, hc_post_graph->post); + ggml_backend_tensor_copy(ffn_comb_backend, hc_post_graph->comb); + } else { + ggml_backend_tensor_set(hc_post_graph->post, + hc_post.data(), 0, + sizeof(float) * hc_post.size()); + ggml_backend_tensor_set(hc_post_graph->comb, + hc_comb.data(), 0, + sizeof(float) * hc_comb.size()); } - ggml_backend_tensor_copy(ffn_out, cached_decode_hc_post_graph.block_out); - ggml_backend_tensor_copy(ffn_post_backend, cached_decode_hc_post_graph.post); - ggml_backend_tensor_copy(ffn_comb_backend, cached_decode_hc_post_graph.comb); const auto hc_post_ffn_t0 = Ds4TimingClock::now(); - if (ggml_backend_graph_compute(backend, cached_decode_hc_post_graph.sg.gf) != GGML_STATUS_SUCCESS) { + if (ggml_backend_graph_compute(backend, hc_post_graph->sg.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] cached hc-post compute failed layer %d ffn\n", il); return false; } - hc_state_backend = cached_decode_hc_post_graph.sg.hidden_states; + hc_state_backend = hc_post_graph->sg.hidden_states; + hc_state_backend_current = true; + hc_state_host_current = false; + if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); + } else if (use_backend_decode_hc_post_direct_ffn) { + const auto hc_post_ffn_t0 = Ds4TimingClock::now(); + ggml_tensor * hc_next_backend = + hc_state_backend == hc_post_graph->residual_hc + ? hc_post_graph->sg.hidden_states + : hc_post_graph->residual_hc; + if (!ds4_try_gpu_hc_post_device(backend, + hc_state_backend, + ffn_out, + ffn_post_backend, + ffn_comb_backend, + n_embd, + n_hc, + hc_next_backend)) { + std::fprintf(stderr, "[deepseek4] direct hc-post compute failed layer %d ffn\n", il); + return false; + } + if (sync_backend_decode_hc_post_direct) { + ggml_backend_synchronize(backend); + } + hc_state_backend = hc_next_backend; + hc_state_backend_current = true; + hc_state_host_current = false; if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); } else { const auto ffn_read_t0 = Ds4TimingClock::now(); - if (use_backend_decode_hc_direct) { + if (ffn_post_backend && ffn_comb_backend) { const Ds4TensorReadback readbacks[] = { {ffn_out, ffn_out_host.data(), sizeof(float) * ffn_out_host.size()}, @@ -5352,8 +6373,9 @@ bool deepseek4_step_layer_range( } // ── HC post (FFN) ─────────────────────────────────────── - if (!use_backend_decode_hc_graph) { + if (!use_backend_hc_post_graph && !use_backend_decode_hc_post_direct_ffn) { const auto hc_post_ffn_t0 = Ds4TimingClock::now(); + sync_hc_state_from_backend_for_host_pre(); hc_post_batch(next_hc, ffn_out_host.data(), hc_state.data(), @@ -5363,6 +6385,8 @@ bool deepseek4_step_layer_range( n_embd, n_hc); std::memcpy(hc_state.data(), next_hc.data(), next_hc.size() * sizeof(float)); + hc_state_backend_current = false; + hc_state_host_current = true; if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); } } @@ -5372,8 +6396,11 @@ bool deepseek4_step_layer_range( w.output_hc_fn && w.output_hc_scale && w.output_hc_base; if ((!is_last_shard || !reuse_decode_graphs || !use_output_hc) && - use_backend_decode_hc_graph && hc_state_backend) { + !hc_state_host_current && + hc_state_backend_current && + hc_state_backend) { ggml_backend_tensor_get(hc_state_backend, hc_state.data(), 0, sizeof(float) * hc_state.size()); + hc_state_host_current = true; } // ── Output: HC pre → norm → lm_head (or return hidden state) ──────── @@ -5400,9 +6427,19 @@ bool deepseek4_step_layer_range( return false; } } - if (use_output_hc && hc_state_backend) { + if (use_output_hc && hc_state_backend_current && hc_state_backend) { ggml_backend_tensor_copy(hc_state_backend, cached_decode_output_graph.sg.hidden_input); + } else if (use_output_hc) { + if (!hc_state_host_current && hc_state_backend_current && hc_state_backend) { + ggml_backend_tensor_get(hc_state_backend, + hc_state.data(), + 0, + sizeof(float) * hc_state.size()); + hc_state_host_current = true; + } + ggml_backend_tensor_set(cached_decode_output_graph.sg.hidden_input, + hc_state.data(), 0, sizeof(float) * hc_state.size()); } else { ggml_backend_tensor_set(cached_decode_output_graph.sg.hidden_input, final_embd.data(), 0, sizeof(float) * final_embd.size()); diff --git a/server/src/deepseek4/deepseek4_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index f2c435241..57c79754f 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -1,12 +1,16 @@ #include "deepseek4_hc_cuda.h" #include "common/gpu_runtime_compat.h" +#include "ggml-backend-impl.h" +#include "ggml-cuda/common.cuh" #include #include +#include #include #include +#include #include #include #include @@ -24,42 +28,62 @@ constexpr int kMaxMixDim = 2 * kMaxHc + kMaxHc * kMaxHc; struct HcCudaScratch { float * d_state = nullptr; float * d_sums = nullptr; + float * d_inv_rms = nullptr; float * d_mix = nullptr; float * d_scale = nullptr; float * d_base = nullptr; float * d_working = nullptr; float * d_post = nullptr; float * d_comb = nullptr; + float * d_pre = nullptr; size_t state_cap = 0; + size_t sums_cap = 0; + size_t inv_rms_cap = 0; + size_t mix_cap = 0; + size_t working_cap = 0; + size_t post_cap = 0; + size_t comb_cap = 0; + size_t pre_cap = 0; ~HcCudaScratch() { if (d_state) cudaFree(d_state); if (d_sums) cudaFree(d_sums); + if (d_inv_rms) cudaFree(d_inv_rms); if (d_mix) cudaFree(d_mix); if (d_scale) cudaFree(d_scale); if (d_base) cudaFree(d_base); if (d_working) cudaFree(d_working); if (d_post) cudaFree(d_post); if (d_comb) cudaFree(d_comb); + if (d_pre) cudaFree(d_pre); } - bool ensure(size_t hc_dim, size_t n_embd, size_t n_hc) { - if (!d_sums && cudaMalloc(&d_sums, sizeof(float) * kSums) != cudaSuccess) return false; - if (!d_mix && cudaMalloc(&d_mix, sizeof(float) * kMixDim) != cudaSuccess) return false; + bool ensure(size_t hc_dim, size_t n_embd, size_t n_hc, size_t n_tokens = 1) { + auto ensure_buffer = [](float *& ptr, size_t & cap, size_t elems) { + if (cap >= elems) { + return true; + } + if (ptr) { + cudaFree(ptr); + ptr = nullptr; + cap = 0; + } + if (cudaMalloc(&ptr, sizeof(float) * elems) != cudaSuccess) { + return false; + } + cap = elems; + return true; + }; + if (!ensure_buffer(d_sums, sums_cap, (size_t) kSums * n_tokens)) return false; + if (!ensure_buffer(d_inv_rms, inv_rms_cap, n_tokens)) return false; + if (!ensure_buffer(d_mix, mix_cap, (size_t) kMaxMixDim * n_tokens)) return false; if (!d_scale && cudaMalloc(&d_scale, sizeof(float) * kMaxMixDim) != cudaSuccess) return false; if (!d_base && cudaMalloc(&d_base, sizeof(float) * kMaxMixDim) != cudaSuccess) return false; - if (!d_working && cudaMalloc(&d_working, sizeof(float) * n_embd) != cudaSuccess) return false; - if (!d_post && cudaMalloc(&d_post, sizeof(float) * n_hc) != cudaSuccess) return false; - if (!d_comb && cudaMalloc(&d_comb, sizeof(float) * n_hc * n_hc) != cudaSuccess) return false; - if (state_cap < hc_dim) { - if (d_state) { - cudaFree(d_state); - d_state = nullptr; - state_cap = 0; - } - if (cudaMalloc(&d_state, sizeof(float) * hc_dim) != cudaSuccess) return false; - state_cap = hc_dim; - } + if (!ensure_buffer(d_working, working_cap, n_embd * n_tokens)) return false; + if (!ensure_buffer(d_post, post_cap, n_hc * n_tokens)) return false; + if (!ensure_buffer(d_comb, comb_cap, n_hc * n_hc * n_tokens)) return false; + if (!ensure_buffer(d_pre, pre_cap, n_hc * n_tokens)) return false; + if (!ensure_buffer(d_state, state_cap, hc_dim)) return false; return true; } }; @@ -83,6 +107,11 @@ void hc_log_cuda_error(const char * label, cudaError_t err) { } } +bool hc_env_flag(const char * name) { + const char * value = std::getenv(name); + return value && value[0] && std::strcmp(value, "0") != 0; +} + __global__ void hc_sumsq_kernel(const float * x, int n, float * sums) { __shared__ float smem[kThreads]; const int tid = threadIdx.x; @@ -147,6 +176,56 @@ __global__ void hc_finalize_inv_rms_kernel(const float * sums, } } +__global__ void hc_sumsq_batch_kernel(const float * x, + int cols, + float * sums) { + __shared__ float smem[kThreads]; + const int tid = threadIdx.x; + const int bid = blockIdx.x; + const int token = blockIdx.y; + const float * x_token = x + (size_t) token * (size_t) cols; + float acc = 0.0f; + for (int i = bid * blockDim.x + tid; i < cols; i += gridDim.x * blockDim.x) { + const float v = x_token[i]; + acc += v * v; + } + smem[tid] = acc; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] += smem[tid + stride]; + __syncthreads(); + } + if (tid == 0) { + sums[(size_t) token * kSums + bid] = smem[0]; + } +} + +__global__ void hc_finalize_inv_rms_batch_kernel(const float * sums, + int n_sums, + int cols, + float eps, + float * inv_rms_out) { + __shared__ float smem[kThreads]; + const int tid = threadIdx.x; + const int token = blockIdx.x; + const float * sums_token = sums + (size_t) token * (size_t) n_sums; + float acc = 0.0f; + for (int i = tid; i < n_sums; i += blockDim.x) { + acc += sums_token[i]; + } + smem[tid] = acc; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + smem[tid] += smem[tid + stride]; + } + __syncthreads(); + } + if (tid == 0) { + inv_rms_out[token] = rsqrtf(smem[0] / (float) cols + eps); + } +} + __global__ void hc_mix_kernel_device_inv(const float * x, const __half * fn, int cols, @@ -174,56 +253,40 @@ __global__ void hc_mix_kernel_device_inv(const float * x, } } -__device__ float hc_sigmoid(float x) { - return 1.0f / (1.0f + expf(-x)); -} - -__global__ void hc_mix_norm_kernel(const float * x, - const __half * fn, - int cols, - int rows, - float eps, - float * mix) { +__global__ void hc_mix_batch_kernel_device_inv(const float * x, + const __half * fn, + int cols, + const float * inv_rms_ptr, + int mix_dim, + float * mix) { __shared__ float smem[kThreads]; - __shared__ float inv_rms; const int row = blockIdx.x; + const int token = blockIdx.y; const int tid = threadIdx.x; - if (row >= rows) { - return; - } - - float sumsq = 0.0f; - for (int c = tid; c < cols; c += blockDim.x) { - const float v = x[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(); - } - if (tid == 0) { - inv_rms = rsqrtf(smem[0] / (float) cols + eps); - } - __syncthreads(); - + const float * x_token = x + (size_t) token * (size_t) cols; + const float inv_rms = inv_rms_ptr[token]; + float acc = 0.0f; const __half * w = fn + (size_t) row * (size_t) cols; - float dot = 0.0f; for (int c = tid; c < cols; c += blockDim.x) { - dot += __half2float(w[c]) * x[c]; + acc += __half2float(w[c]) * (x_token[c] * inv_rms); } - smem[tid] = dot * inv_rms; + smem[tid] = acc; __syncthreads(); for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) smem[tid] += smem[tid + stride]; + if (tid < stride) { + smem[tid] += smem[tid + stride]; + } __syncthreads(); } if (tid == 0) { - mix[row] = smem[0]; + mix[(size_t) token * (size_t) mix_dim + row] = smem[0]; } } +__device__ float hc_sigmoid(float x) { + return 1.0f / (1.0f + expf(-x)); +} + __global__ void hc_finish_kernel(const float * hc_state, const float * mix, const float * scale, @@ -309,6 +372,143 @@ __global__ void hc_finish_kernel(const float * hc_state, } } +__global__ void hc_post_kernel(const float * residual_hc, + const float * block_out, + const float * post, + const float * comb, + int n_embd, + int n_hc, + float * out_hc) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int total = n_embd * n_hc; + if (idx >= total) { + return; + } + + const int d = idx % n_embd; + const int dst = idx / n_embd; + float acc = block_out[d] * post[dst]; + for (int src = 0; src < n_hc; ++src) { + acc += comb[dst + src * n_hc] * + residual_hc[(size_t) src * (size_t) n_embd + (size_t) d]; + } + out_hc[(size_t) dst * (size_t) n_embd + (size_t) d] = acc; +} + +__global__ void hc_finish_batch_kernel(const float * hc_state, + const float * mix, + const float * scale, + const float * base, + int n_tokens, + int n_embd, + int n_hc, + int sinkhorn_iters, + float * working, + float * post, + float * comb, + float * pre) { + (void) n_tokens; + __shared__ float split[kMaxMixDim]; + const int token = blockIdx.x; + const int tid = threadIdx.x; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + const float * hc_state_token = + hc_state + (size_t) token * (size_t) n_embd * (size_t) n_hc; + const float * mix_token = mix + (size_t) token * (size_t) mix_dim; + float * working_token = working + (size_t) token * (size_t) n_embd; + float * post_token = post + (size_t) token * (size_t) n_hc; + float * comb_token = comb + (size_t) token * (size_t) n_hc * (size_t) n_hc; + float * pre_token = pre ? pre + (size_t) token * (size_t) n_hc : nullptr; + + 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_token[i] * pre_scale + base[i]) + sinkhorn_eps; + if (pre_token) { + pre_token[i] = split[i]; + } + } + for (int i = 0; i < n_hc; ++i) { + split[n_hc + i] = + 2.0f * hc_sigmoid(mix_token[n_hc + i] * post_scale + base[n_hc + i]); + post_token[i] = split[n_hc + i]; + } + + 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_token[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; + } + } + } + + for (int i = 0; i < n_hc * n_hc; ++i) { + split[2 * n_hc + i] = c[i]; + comb_token[i] = c[i]; + } + } + __syncthreads(); + + 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] * hc_state_token[(size_t) h * n_embd + d]; + } + working_token[d] = acc; + } +} + bool hc_pre_device_locked(const void * hc_state_device, const void * fn_device, const void * scale_device, @@ -381,7 +581,7 @@ bool hc_pre_device_locked(const void * hc_state_device, return fail("finish kernel", err); } if (sync_device) { - err = cudaDeviceSynchronize(); + err = stream ? cudaStreamSynchronize(stream) : cudaDeviceSynchronize(); if (err != cudaSuccess) { return fail("finish sync", err); } @@ -390,6 +590,116 @@ bool hc_pre_device_locked(const void * hc_state_device, return true; } +bool hc_pre_batch_device_locked(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, + void * pre_device, + HcCudaScratch & scratch, + cudaStream_t stream, + bool sync_device, + bool log_errors) { + const int hc_dim = n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + auto fail = [log_errors](const char * label, cudaError_t err) { + if (log_errors) { + hc_log_cuda_error(label, err); + } + return false; + }; + if (mix_dim > kMixDim) { + if (log_errors) { + std::fprintf(stderr, + "[deepseek4-hc-direct] batch mix dim too large: %d > %d\n", + mix_dim, kMixDim); + } + return false; + } + const float * hc_state_ptr = static_cast(hc_state_device); + + hc_sumsq_batch_kernel<<>>( + hc_state_ptr, + hc_dim, + scratch.d_sums); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + return fail("batch sumsq kernel", err); + } + + hc_finalize_inv_rms_batch_kernel<<>>( + scratch.d_sums, + kSums, + hc_dim, + eps, + scratch.d_inv_rms); + err = cudaGetLastError(); + if (err != cudaSuccess) { + return fail("batch inv_rms kernel", err); + } + + hc_mix_batch_kernel_device_inv<<>>( + hc_state_ptr, + static_cast(fn_device), + hc_dim, + scratch.d_inv_rms, + mix_dim, + scratch.d_mix); + err = cudaGetLastError(); + if (err != cudaSuccess) { + return fail("batch mix kernel", err); + } + + hc_finish_batch_kernel<<>>( + hc_state_ptr, + scratch.d_mix, + static_cast(scale_device), + static_cast(base_device), + n_tokens, + n_embd, + n_hc, + sinkhorn_iters, + static_cast(working_device), + static_cast(post_device), + static_cast(comb_device), + static_cast(pre_device)); + err = cudaGetLastError(); + if (err != cudaSuccess) { + return fail("batch finish kernel", err); + } + if (sync_device) { + err = stream ? cudaStreamSynchronize(stream) : cudaDeviceSynchronize(); + if (err != cudaSuccess) { + return fail("batch finish sync", err); + } + } + + return true; +} + +static cudaStream_t hc_backend_stream_or_default(ggml_backend_t backend) { + if (!backend || !backend->context) { + return nullptr; + } + if (!backend->device) { + return nullptr; + } + const auto dev_type = ggml_backend_dev_type(backend->device); + if (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && + dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU) { + return nullptr; + } + auto * cuda_ctx = static_cast(backend->context); + return cuda_ctx ? cuda_ctx->stream() : nullptr; +} + } // namespace bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, @@ -552,6 +862,52 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, false); } +bool deepseek4_cuda_hc_pre_device_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const void * scale_device, + const void * base_device, + 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_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { + return false; + } + const int hc_dim = n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + if (mix_dim > kMaxMixDim) { + return false; + } + + cudaStream_t stream = hc_backend_stream_or_default(backend); + std::lock_guard lock(g_mu); + HcCudaScratch & scratch = hc_scratch_for_stream(stream); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + return false; + } + return hc_pre_device_locked(hc_state_device, + fn_device, + scale_device, + base_device, + n_embd, + n_hc, + sinkhorn_iters, + eps, + working_device, + post_device, + comb_device, + scratch, + stream, + true, + false); +} + bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, const void * fn_device, const float * scale_host, @@ -607,4 +963,221 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, true); } +bool deepseek4_cuda_hc_pre_device_params_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const float * scale_host, + const float * base_host, + 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_host || !base_host || + !working_device || !post_device || !comb_device || + n_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { + return false; + } + const int hc_dim = n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + if (mix_dim > kMaxMixDim) { + return false; + } + + cudaStream_t stream = hc_backend_stream_or_default(backend); + std::lock_guard lock(g_mu); + HcCudaScratch & scratch = hc_scratch_for_stream(stream); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc)) { + std::fprintf(stderr, "[deepseek4-hc-direct] ensure failed\n"); + return false; + } + if (cudaMemcpy(scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, + cudaMemcpyHostToDevice) != cudaSuccess) { + hc_log_cuda_error("copy scale", cudaGetLastError()); + return false; + } + if (cudaMemcpy(scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, + cudaMemcpyHostToDevice) != cudaSuccess) { + hc_log_cuda_error("copy base", cudaGetLastError()); + return false; + } + return hc_pre_device_locked(hc_state_device, + fn_device, + scratch.d_scale, + scratch.d_base, + n_embd, + n_hc, + sinkhorn_iters, + eps, + working_device, + post_device, + comb_device, + scratch, + stream, + true, + true); +} + +bool deepseek4_cuda_hc_post_device_on_backend(ggml_backend_t backend, + const void * residual_hc_device, + const void * block_out_device, + const void * post_device, + const void * comb_device, + 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_embd <= 0 || n_hc <= 0 || + n_hc > kMaxHc) { + return false; + } + + cudaStream_t stream = hc_backend_stream_or_default(backend); + const int total = n_embd * n_hc; + const int blocks = (total + kThreads - 1) / kThreads; + hc_post_kernel<<>>( + static_cast(residual_hc_device), + static_cast(block_out_device), + static_cast(post_device), + static_cast(comb_device), + n_embd, + n_hc, + static_cast(out_hc_device)); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + hc_log_cuda_error("hc_post kernel", err); + return false; + } + return true; +} + +bool deepseek4_cuda_hc_pre_batch_device_on_backend(ggml_backend_t backend, + 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, + void * pre_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 hc_dim = n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + if (mix_dim > kMaxMixDim) { + return false; + } + + const bool use_default_stream = + hc_env_flag("DFLASH_DS4_BATCH_HC_PRE_DIRECT_DEFAULT_STREAM"); + cudaStream_t stream = use_default_stream + ? nullptr + : hc_backend_stream_or_default(backend); + std::lock_guard lock(g_mu); + HcCudaScratch & scratch = hc_scratch_for_stream(stream); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc, (size_t) n_tokens)) { + return false; + } + return hc_pre_batch_device_locked(hc_state_device, + fn_device, + scale_device, + base_device, + n_tokens, + n_embd, + n_hc, + sinkhorn_iters, + eps, + working_device, + post_device, + comb_device, + pre_device, + scratch, + stream, + !use_default_stream, + false); +} + +bool deepseek4_cuda_hc_pre_batch_device_params_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const float * scale_host, + const float * base_host, + int n_tokens, + int n_embd, + int n_hc, + int sinkhorn_iters, + float eps, + void * working_device, + void * post_device, + void * comb_device, + void * pre_device) { + if (!hc_state_device || !fn_device || !scale_host || !base_host || + !working_device || !post_device || !comb_device || + n_tokens <= 0 || n_embd <= 0 || n_hc <= 0 || n_hc > kMaxHc) { + return false; + } + const int hc_dim = n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + if (mix_dim > kMaxMixDim) { + return false; + } + + const bool use_default_stream = + hc_env_flag("DFLASH_DS4_BATCH_HC_PRE_DIRECT_DEFAULT_STREAM"); + cudaStream_t stream = use_default_stream + ? nullptr + : hc_backend_stream_or_default(backend); + std::lock_guard lock(g_mu); + HcCudaScratch & scratch = hc_scratch_for_stream(stream); + if (!scratch.ensure((size_t) hc_dim, (size_t) n_embd, (size_t) n_hc, (size_t) n_tokens)) { + std::fprintf(stderr, "[deepseek4-hc-direct] ensure failed\n"); + return false; + } + cudaError_t err = cudaMemcpyAsync(scratch.d_scale, + scale_host, + sizeof(float) * (size_t) mix_dim, + cudaMemcpyHostToDevice, + stream); + if (err != cudaSuccess) { + hc_log_cuda_error("batch copy scale", err); + return false; + } + err = cudaMemcpyAsync(scratch.d_base, + base_host, + sizeof(float) * (size_t) mix_dim, + cudaMemcpyHostToDevice, + stream); + if (err != cudaSuccess) { + hc_log_cuda_error("batch copy base", err); + return false; + } + return hc_pre_batch_device_locked(hc_state_device, + fn_device, + scratch.d_scale, + scratch.d_base, + n_tokens, + n_embd, + n_hc, + sinkhorn_iters, + eps, + working_device, + post_device, + comb_device, + pre_device, + scratch, + stream, + !use_default_stream, + true); +} + } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_hc_cuda.h b/server/src/deepseek4/deepseek4_hc_cuda.h index d2268ac12..3102d8fce 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.h +++ b/server/src/deepseek4/deepseek4_hc_cuda.h @@ -2,6 +2,8 @@ #include +#include "ggml-backend.h" + namespace dflash::common { bool deepseek4_cuda_hc_pre_mix(const float * hc_state_host, @@ -35,6 +37,19 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, void * post_device, void * comb_device); +bool deepseek4_cuda_hc_pre_device_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const void * scale_device, + const void * base_device, + int n_embd, + int n_hc, + int sinkhorn_iters, + float eps, + void * working_device, + void * post_device, + void * comb_device); + bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, const void * fn_device, const float * scale_host, @@ -47,4 +62,56 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, void * post_device, void * comb_device); +bool deepseek4_cuda_hc_pre_device_params_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const float * scale_host, + const float * base_host, + 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_device_on_backend(ggml_backend_t backend, + const void * residual_hc_device, + const void * block_out_device, + const void * post_device, + const void * comb_device, + int n_embd, + int n_hc, + void * out_hc_device); + +bool deepseek4_cuda_hc_pre_batch_device_on_backend(ggml_backend_t backend, + 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, + void * pre_device = nullptr); + +bool deepseek4_cuda_hc_pre_batch_device_params_on_backend(ggml_backend_t backend, + const void * hc_state_device, + const void * fn_device, + const float * scale_host, + const float * base_host, + int n_tokens, + int n_embd, + int n_hc, + int sinkhorn_iters, + float eps, + void * working_device, + void * post_device, + void * comb_device, + void * pre_device = nullptr); + } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index ef3aa9fbf..3f134529e 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -304,7 +304,22 @@ bool load_deepseek4_gguf_partial(const std::string & path, const uint32_t n_lora_o = get_u32_or(gctx, "deepseek4.attention.output_lora_rank", 1024); const uint32_t n_out_group = get_u32_or(gctx, "deepseek4.attention.output_group_count", 8); const uint32_t n_expert = get_u32_or(gctx, "deepseek4.expert_count", 256); - const uint32_t n_expert_used = get_u32_or(gctx, "deepseek4.expert_used_count", 6); + const uint32_t n_expert_used_meta = get_u32_or(gctx, "deepseek4.expert_used_count", 6); + uint32_t n_expert_used = n_expert_used_meta; + if (const char * topk_env = std::getenv("DFLASH_DS4_TOPK")) { + char * end = nullptr; + const long parsed = std::strtol(topk_env, &end, 10); + if (end != topk_env && parsed > 0) { + const uint32_t requested = (uint32_t) parsed; + n_expert_used = std::max( + 1, std::min(requested, n_expert_used_meta)); + if (n_expert_used != n_expert_used_meta) { + std::fprintf(stderr, + "[deepseek4] DFLASH_DS4_TOPK=%u: routed experts %u -> %u\n", + requested, n_expert_used_meta, n_expert_used); + } + } + } const uint32_t n_expert_shared = get_u32_or(gctx, "deepseek4.expert_shared_count", 1); const uint32_t n_ff_exp = get_u32_or(gctx, "deepseek4.expert_feed_forward_length", 2048); const uint32_t n_hash_layer = get_u32_or(gctx, "deepseek4.hash_layer_count", 3); diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 8d4bdab24..efe88c4e5 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -108,6 +108,89 @@ static bool prompt_ends_in_open_think(const std::string & prompt) { prompt.compare(end - kThinkOpenLen, kThinkOpenLen, kThinkOpen) == 0; } +class Utf8OutputFilter { +public: + std::string push(const std::string & text) { + pending_ += text; + return drain(false); + } + + std::string finish() { + std::string out = drain(true); + pending_.clear(); + return out; + } + +private: + static bool is_cont(uint8_t c) { + return (c & 0xC0) == 0x80; + } + + std::string drain(bool final) { + std::string out; + size_t i = 0; + while (i < pending_.size()) { + const uint8_t c0 = (uint8_t) pending_[i]; + if (c0 < 0x80) { + out.push_back((char)c0); + ++i; + continue; + } + + int need = 0; + uint32_t cp = 0; + uint32_t min_cp = 0; + if (c0 >= 0xC2 && c0 <= 0xDF) { + need = 2; + cp = c0 & 0x1F; + min_cp = 0x80; + } else if (c0 >= 0xE0 && c0 <= 0xEF) { + need = 3; + cp = c0 & 0x0F; + min_cp = 0x800; + } else if (c0 >= 0xF0 && c0 <= 0xF4) { + need = 4; + cp = c0 & 0x07; + min_cp = 0x10000; + } else { + ++i; + continue; + } + + if (i + (size_t)need > pending_.size()) { + if (!final) break; + i = pending_.size(); + break; + } + + bool ok = true; + for (int j = 1; j < need; ++j) { + const uint8_t cj = (uint8_t) pending_[i + (size_t)j]; + if (!is_cont(cj)) { + ok = false; + break; + } + cp = (cp << 6) | (cj & 0x3F); + } + if (!ok || cp < min_cp || (cp >= 0xD800 && cp <= 0xDFFF) || + cp > 0x10FFFF) { + ++i; + continue; + } + + // Do not surface replacement characters from byte-fallback noise. + if (cp != 0xFFFD) { + out.append(pending_, i, (size_t)need); + } + i += (size_t)need; + } + pending_.erase(0, i); + return out; + } + + std::string pending_; +}; + // ─── piecewise keep-ratio curve ───────────────────────────────────────── static float pflash_keep_ratio(const ServerConfig & cfg, int n_tokens) { @@ -2755,6 +2838,28 @@ void HttpServer::worker_loop() { int completion_tokens = 0; bool visible_output_seen = false; bool client_disconnected = false; + Utf8OutputFilter live_output_filter; + + auto send_live_text = [&](const std::string & text) -> bool { + if (text.empty()) return true; + visible_output_seen = true; + broadcast_token(text); + if (req.stream) { + auto chunks = emitter.emit_token(text); + for (const auto & chunk : chunks) { + if (!send_all(fd, chunk.data(), chunk.size())) { + client_disconnected = true; + return false; + } + } + if (emitter.stop_hit()) return false; + } + return true; + }; + + auto emit_live_text = [&](const std::string & text) -> bool { + return send_live_text(live_output_filter.push(text)); + }; io.on_token = [&](int32_t token) -> bool { if (client_disconnected) return false; @@ -2775,23 +2880,11 @@ void HttpServer::worker_loop() { // Gemma4 thinking channel: map <|channel> → , → \n if (raw == "<|channel>") { - visible_output_seen = true; - broadcast_token(""); - if (req.stream) { - auto chunks = emitter.emit_token(""); - for (const auto & chunk : chunks) - if (!send_all(fd, chunk.data(), chunk.size())) { client_disconnected = true; return false; } - } + if (!emit_live_text("")) return false; return true; } if (raw == "") { - visible_output_seen = true; - broadcast_token("\n"); - if (req.stream) { - auto chunks = emitter.emit_token("\n"); - for (const auto & chunk : chunks) - if (!send_all(fd, chunk.data(), chunk.size())) { client_disconnected = true; return false; } - } + if (!emit_live_text("\n")) return false; return true; } @@ -2803,14 +2896,8 @@ void HttpServer::worker_loop() { // reasoning_content with empty visible content. Forward the text // form into the emitter so parse_reasoning() can split correctly. if (raw == "" || raw == "") { - visible_output_seen = true; - broadcast_token(raw == "" ? "\n" : ""); - if (req.stream) { - auto chunks = emitter.emit_token( - raw == "" ? "\n" : ""); - for (const auto & chunk : chunks) - if (!send_all(fd, chunk.data(), chunk.size())) { client_disconnected = true; return false; } - } + if (!emit_live_text(raw == "" ? "\n" : "")) + return false; return true; } @@ -2822,24 +2909,7 @@ void HttpServer::worker_loop() { } std::string text = tokenizer_.token_text(token); - - // Send token text to status page clients (browser accumulates). - if (!text.empty()) { - visible_output_seen = true; - broadcast_token(text); - } - - if (req.stream && !text.empty()) { - auto chunks = emitter.emit_token(text); - for (const auto & chunk : chunks) { - if (!send_all(fd, chunk.data(), chunk.size())) { - client_disconnected = true; - return false; - } - } - // Stop generation if a stop sequence was hit. - if (emitter.stop_hit()) return false; - } + if (!emit_live_text(text)) return false; return true; }; @@ -2871,6 +2941,9 @@ void HttpServer::worker_loop() { } else { result = backend_.generate(gen_req, io); } + if (!client_disconnected) { + (void) send_live_text(live_output_filter.finish()); + } if (dflash_residency == DraftResidencyAction::ReleaseAfterUse && !config_.draft_path.empty()) { @@ -3022,14 +3095,27 @@ void HttpServer::worker_loop() { } else if (!req.stream && !client_disconnected) { // Non-streaming: build complete response using emitter state. // Feed all tokens through emitter (skip specials like streaming path). + Utf8OutputFilter response_output_filter; + auto emit_response_text = [&](const std::string & text) -> bool { + std::string filtered = response_output_filter.push(text); + if (filtered.empty()) return true; + emitter.emit_token(filtered); + return !emitter.stop_hit(); + }; auto feed_tokens = [&](const std::vector & toks) -> bool { for (int32_t tok : toks) { const std::string & raw = tokenizer_.raw_token(tok); if (tok == tokenizer_.eos_id()) continue; if (tok == tokenizer_.eos_chat_id()) continue; // Gemma4 channel → think mapping - if (raw == "<|channel>") { emitter.emit_token(""); continue; } - if (raw == "") { emitter.emit_token("\n"); continue; } + if (raw == "<|channel>") { + if (!emit_response_text("")) return false; + continue; + } + if (raw == "") { + if (!emit_response_text("\n")) return false; + continue; + } // Qwen3.6 thinking tokens (id 248068 / 248069) — must // forward as text so the emitter transitions // reasoning→content. Without this the generic <...> @@ -3037,21 +3123,32 @@ void HttpServer::worker_loop() { // empty and the model's whole answer wedged in // reasoning_content. Mirrors the streaming-path fix // above. - if (raw == "") { emitter.emit_token(""); continue; } - if (raw == "") { emitter.emit_token("\n"); continue; } + if (raw == "") { + if (!emit_response_text("")) return false; + continue; + } + if (raw == "") { + if (!emit_response_text("\n")) return false; + continue; + } if (raw.size() >= 2 && raw[0] == '<' && raw[1] == '|') continue; if (raw.size() >= 2 && raw[0] == '<' && raw.back() == '>') { if (!(raw.size() == 6 && raw[1] == '0' && raw[2] == 'x')) continue; } std::string text = tokenizer_.token_text(tok); - emitter.emit_token(text); - if (emitter.stop_hit()) return false; + if (!emit_response_text(text)) return false; } return true; }; - feed_tokens(result.tokens); + const bool feed_ok = feed_tokens(result.tokens); + if (feed_ok) { + std::string tail = response_output_filter.finish(); + if (!tail.empty()) { + emitter.emit_token(tail); + } + } const int total_completion_tokens = (int)result.tokens.size(); emitter.emit_finish(total_completion_tokens); From 3e11f02893401ac9777a4309b41280ff6a27bb3d Mon Sep 17 00:00:00 2001 From: weicj Date: Fri, 10 Jul 2026 09:02:25 +0800 Subject: [PATCH 7/7] fix(deepseek4): address cubic review comments --- server/src/deepseek4/deepseek4_graph.cpp | 84 ++++++++++------------- server/src/deepseek4/deepseek4_hc_cuda.cu | 4 +- server/src/server/http_server.cpp | 5 +- 3 files changed, 40 insertions(+), 53 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index e6ab64734..68687005d 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -3477,6 +3477,7 @@ static bool deepseek4_step_hybrid( const int32_t * token_ids, ExpertIpcClient * expert_worker, bool worker_owns_hot_ids, + bool want_logits, bool disable_cached_decode, MoeExpertCompute * expert_compute, const MoeExpertLayer * expert_layers, @@ -4732,6 +4733,14 @@ static bool deepseek4_step_hybrid( if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); + if (!want_logits) { + cache.cur_pos = kv_start + n_tokens; + if (telemetry) { + telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + } + return true; + } + const bool use_output_hc = w.output_hc_fn && w.output_hc_scale && w.output_hc_base; @@ -5015,6 +5024,7 @@ bool deepseek4_step( const int32_t * token_ids, ExpertIpcClient * expert_worker, bool worker_owns_hot_ids, + bool want_logits, bool disable_cached_decode, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats, @@ -5043,7 +5053,7 @@ bool deepseek4_step( kv_start, /*layer_begin=*/0, /*layer_end=*/w.n_layer, - &out_logits, + want_logits ? &out_logits : nullptr, token_ids, telemetry); } @@ -5052,6 +5062,7 @@ bool deepseek4_step( return deepseek4_step_hybrid(backend, w, cache, *moe_hybrid, embed, n_tokens, kv_start, out_logits, token_ids, expert_worker, worker_owns_hot_ids, + want_logits, disable_cached_decode, expert_compute, expert_layers, telemetry, routing_stats); @@ -5165,22 +5176,23 @@ bool deepseek4_step( cur = ggml_add(ctx, cur, ffn_out); } - // ── Output head ───────────────────────────────────────────────────── - // TODO: HC output pre (merge residual streams for final projection) - - // Final RMSNorm - cur = build_rms_norm(ctx, cur, w.out_norm, w.rms_eps); - - // lm_head projection - ggml_tensor * logits = ggml_mul_mat(ctx, w.output, cur); - ggml_set_name(logits, "logits"); - ggml_set_output(logits); - if (cached_sg) { - cached_sg->logits = logits; + ggml_tensor * logits = nullptr; + ggml_tensor * graph_output = cur; + if (want_logits) { + // ── Output head ───────────────────────────────────────────────────── + // TODO: HC output pre (merge residual streams for final projection) + cur = build_rms_norm(ctx, cur, w.out_norm, w.rms_eps); + logits = ggml_mul_mat(ctx, w.output, cur); + ggml_set_name(logits, "logits"); + graph_output = logits; + if (cached_sg) { + cached_sg->logits = logits; + } } + ggml_set_output(graph_output); // ── Build and run graph ───────────────────────────────────────────── - ggml_build_forward_expand(gf, logits); + ggml_build_forward_expand(gf, graph_output); // Allocate if (cached_sg) { @@ -5232,14 +5244,16 @@ bool deepseek4_step( telemetry->full_graph_compute_us += ds4_elapsed_us(full_compute_t0, Ds4TimingClock::now()); } - // Read logits (only last token for generation) - const auto full_read_t0 = Ds4TimingClock::now(); - out_logits.resize(w.n_vocab); - const size_t logits_offset = (size_t)(n_tokens - 1) * w.n_vocab * sizeof(float); - ggml_backend_tensor_get(logits, out_logits.data(), logits_offset, - w.n_vocab * sizeof(float)); - if (telemetry) { - telemetry->full_graph_read_us += ds4_elapsed_us(full_read_t0, Ds4TimingClock::now()); + if (want_logits) { + // Read logits (only last token for generation) + const auto full_read_t0 = Ds4TimingClock::now(); + out_logits.resize(w.n_vocab); + const size_t logits_offset = (size_t)(n_tokens - 1) * w.n_vocab * sizeof(float); + ggml_backend_tensor_get(logits, out_logits.data(), logits_offset, + w.n_vocab * sizeof(float)); + if (telemetry) { + telemetry->full_graph_read_us += ds4_elapsed_us(full_read_t0, Ds4TimingClock::now()); + } } release_full_step(); @@ -5264,28 +5278,6 @@ bool deepseek4_step( return true; } -bool deepseek4_step( - ggml_backend_t backend, - const DeepSeek4Weights & w, - DeepSeek4Cache & cache, - const float * embed, - int n_tokens, - int kv_start, - std::vector & out_logits, - MoeHybridStorage * moe_hybrid, - const int32_t * token_ids, - ExpertIpcClient * expert_worker, - bool worker_owns_hot_ids, - DeepSeek4StepTelemetry * telemetry, - MoeHybridRoutingStats * routing_stats, - MoeExpertCompute * expert_compute, - const MoeExpertLayer * expert_layers) { - return deepseek4_step(backend, w, cache, embed, n_tokens, kv_start, - out_logits, moe_hybrid, token_ids, expert_worker, - worker_owns_hot_ids, false, telemetry, - routing_stats, expert_compute, expert_layers); -} - bool deepseek4_step( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -5299,15 +5291,13 @@ bool deepseek4_step( ExpertIpcClient * expert_worker, bool worker_owns_hot_ids, bool want_logits, - bool disable_cached_decode, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats, MoeExpertCompute * expert_compute, const MoeExpertLayer * expert_layers) { - (void) want_logits; return deepseek4_step(backend, w, cache, embed, n_tokens, kv_start, out_logits, moe_hybrid, token_ids, expert_worker, - worker_owns_hot_ids, disable_cached_decode, telemetry, + worker_owns_hot_ids, want_logits, false, telemetry, routing_stats, expert_compute, expert_layers); } diff --git a/server/src/deepseek4/deepseek4_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index 57c79754f..2dce8e712 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -615,11 +615,11 @@ bool hc_pre_batch_device_locked(const void * hc_state_device, } return false; }; - if (mix_dim > kMixDim) { + if (mix_dim > kMaxMixDim) { if (log_errors) { std::fprintf(stderr, "[deepseek4-hc-direct] batch mix dim too large: %d > %d\n", - mix_dim, kMixDim); + mix_dim, kMaxMixDim); } return false; } diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index efe88c4e5..b03f8919b 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -178,10 +178,7 @@ class Utf8OutputFilter { continue; } - // Do not surface replacement characters from byte-fallback noise. - if (cp != 0xFFFD) { - out.append(pending_, i, (size_t)need); - } + out.append(pending_, i, (size_t)need); i += (size_t)need; } pending_.erase(0, i);