diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index b693e1041..2e54ce514 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -271,7 +271,6 @@ 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 @@ -280,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 @@ -379,7 +379,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 +419,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 +438,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. @@ -643,8 +636,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 +650,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 +773,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 +870,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/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..8047429f1 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3,6 +3,7 @@ #include "deepseek4_backend.h" #include "deepseek4_internal.h" #include "common/sampler.h" +#include "../common/moe_hybrid_types.h" #include "ggml.h" #include "ggml-backend.h" @@ -10,6 +11,7 @@ #include #include +#include #include #include #include @@ -40,6 +42,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; @@ -61,10 +68,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; @@ -85,22 +116,44 @@ 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 " + "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", 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), - 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), @@ -241,24 +294,21 @@ static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, std::string * err) { out = {}; ggml_backend_cuda_get_device_memory(gpu, &out.gpu_free, &out.gpu_total); - if (out.gpu_total == 0) { - if (err) *err = "could not query GPU memory"; - return false; - } if (!compute_ds4_expert_memory_info(w, nullptr, out.mem, err)) { return false; } - out.core_bytes = out.gpu_total - out.gpu_free; + out.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; + 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.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) { @@ -270,7 +320,9 @@ 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)); + out.max_hot_per_layer = std::min( + w.n_expert, + (int) (out.expert_budget / out.mem.bytes_per_uniform_round)); if (out.max_hot_per_layer <= 0) { if (err) *err = "expert budget is smaller than one uniform expert round"; return false; @@ -288,16 +340,23 @@ 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; } -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 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; } } // namespace @@ -309,6 +368,44 @@ 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; + + 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_runtime_.compute_ptr(), + expert_runtime_.layer_ptr(0)); +} + +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_) { @@ -319,25 +416,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; } } @@ -347,11 +440,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, @@ -364,7 +453,8 @@ 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, cfg_.device.gpu, max_ctx, budget, err)) { return false; } @@ -390,6 +480,38 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & return true; } +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() { TargetLoadPlan plan; plan.skip_expert_tensors = true; @@ -399,6 +521,16 @@ 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)) { @@ -418,42 +550,28 @@ 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)) { + if (!build_deepseek4_moe_hybrid_storage_from_file( + 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; } - if (hybrid->has_mmap() && !hybrid->materialized_cold_experts) { - size_t max_expert_bytes = 0; - for (const auto & layer : hybrid->layers) { - const size_t per_expert_bytes = layer.fused_gate_up - ? layer.gate_up_expert_bytes + layer.down_expert_bytes - : layer.gate_expert_bytes + layer.up_expert_bytes + layer.down_expert_bytes; - max_expert_bytes = std::max(max_expert_bytes, per_expert_bytes); - } - if (max_expert_bytes == 0) { - std::fprintf(stderr, "[deepseek4] failed to compute streaming expert size\n"); - return false; - } - if (!stream_engine_.init(backend_, max_expert_bytes, &err)) { - std::fprintf(stderr, "[deepseek4] failed to init cold-expert stream engine: %s\n", - err.c_str()); - return false; - } - std::fprintf(stderr, - "[deepseek4] cold-expert stream engine ready: pinned=%.1f MiB scratch=%.1f MiB\n", - stream_engine_.pinned_bytes() / 1024.0 / 1024.0, - stream_engine_.scratch_bytes() / 1024.0 / 1024.0); - } - moe_hybrid_ = std::move(hybrid); + if (!init_single_target_expert_runtime(&err)) { + std::fprintf(stderr, "[deepseek4] failed to initialize hybrid expert runtime: %s\n", + err.c_str()); + return false; + } 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\n", + moe_placement_.total_hot, + total_cold, + cold_backend); return true; } @@ -464,105 +582,65 @@ 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 int requested_chunk = cfg_.chunk > 0 ? cfg_.chunk : 512; + const int chunk = moe_hybrid_ ? 1 : requested_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=%s intermediate_logits=%s backend=%s\n", + n_total, + requested_chunk, + chunk, + moe_hybrid_ ? "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; } @@ -570,12 +648,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; } @@ -589,10 +674,18 @@ 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 bool disable_cached_decode = false; 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) { @@ -607,24 +700,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; } @@ -648,10 +748,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; } @@ -665,10 +776,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; @@ -676,6 +789,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; @@ -687,12 +801,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); @@ -757,8 +872,11 @@ void DeepSeek4Backend::shutdown() { free_deepseek4_snapshot(snapshots_[i]); } free_deepseek4_cache(cache_); - stream_engine_.destroy(); moe_hybrid_.reset(); + expert_runtime_.reset(); + 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..5c26d21b5 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -11,7 +11,7 @@ #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/moe_expert_compute.h" #include "deepseek4_internal.h" #include "ggml.h" @@ -60,6 +60,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 +83,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 +102,13 @@ class DeepSeek4Backend : public ModelBackend { int max_ctx, MoeHybridPlacement & out, std::string * err) const; + bool init_single_target_expert_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. + MoeExpertComputeRuntime expert_runtime_; + std::vector layer_expert_bytes_; 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 06fbe8a98..68687005d 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 @@ -35,17 +38,73 @@ 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 + // 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 +114,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; @@ -97,8 +156,59 @@ 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 +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; @@ -137,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; @@ -154,6 +267,68 @@ struct DeepSeek4CachedDecodeFfnGraph { layer_idx = -1; n_tokens = 0; hash_routed = false; + estimated_bytes = 0; + last_used_epoch = 0; + reuse_count = 0; + } +}; + +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; } }; @@ -176,6 +351,22 @@ struct DeepSeek4CachedDecodeOutputGraph { } }; +struct DeepSeek4CachedDecodeGpuFfnState { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + GpuResidentState state; + + bool valid() const { + return owner_ctx && backend && state.valid(); + } + + void free() { + state.destroy(); + owner_ctx = nullptr; + backend = nullptr; + } +}; + struct DeepSeek4AttentionGraphInputs { ggml_tensor * rope_pos = nullptr; ggml_tensor * neg_pos = nullptr; @@ -198,8 +389,10 @@ 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; + bool uses_shared_inputs = false; StepGraph sg; DeepSeek4AttentionGraphInputs inputs; @@ -224,8 +417,10 @@ struct DeepSeek4CachedDecodeAttnGraph { n_raw = 0; n_comp_attn = 0; n_index_comp = 0; + flush_boundary = false; compressed = false; indexed = false; + uses_shared_inputs = false; } }; @@ -248,12 +443,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; @@ -262,18 +470,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); @@ -284,10 +499,131 @@ 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); } }; +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, @@ -310,7 +646,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; @@ -369,6 +706,93 @@ static bool build_cached_decode_ffn_graph( 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; + 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; +} + +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; @@ -394,12 +818,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)); @@ -414,6 +860,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, @@ -1061,15 +1580,18 @@ 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 * pre = nullptr; ggml_tensor * post = nullptr; ggml_tensor * comb = 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 && - post && comb; + pre && post && comb; } void free() { @@ -1077,7 +1599,9 @@ struct DeepSeek4CachedDecodeHcPreGraph { owner_ctx = nullptr; backend = nullptr; layer_idx = -1; + n_tokens = 0; ffn = false; + pre = nullptr; post = nullptr; comb = nullptr; } @@ -1086,6 +1610,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; @@ -1094,6 +1619,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; } @@ -1102,6 +1628,7 @@ struct DeepSeek4CachedDecodeHcPostGraph { step_graph_destroy(sg); owner_ctx = nullptr; backend = nullptr; + n_tokens = 0; residual_hc = nullptr; block_out = nullptr; post = nullptr; @@ -1109,67 +1636,221 @@ struct DeepSeek4CachedDecodeHcPostGraph { } }; -static bool build_cached_decode_attn_graph( - DeepSeek4CachedDecodeAttnGraph & out, - ggml_backend_t backend, - const DeepSeek4Weights & w, - const DeepSeek4Layer & L, - DeepSeek4LayerCache & lc, - int layer_idx, - int kv_start, - int raw_attn_count, - int comp_attn_count, - int index_comp_count) { - out.free(); - - 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; - out.sg.ctx = ggml_init(params); - if (!out.sg.ctx) { - return false; - } +// 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, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + DeepSeek4LayerCache & lc, + int layer_idx, + int kv_start, + int raw_attn_count, + int comp_attn_count, + int index_comp_count, + bool flush_boundary, + const Ds4DecodeSharedInputs * shared_inputs = nullptr) { + out.free(); + + 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; + out.sg.ctx = ggml_init(params); + if (!out.sg.ctx) { + return false; + } const int ratio = w.compress_ratios[layer_idx]; out.n_tokens = 1; 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.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); + out.sg.gf = ggml_new_graph_custom( + out.sg.ctx, DS4_CACHED_DECODE_ATTN_GRAPH_SIZE, false); + + 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; @@ -1289,36 +1970,38 @@ 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 && !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; @@ -1340,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, @@ -1347,7 +2145,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; @@ -1364,20 +2163,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); @@ -1390,17 +2192,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) { @@ -1408,17 +2210,53 @@ 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 * 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); @@ -1431,6 +2269,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; } @@ -1438,7 +2277,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; @@ -1452,29 +2292,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); @@ -1487,6 +2328,7 @@ static bool build_cached_decode_hc_post_graph( out.owner_ctx = w.ctx; out.backend = backend; + out.n_tokens = n_tokens; return true; } @@ -1507,6 +2349,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; } @@ -1534,140 +2378,333 @@ 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); + 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); + 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; + } - 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; - } + 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) { + 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; + } - float * dst = ffn_out_host.data() + (size_t)ti * (size_t)n_embd; + 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, + 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, 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] = hot_out[(size_t)i] + cold_out[(size_t)i]; + dst[i] += worker_out[(size_t)i]; } } - if (step_tel) step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); + } + 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; } - 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); - } + 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; + } + (void) selected_weights_gpu; + if (!enabled || !ffn_post_gpu) { return true; } - 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, - 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)) { + 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, + 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); + } + 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; } - 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; } @@ -1801,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; }; @@ -1811,7 +2850,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]; @@ -1819,23 +2950,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); }); } @@ -2072,24 +3223,149 @@ 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); + } + }); +} + +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, @@ -2178,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; @@ -2194,7 +3475,12 @@ 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 want_logits, + bool disable_cached_decode, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layers, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats) { const auto step_t0 = Ds4TimingClock::now(); @@ -2219,239 +3505,1038 @@ 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; + const bool use_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_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_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 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_backend_decode_hc; + 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 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; + 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(); + } } - - // ── 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]; + 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_shared_inputs.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; + } + 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; + 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->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()); - const auto route_compute_t0 = Ds4TimingClock::now(); - 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); - 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 (telemetry) telemetry->route_read_us += ds4_elapsed_us(route_read_t0, Ds4TimingClock::now()); + 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()); } - ggml_gallocr_free(ffn_alloc); - ggml_free(ffn_ctx); - if (!ok) { + 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; } - - 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) { - 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 (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 && + candidate.uses_shared_inputs == (decode_shared_inputs_ptr != nullptr); + }); + 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, + decode_shared_inputs_ptr)) { 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; - 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]; - } - 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->attn_build_us += ds4_elapsed_us(attn_build_t0, Ds4TimingClock::now()); + it = std::prev(per_layer.end()); } - if (telemetry) telemetry->route_select_us += ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); - if (routing_stats) { + cached_attn = &*it; + gf = cached_attn->sg.gf; + attn_out = cached_attn->sg.hidden_states; + + 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()); + } + 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)); + } + } + } + } 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]; + for (int ti = 0; ti < n_tokens; ++ti) { + const int32_t tok = token_ids[ti]; + 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.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); + } + 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); + 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()); + 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(); + 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()); + 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; + if (ok) { + const auto route_read_t0 = Ds4TimingClock::now(); + 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); + ggml_free(ffn_ctx); + if (!ok) { + if (hot_alloc) ggml_gallocr_free(hot_alloc); + if (cold_alloc) ggml_gallocr_free(cold_alloc); + return false; + } + + const auto route_select_t0 = Ds4TimingClock::now(); + 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 >= 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; + } + } + 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.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]; + 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; + } + } + }); + 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); } } + 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]; - if (!eval_ds4_hybrid( - backend, cpu_backend, hybrid_cfg, desc, &moe_hybrid, storage, stream_engine, + 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); + 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, + 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, - 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; @@ -2485,9 +4570,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); @@ -2496,76 +4585,95 @@ 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; 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; @@ -2574,83 +4682,331 @@ 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 (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); + + 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; + + 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()); + } + + // ── Output HC pre → norm → logits ─────────────────────────────────── + const auto output_t0 = Ds4TimingClock::now(); + 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); + } 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)); + } + } + + 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; + } + } + 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 { + 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); + } + if (!final_ok) return false; + if (telemetry) { + telemetry->output_us += ds4_elapsed_us(output_t0, Ds4TimingClock::now()); + telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + } + + cache.cur_pos = kv_start + n_tokens; + 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; + } } - if (hot_alloc) ggml_gallocr_free(hot_alloc); - if (cold_alloc) ggml_gallocr_free(cold_alloc); + out.cur_pos = cache.cur_pos; + return true; +} - // ── 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; - } - 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]; - } - final_embd[d] = acc; - } - } 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)) { - ggml_gallocr_free(final_alloc); - ggml_free(ctx2); +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; } - 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); + if (!tensors_compatible(snap.hc_state_snap, cache.hc_state) || + !copy_tensor_to_backend(snap.hc_state_snap, cache.hc_state)) { + return false; } - 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()); - telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + + 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 = kv_start + n_tokens; + cache.cur_pos = snap.cur_pos; return true; } @@ -2666,35 +5022,122 @@ 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 want_logits, + 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, + want_logits ? &out_logits : nullptr, + 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, + want_logits, + 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 && + 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 (cached_sg->meta_arena.size() < ctx_size) { + cached_sg->meta_arena.resize(ctx_size); + } + } + ggml_init_params params{}; - params.mem_size = ctx_size; - params.mem_buffer = 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); if (!ctx) return false; + 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); + } else { + 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); + ggml_cgraph * gf = + ggml_new_graph_custom(ctx, ds4_full_prefill_graph_size(n_tokens), false); + if (cached_sg) { + cached_sg->gf = gf; + } std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; @@ -2733,30 +5176,47 @@ 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); + 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 - 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 +5229,34 @@ 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) - 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 (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()); + } + } - 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,9 +5272,35 @@ bool deepseek4_step( } cache.cur_pos = next_pos; + if (telemetry) { + telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + } 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, + bool want_logits, + 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, want_logits, false, telemetry, + routing_stats, expert_compute, expert_layers); +} + bool deepseek4_step_layer_range( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -2856,11 +5353,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); @@ -2877,18 +5441,44 @@ 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(); } } 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(); + 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); @@ -2903,7 +5493,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; @@ -2913,24 +5503,149 @@ 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_backend_decode_hc_direct = use_backend_decode_hc && ds4_backend_is_hip(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; + 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]; @@ -2947,8 +5662,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() || @@ -2970,7 +5692,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, @@ -2987,38 +5709,82 @@ 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); } 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 ───────────────────────────── { @@ -3033,6 +5799,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) { @@ -3042,7 +5815,9 @@ 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 && + candidate.uses_shared_inputs == (decode_shared_inputs_ptr != nullptr); }); if (it == per_layer.end()) { if (per_layer.size() >= 20) { @@ -3053,7 +5828,9 @@ 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, + decode_shared_inputs_ptr)) { std::fprintf(stderr, "[deepseek4] cached attn graph alloc failed layer %d\n", il); return false; } @@ -3064,43 +5841,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 + 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) { - 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 + 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) { - 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 { @@ -3162,35 +5939,81 @@ 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(); - 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 auto attn_read_t0 = Ds4TimingClock::now(); + if (attn_post_backend && attn_comb_backend) { + 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()); } 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(), @@ -3200,16 +6023,26 @@ 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()); } } + ds4_trace_decode_marker(trace_decode, kv_start, il, "attn_done"); // ── HC pre (FFN) ──────────────────────────────────────────── 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() || @@ -3231,7 +6064,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, @@ -3248,33 +6081,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); @@ -3289,79 +6165,207 @@ 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]; - 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); + 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 && split_layer_range_route_expert) { + 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()); - } + 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); + } + } - 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) { - 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(); - 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()); + if (ffn_post_backend && ffn_comb_backend) { + 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()); } // ── 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(), @@ -3371,13 +6375,22 @@ 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()); } } } - 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) && + !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) ──────── @@ -3385,13 +6398,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() || @@ -3402,8 +6417,23 @@ 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_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()); + } if (ggml_backend_graph_compute(backend, cached_decode_output_graph.sg.gf) != GGML_STATUS_SUCCESS) { return false; } @@ -3564,196 +6594,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_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index 8f784d1f2..2dce8e712 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -1,13 +1,19 @@ #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 #include namespace dflash::common { @@ -22,48 +28,78 @@ 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; } }; 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) { @@ -71,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; @@ -111,49 +152,100 @@ __global__ void hc_mix_kernel(const float * x, if (tid == 0) mix[row] = smem[0]; } -__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_finalize_inv_rms_kernel(const float * sums, + int n_sums, + int cols, + float eps, + float * inv_rms_out) { __shared__ float smem[kThreads]; - __shared__ float inv_rms; - const int row = blockIdx.x; const int tid = threadIdx.x; - if (row >= rows) { - return; + 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); } +} - float sumsq = 0.0f; - for (int c = tid; c < cols; c += blockDim.x) { - const float v = x[c]; - sumsq += v * v; +__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] = sumsq; + 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 = rsqrtf(smem[0] / (float) cols + eps); + 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, + 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; - float dot = 0.0f; for (int c = tid; c < cols; c += blockDim.x) { - dot += __half2float(w[c]) * x[c]; + acc += __half2float(w[c]) * (x[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) { @@ -161,6 +253,40 @@ __global__ void hc_mix_norm_kernel(const float * x, } } +__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]; + const int row = blockIdx.x; + const int token = blockIdx.y; + const int tid = threadIdx.x; + 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; + for (int c = tid; c < cols; c += blockDim.x) { + acc += __half2float(w[c]) * (x_token[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[(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, @@ -246,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, @@ -257,6 +520,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,82 +532,174 @@ 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 = stream ? cudaStreamSynchronize(stream) : 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; +} + +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); } - } - 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); + return false; + }; + if (mix_dim > kMaxMixDim) { + if (log_errors) { + std::fprintf(stderr, + "[deepseek4-hc-direct] batch mix dim too large: %d > %d\n", + mix_dim, kMaxMixDim); } + return false; } - err = cudaDeviceSynchronize(); + 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("post-copy sync", err); + 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, @@ -355,30 +713,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 +767,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 +841,8 @@ 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, @@ -492,6 +856,55 @@ bool deepseek4_cuda_hc_pre_device(const void * hc_state_device, 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, + fn_device, + scale_device, + base_device, + n_embd, + n_hc, + sinkhorn_iters, + eps, + working_device, + post_device, + comb_device, + scratch, + stream, + true, false); } @@ -518,29 +931,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(scratch.d_base, base_host, sizeof(float) * (size_t) mix_dim, + cudaMemcpyHostToDevice) != cudaSuccess) { + hc_log_cuda_error("copy base", cudaGetLastError()); return false; } - if (cudaMemcpy(g_scratch.d_scale, scale_host, sizeof(float) * (size_t) mix_dim, + 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,7 +1014,170 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, 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_internal.h b/server/src/deepseek4/deepseek4_internal.h index 31c086394..f13612255 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -30,11 +30,18 @@ namespace dflash::common { struct MoeHybridPlacement; struct MoeHybridConfig; struct MoeHybridRoutingStats; -class MoeHybridStreamEngine; +struct MoeExpertLayer; +struct MoeExpertCompute; +class ExpertIpcClient; 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; @@ -56,10 +63,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; @@ -214,7 +245,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; }; @@ -247,6 +278,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 { @@ -303,7 +338,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, @@ -314,9 +350,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, @@ -349,15 +410,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; 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/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 diff --git a/server/src/server/http_server.cpp b/server/src/server/http_server.cpp index 8d4bdab24..b03f8919b 100644 --- a/server/src/server/http_server.cpp +++ b/server/src/server/http_server.cpp @@ -108,6 +108,86 @@ 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; + } + + 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 +2835,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 +2877,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 +2893,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 +2906,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 +2938,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 +3092,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 +3120,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);