diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 1636fff02..35de5c388 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1,44 +1,12 @@ cmake_minimum_required(VERSION 3.21) # 3.21 adds first-class HIP language support (project(LANGUAGES ... HIP)) set(DFLASH27B_GPU_BACKEND "cuda" CACHE STRING "GPU backend to build: cuda or hip") set_property(CACHE DFLASH27B_GPU_BACKEND PROPERTY STRINGS cuda hip) -option(DFLASH27B_ENABLE_MIXED_CUDA_HIP - "Primary GPU runtime plus an isolated in-process CUDA/HIP peer (Linux only)" - OFF) string(TOLOWER "${DFLASH27B_GPU_BACKEND}" DFLASH27B_GPU_BACKEND) -if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - # Reject unsupported targets before project() tries to discover both GPU - # toolchains. Toolchain files define CMAKE_SYSTEM_NAME for cross builds; - # otherwise the host system is the native target. - if(CMAKE_SYSTEM_NAME) - set(_dflash_mixed_target_system "${CMAKE_SYSTEM_NAME}") - else() - set(_dflash_mixed_target_system "${CMAKE_HOST_SYSTEM_NAME}") - endif() - if(NOT _dflash_mixed_target_system STREQUAL "Linux") - message(FATAL_ERROR - "The in-process CUDA+HIP runtime is currently supported on Linux only") - endif() - unset(_dflash_mixed_target_system) -endif() -# These are internal build-shape switches rather than user-selectable backend -# policy. Reset both on every configure so reusing a build directory after -# changing the primary backend cannot leave the old peer as a module. -set(GGML_CUDA_MODULE OFF CACHE BOOL "Build CUDA as a runtime-loadable module" FORCE) -set(GGML_HIP_MODULE OFF CACHE BOOL "Build HIP as a runtime-loadable module" FORCE) - if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") set(DFLASH27B_USER_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}") project(dflash LANGUAGES C CXX CUDA) elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - set(DFLASH27B_CUDA_ARCHITECTURES "86" CACHE STRING - "Secondary CUDA GPU targets, e.g. 86 for RTX 3090") - set(CMAKE_CUDA_ARCHITECTURES "${DFLASH27B_CUDA_ARCHITECTURES}" - CACHE STRING "" FORCE) - project(dflash LANGUAGES C CXX HIP CUDA) - else() - project(dflash LANGUAGES C CXX HIP) - endif() + project(dflash LANGUAGES C CXX HIP) else() message(FATAL_ERROR "DFLASH27B_GPU_BACKEND must be 'cuda' or 'hip', got '${DFLASH27B_GPU_BACKEND}'") endif() @@ -74,9 +42,8 @@ endif() # If we do not set this, ggml will output to bin and DLLs will not load set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}") -# ROCm root for HIP and mixed CUDA+HIP builds (rpath + header discovery). -if(DFLASH27B_GPU_BACKEND STREQUAL "hip" OR - DFLASH27B_ENABLE_MIXED_CUDA_HIP) +# ROCm root for HIP builds (rpath + rocwmma header discovery). +if(DFLASH27B_GPU_BACKEND STREQUAL "hip") if(DEFINED ROCM_PATH) set(_dflash_rocm_root "${ROCM_PATH}") elseif(DEFINED ENV{ROCM_PATH}) @@ -90,8 +57,7 @@ endif() # Bake portable rpath into all executables so bundled ggml backend libs / libggml-base # are found regardless of LD_LIBRARY_PATH or stale /usr/local/lib (closes #31). set(CMAKE_INSTALL_RPATH "$ORIGIN/deps/llama.cpp/ggml/src;$ORIGIN/deps/llama.cpp/ggml/src/ggml-cuda;$ORIGIN/deps/llama.cpp/ggml/src/ggml-hip;$ORIGIN/../deps/llama.cpp/ggml/src;$ORIGIN/../deps/llama.cpp/ggml/src/ggml-cuda;$ORIGIN/../deps/llama.cpp/ggml/src/ggml-hip") -if((DFLASH27B_GPU_BACKEND STREQUAL "hip" OR - DFLASH27B_ENABLE_MIXED_CUDA_HIP) AND _dflash_rocm_root) +if(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND _dflash_rocm_root) list(APPEND CMAKE_BUILD_RPATH "${_dflash_rocm_root}/lib" "${_dflash_rocm_root}/lib64") set(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_RPATH};${_dflash_rocm_root}/lib;${_dflash_rocm_root}/lib64") endif() @@ -112,52 +78,20 @@ endif() if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") set(GGML_CUDA ON CACHE BOOL "" FORCE) - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - set(GGML_HIP ON CACHE BOOL "" FORCE) - set(GGML_HIP_RCCL OFF CACHE BOOL "" FORCE) - # Keep CUDA/CPU linked normally so every existing code path is - # unchanged. Build only HIP as an RTLD_LOCAL module: both GPU - # implementations use the historical ggml_backend_cuda_* symbol - # family and therefore cannot safely share the global symbol scope. - set(GGML_BACKEND_DL OFF CACHE BOOL "" FORCE) - set(GGML_HIP_MODULE ON CACHE BOOL - "Build HIP as an isolated runtime-loadable backend" FORCE) - # The peer module must link to a shared ggml core. Keep this as a - # scoped build requirement; do not overwrite the user's cached - # BUILD_SHARED_LIBS choice for later non-mixed reconfiguration. - set(DFLASH27B_MIXED_GGML_SHARED ON) - set(DFLASH27B_GGML_BACKEND_TARGET ggml-cuda) - set(DFLASH27B_HIP_ARCHITECTURES "" CACHE STRING - "Secondary HIP GPU targets, e.g. gfx1151") - else() - set(GGML_HIP OFF CACHE BOOL "" FORCE) - set(GGML_BACKEND_DL OFF CACHE BOOL "" FORCE) - set(DFLASH27B_GGML_BACKEND_TARGET ggml-cuda) - endif() + set(GGML_HIP OFF CACHE BOOL "" FORCE) set(GGML_CUDA_GRAPHS ON CACHE BOOL "Enable CUDA graphs for AR-decode replay (lucebox)" FORCE) + set(DFLASH27B_GGML_BACKEND_TARGET ggml-cuda) elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - set(GGML_CUDA ON CACHE BOOL "" FORCE) - set(GGML_CUDA_MODULE ON CACHE BOOL - "Build CUDA as an isolated runtime-loadable backend" FORCE) - set(DFLASH27B_MIXED_GGML_SHARED ON) - else() - set(GGML_CUDA OFF CACHE BOOL "" FORCE) - set(GGML_CUDA_MODULE OFF CACHE BOOL "" FORCE) - endif() + set(GGML_CUDA OFF CACHE BOOL "" FORCE) set(GGML_HIP ON CACHE BOOL "" FORCE) set(GGML_HIP_RCCL OFF CACHE BOOL "" FORCE) set(DFLASH27B_GGML_BACKEND_TARGET ggml-hip) set(DFLASH27B_HIP_ARCHITECTURES "" CACHE STRING "HIP GPU targets, e.g. gfx906;gfx1100") - set(GGML_BACKEND_DL OFF CACHE BOOL "" FORCE) if(DFLASH27B_HIP_ARCHITECTURES AND NOT CMAKE_HIP_ARCHITECTURES) set(CMAKE_HIP_ARCHITECTURES "${DFLASH27B_HIP_ARCHITECTURES}" CACHE STRING "" FORCE) endif() - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - set(GGML_CUDA_GRAPHS ON CACHE BOOL - "Enable CUDA graphs in the secondary CUDA backend" FORCE) - endif() endif() +set(GGML_BACKEND_DL OFF CACHE BOOL "" FORCE) set(GGML_METAL OFF CACHE BOOL "" FORCE) set(GGML_VULKAN OFF CACHE BOOL "" FORCE) set(GGML_BLAS OFF CACHE BOOL "" FORCE) @@ -190,20 +124,14 @@ option(DFLASH27B_HIP_SM80_EQUIV # (defaulting to gfx1151 / Strix Halo) and pass it through to the HIP # toolchain and ggml backend. if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") - # Pascal (60–62) lacks tensor-core WMMA support and is excluded by default. - # Volta/Turing (70/75) support F16 WMMA (m16n8k16); Ampere (86)+ supports - # BF16 WMMA (m16n16k16). Add Pascal via -DDFLASH27B_USER_CUDA_ARCHITECTURES - # if needed (scalar fallback kernels compile for sm_60–69). + # Pascal (60/61/62), Volta (70), Turing (75) and Ampere (86) always; # Blackwell consumer (120) and Thor # (110 on CUDA 13+) added when nvcc supports them. DGX Spark / # GB10 is compute capability 12.1 (121), added at CUDA 12.9+. if(DFLASH27B_USER_CUDA_ARCHITECTURES) set(_dflash_archs "${DFLASH27B_USER_CUDA_ARCHITECTURES}") else() - # WMMA (warp matrix multiply) requires sm_70+. Pascal (60–62) lacks - # tensor cores and is excluded from the default arch list. Volta/Turing - # (70/75) support F16 WMMA; Ampere (86)+ supports BF16 WMMA. - set(_dflash_archs "70;75;86") + set(_dflash_archs "60;61;62;70;75;86") if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "12.8") list(APPEND _dflash_archs "120") endif() @@ -219,16 +147,6 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") # which triggers massive first-request PTX JIT on newer GPUs even though # dflash_common itself is compiled for the intended arch set. set(CMAKE_CUDA_ARCHITECTURES "${_dflash_archs}" CACHE STRING "" FORCE) - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - if(DFLASH27B_HIP_ARCHITECTURES) - set(_dflash_mixed_hip_archs "${DFLASH27B_HIP_ARCHITECTURES}") - elseif(AMDGPU_TARGETS) - set(_dflash_mixed_hip_archs "${AMDGPU_TARGETS}") - else() - set(_dflash_mixed_hip_archs "gfx1151") - endif() - set(CMAKE_HIP_ARCHITECTURES "${_dflash_mixed_hip_archs}" CACHE STRING "" FORCE) - endif() elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") # User override precedence: -DDFLASH27B_HIP_ARCHITECTURES → -DAMDGPU_TARGETS # → gfx1151 default (Strix Halo). @@ -282,44 +200,10 @@ if(WIN32 AND NOT CMAKE_ASM_COMPILER) set(CMAKE_ASM_COMPILER "${CMAKE_CXX_COMPILER}" CACHE FILEPATH "" FORCE) endif() -# Resolve GPU runtime packages before creating the ggml targets so their build -# rpaths include non-system toolkit installations. This matters when a peer -# module is loaded with dlopen and cannot inherit link-time search paths from -# the primary executable. -if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR - DFLASH27B_ENABLE_MIXED_CUDA_HIP) - find_package(CUDAToolkit REQUIRED) - if(UNIX AND CUDAToolkit_LIBRARY_DIR) - list(APPEND CMAKE_BUILD_RPATH "${CUDAToolkit_LIBRARY_DIR}") - list(APPEND CMAKE_INSTALL_RPATH "${CUDAToolkit_LIBRARY_DIR}") - endif() -endif() -# Use only the ggml subtree of llama.cpp (skip libllama). Mixed builds need a -# shared ggml core for the isolated peer module, but that requirement belongs -# only to this sub-build. Restoring the previous value prevents a reused build -# directory from silently changing ordinary builds after mixed mode is turned -# off. -if(DFLASH27B_MIXED_GGML_SHARED) - set(_dflash_build_shared_libs_was_defined OFF) - if(DEFINED BUILD_SHARED_LIBS) - set(_dflash_build_shared_libs_was_defined ON) - set(_dflash_saved_build_shared_libs "${BUILD_SHARED_LIBS}") - endif() - set(BUILD_SHARED_LIBS ON) -endif() +# Use only the ggml subtree of llama.cpp (skip libllama). add_subdirectory(deps/llama.cpp/ggml EXCLUDE_FROM_ALL) -if(DFLASH27B_MIXED_GGML_SHARED) - if(_dflash_build_shared_libs_was_defined) - set(BUILD_SHARED_LIBS "${_dflash_saved_build_shared_libs}") - else() - unset(BUILD_SHARED_LIBS) - endif() - unset(_dflash_saved_build_shared_libs) - unset(_dflash_build_shared_libs_was_defined) -endif() -if(DFLASH27B_GPU_BACKEND STREQUAL "hip" OR - DFLASH27B_ENABLE_MIXED_CUDA_HIP) +if(DFLASH27B_GPU_BACKEND STREQUAL "hip") # The vendored ggml HIP shim still uses a few CUDA spellings that are not # mapped in this upstream snapshot. Keep the compatibility layer in this # repo so the build stays reproducible from a clean checkout. @@ -333,22 +217,11 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip" OR ${CMAKE_CURRENT_SOURCE_DIR}/src/hip_compat) endif() -if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - # CUDA and HIP are compiled from the same ggml-cuda implementation and - # therefore share C++ typeinfo, vtable, template, and data-symbol names. - # Bind every definition in the secondary module locally: limiting this to - # functions still permits the primary runtime's pool/vtable state to - # interpose and hand an allocation to the wrong vendor kernel. - if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") - target_link_options(ggml-hip PRIVATE "LINKER:-Bsymbolic") - else() - target_link_options(ggml-cuda PRIVATE "LINKER:-Bsymbolic") - endif() -endif() - -if(DFLASH27B_GPU_BACKEND STREQUAL "hip" OR - DFLASH27B_ENABLE_MIXED_CUDA_HIP) - # The ggml HIP subdirectory establishes ROCm's package search roots. +if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + # The CUDA-only sources include directly, so the toolkit + # headers must be available when compiling the library. + find_package(CUDAToolkit REQUIRED) +elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") find_package(hip REQUIRED) endif() @@ -408,7 +281,6 @@ add_library(dflash_common STATIC src/laguna/laguna_layer_split_adapter.cpp src/laguna/laguna_dflash_target.cpp src/common/backend_ipc.cpp - src/common/dynamic_backend.cpp src/common/domino_head.cpp src/common/dspark_head.cpp src/common/target_shard_ipc.cpp @@ -515,19 +387,11 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") DFLASH27B_BACKEND_CUDA=1 DFLASH27B_CUDA_MIN_SM=${_dflash_cuda_min_sm} DFLASH27B_MIN_SM=${_dflash_cuda_min_sm}) - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - target_compile_definitions(dflash_common PRIVATE - DFLASH27B_BACKEND_MIXED=1) - endif() elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip") target_sources(dflash_common PRIVATE src/deepseek4/deepseek4_hc_cuda.cu) set_source_files_properties(src/deepseek4/deepseek4_hc_cuda.cu PROPERTIES LANGUAGE HIP) set_target_properties(dflash_common PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") target_compile_definitions(dflash_common PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - target_compile_definitions(dflash_common PRIVATE - DFLASH27B_BACKEND_MIXED=1) - endif() # 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) @@ -831,14 +695,6 @@ if(DFLASH27B_TESTS) add_test(NAME cuda_comm_api COMMAND test_cuda_comm_api) endif() - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - add_executable(test_mixed_cuda_hip test/test_mixed_cuda_hip.cpp) - target_include_directories(test_mixed_cuda_hip PRIVATE - ${DFLASH27B_SRC_INCLUDE_DIRS}) - target_link_libraries(test_mixed_cuda_hip PRIVATE dflash_common) - add_test(NAME mixed_cuda_hip COMMAND test_mixed_cuda_hip) - endif() - add_executable(test_rocmfp4 deps/llama.cpp/ggml/rocmfp4/test_rocmfp4.c) target_link_libraries(test_rocmfp4 PRIVATE ggml-base) if(UNIX) @@ -947,6 +803,8 @@ if(DFLASH27B_TESTS) endif() target_link_libraries(test_deepseek4_mmid_grouped_cuda PRIVATE ggml ${DFLASH27B_GGML_BACKEND_TARGET}) + target_include_directories(test_deepseek4_mmid_grouped_cuda PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/rocmfpx) list(APPEND _raw_unit_test_targets test_deepseek4_mmid_grouped_cuda) endif() # HIP-only standalone build; CUDA backend is covered by aggregated test_server_unit. @@ -1230,7 +1088,6 @@ if(DFLASH27B_TESTS) test/test_drafter_early_exit_score_range.cpp test/test_drafter_tail_capture_guard.cpp test/test_drafter_warm_path_regression.cpp - test/test_qwen3_buffer_plan.cpp test/test_gguf_mmap.cpp test/test_kv_quant.cpp test/test_kvflash_placement.cpp @@ -1453,18 +1310,10 @@ if(DFLASH27B_SERVER) target_include_directories(dflash_server PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) if(DFLASH27B_GPU_BACKEND STREQUAL "hip") target_compile_definitions(dflash_server PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - target_compile_definitions(dflash_server PRIVATE - DFLASH27B_BACKEND_MIXED=1) - endif() else() target_compile_definitions(dflash_server PRIVATE DFLASH27B_BACKEND_CUDA=1 DFLASH27B_CUDA_MIN_SM=${_dflash_cuda_min_sm}) - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - target_compile_definitions(dflash_server PRIVATE - DFLASH27B_BACKEND_MIXED=1) - endif() endif() if(NOT WIN32) target_link_libraries(dflash_server PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET} pthread) @@ -1505,18 +1354,10 @@ if(DFLASH27B_SERVER) target_include_directories(backend_ipc_daemon PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) if(DFLASH27B_GPU_BACKEND STREQUAL "hip") target_compile_definitions(backend_ipc_daemon PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - target_compile_definitions(backend_ipc_daemon PRIVATE - DFLASH27B_BACKEND_MIXED=1) - endif() else() target_compile_definitions(backend_ipc_daemon PRIVATE DFLASH27B_BACKEND_CUDA=1 DFLASH27B_CUDA_MIN_SM=${_dflash_cuda_min_sm}) - if(DFLASH27B_ENABLE_MIXED_CUDA_HIP) - target_compile_definitions(backend_ipc_daemon PRIVATE - DFLASH27B_BACKEND_MIXED=1) - endif() endif() if(NOT WIN32) target_link_libraries(backend_ipc_daemon PRIVATE dflash_common ggml ${DFLASH27B_GGML_BACKEND_TARGET} pthread) diff --git a/server/deps/llama.cpp/ggml/include/ggml-alloc.h b/server/deps/llama.cpp/ggml/include/ggml-alloc.h index 0a8979139..23f48128b 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-alloc.h +++ b/server/deps/llama.cpp/ggml/include/ggml-alloc.h @@ -46,11 +46,10 @@ GGML_API enum ggml_status ggml_tallocr_alloc(struct ggml_tallocr * talloc, st typedef struct ggml_gallocr * ggml_gallocr_t; GGML_API ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft); -// Uses max_chunk_size as the preferred backing-allocation limit while -// preserving a single logical graph allocator. Individual tensors are never -// split and may exceed the limit. This is useful on devices without virtual -// memory support, where a large contiguous allocation can fail despite -// sufficient aggregate free memory. +// Limits each backing allocation while preserving a single logical graph +// allocator. Tensors are never split across chunks. This is useful on devices +// without virtual memory support, where a large contiguous allocation can fail +// despite sufficient aggregate free memory. GGML_API ggml_gallocr_t ggml_gallocr_new_with_max_chunk_size( ggml_backend_buffer_type_t buft, size_t max_chunk_size); diff --git a/server/deps/llama.cpp/ggml/include/ggml-backend.h b/server/deps/llama.cpp/ggml/include/ggml-backend.h index e6bf95e5e..1187810fe 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-backend.h +++ b/server/deps/llama.cpp/ggml/include/ggml-backend.h @@ -385,6 +385,14 @@ extern "C" { ggml_backend_sched_t sched, bool enabled); + // In a non-pipelined scheduler, protect reusable split-input buffers with + // stream events instead of blocking the host on every destination stream. + // The source waits until the destination has finished the prior generation; + // the existing async-copy contract then publishes the new generation. + GGML_API void ggml_backend_sched_set_single_copy_event_fences( + ggml_backend_sched_t sched, + bool enabled); + // // Meta backend // diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index f5371f5af..116132801 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -677,6 +677,10 @@ extern "C" { GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed + // ...is a persistent fork payload whose cross-device copy may run on + // the destination stream. Transient allocator storage must never set + // this flag because the producer may otherwise reuse it too early. + GGML_TENSOR_FLAG_DST_STREAM_COPY = 32, }; enum ggml_tri_type { diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp index 174f6a45d..2cf4b4114 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp @@ -20,7 +20,6 @@ #include #include #include -#include #include #ifdef __APPLE__ @@ -820,11 +819,6 @@ struct ggml_backend_sched_deferred_peer_copy { struct ggml_backend_sched { bool is_reset; // true if the scheduler has been reset since the last graph split bool is_alloc; - // True only after every backend stream has been synchronized and before - // another graph is submitted. In that state scheduler copy destinations - // are already reusable, so queueing the same generation event before - // every split is redundant (and particularly expensive across HIP/CUDA). - bool backends_synchronized; int n_backends; @@ -882,27 +876,7 @@ struct ggml_backend_sched { int deferred_peer_copies_capacity; bool split_deferred_peer_copies; bool batch_split_copies; - ggml_backend_buffer_t batch_staging_buffers[GGML_SCHED_MAX_BACKENDS]; - uint8_t * batch_staging_bases[GGML_SCHED_MAX_BACKENDS]; - size_t batch_staging_sizes[GGML_SCHED_MAX_BACKENDS]; - bool profile_requested; - int profile_min_splits; - - struct { - bool active; - uint64_t loop_us; - uint64_t destination_wait_us; - uint64_t d2h_submit_us; - uint64_t host_relay_us; - uint64_t h2d_submit_us; - uint64_t compute_submit_us; - uint64_t staged_bytes; - int staged_copies; - uint64_t source_wait_us[GGML_SCHED_MAX_BACKENDS]; - uint64_t final_wait_us[GGML_SCHED_MAX_BACKENDS]; - int source_waits[GGML_SCHED_MAX_BACKENDS]; - int split_counts[GGML_SCHED_MAX_BACKENDS]; - } profile; + bool single_copy_event_fences; int debug; @@ -1775,557 +1749,18 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) { return true; } -struct ggml_backend_sched_tensor_payload_layout { - bool contiguous; - size_t run_bytes; - size_t runs_dim0; - size_t ne1; - size_t ne2; - size_t ne3; - size_t bytes; -}; - -static bool ggml_backend_sched_checked_mul( - size_t first, size_t second, size_t & result) { - if (first != 0 && second > SIZE_MAX / first) { - return false; - } - result = first * second; - return true; -} - -// Describe the logical payload independently of any padding or stride gaps. -// Contiguous tensors retain the one-copy fast path. Strided and permuted -// tensors are transferred as packed runs and restored at the same logical -// coordinates in the destination tensor. -static bool ggml_backend_sched_make_tensor_payload_layout( - const struct ggml_tensor * tensor, - ggml_backend_sched_tensor_payload_layout & layout) { - layout = {}; - if (!tensor) { - return false; - } - const size_t bytes = ggml_nbytes(tensor); - if (bytes == 0) { - layout.contiguous = true; - return true; - } - if (ggml_is_contiguous(tensor)) { - layout.contiguous = true; - layout.run_bytes = bytes; - layout.runs_dim0 = 1; - layout.ne1 = layout.ne2 = layout.ne3 = 1; - layout.bytes = bytes; - return true; - } - - const size_t type_size = ggml_type_size(tensor->type); - const int64_t block_size = ggml_blck_size(tensor->type); - if (type_size == 0 || block_size <= 0 || tensor->ne[0] <= 0 || - tensor->ne[0] % block_size != 0) { - return false; - } - layout.contiguous = false; - if (tensor->nb[0] == type_size) { - layout.run_bytes = ggml_row_size(tensor->type, tensor->ne[0]); - layout.runs_dim0 = 1; - } else { - // A permuted tensor is not contiguous even within dimension zero. - // Copy one storage block at a time while retaining its logical order. - layout.run_bytes = type_size; - layout.runs_dim0 = (size_t) (tensor->ne[0] / block_size); - } - layout.ne1 = (size_t) tensor->ne[1]; - layout.ne2 = (size_t) tensor->ne[2]; - layout.ne3 = (size_t) tensor->ne[3]; - - size_t runs = layout.runs_dim0; - if (!ggml_backend_sched_checked_mul(runs, layout.ne1, runs) || - !ggml_backend_sched_checked_mul(runs, layout.ne2, runs) || - !ggml_backend_sched_checked_mul(runs, layout.ne3, runs) || - !ggml_backend_sched_checked_mul(runs, layout.run_bytes, layout.bytes)) { - return false; - } - return true; -} - -static bool ggml_backend_sched_reserve_staging( - size_t & cursor, size_t bytes, size_t * offset = nullptr) { - constexpr size_t staging_alignment = 64; - if (cursor > SIZE_MAX - (staging_alignment - 1)) { - return false; - } - cursor = (cursor + staging_alignment - 1) & - ~(staging_alignment - 1); - if (bytes > SIZE_MAX - cursor) { - return false; - } - if (offset) { - *offset = cursor; - } - cursor += bytes; - return true; -} - -// The fast staging path is reserved for unlike non-host runtimes. Native -// peers retain their backend copy implementation, while host-backed tensors -// use the established direct fallback without allocating redundant arenas. -static bool ggml_backend_sched_batch_staging_candidate( - ggml_backend_sched_t sched, - int source_backend_id, - int destination_backend_id, - const struct ggml_tensor * source) { - if (!sched || !source || source_backend_id < 0 || - destination_backend_id < 0 || - source_backend_id >= sched->n_backends || - destination_backend_id >= sched->n_backends || - source_backend_id == destination_backend_id) { - return false; - } - ggml_backend_t source_backend = sched->backends[source_backend_id]; - ggml_backend_t destination_backend = - sched->backends[destination_backend_id]; - if (!source_backend->iface.get_tensor_async || - !destination_backend->iface.set_tensor_async || - ggml_guid_matches(ggml_backend_guid(source_backend), - ggml_backend_guid(destination_backend))) { - return false; - } - ggml_backend_buffer_t source_buffer = source->view_src - ? source->view_src->buffer : source->buffer; - const bool source_is_host = source_buffer - ? ggml_backend_buffer_is_host(source_buffer) - : ggml_backend_buft_is_host(sched->bufts[source_backend_id]); - if (source_is_host) { - return false; - } - ggml_backend_dev_t destination_device = - ggml_backend_get_device(destination_backend); - return destination_device && - ggml_backend_dev_type(destination_device) != - GGML_BACKEND_DEVICE_TYPE_CPU; -} - -static bool ggml_backend_sched_prepare_batch_staging( - ggml_backend_sched_t sched) { - if (!sched->batch_split_copies) { - return true; - } - - size_t required[GGML_SCHED_MAX_BACKENDS] = {}; - for (int split_id = 0; split_id < sched->n_splits; ++split_id) { - const struct ggml_backend_sched_split & split = - sched->splits[split_id]; - for (int input_id = 0; input_id < split.n_inputs; ++input_id) { - struct ggml_tensor * input = split.inputs[input_id]; - if (input->flags & GGML_TENSOR_FLAG_INPUT) { - continue; - } - ggml_backend_t source_backend = - ggml_backend_sched_get_tensor_backend(sched, input); - const int source_backend_id = - ggml_backend_sched_backend_id(sched, source_backend); - if (!ggml_backend_sched_batch_staging_candidate( - sched, source_backend_id, split.backend_id, input)) { - continue; - } - ggml_backend_sched_tensor_payload_layout layout; - if (!ggml_backend_sched_make_tensor_payload_layout(input, layout)) { - // Unsupported exotic layouts retain the established blocking - // tensor-copy fallback instead of making graph allocation fail. - continue; - } - if (!ggml_backend_sched_reserve_staging( - required[source_backend_id], layout.bytes) || - !ggml_backend_sched_reserve_staging( - required[split.backend_id], layout.bytes)) { - return false; - } - } - } - - ggml_backend_buffer_t new_buffers[GGML_SCHED_MAX_BACKENDS] = {}; - uint8_t * new_bases[GGML_SCHED_MAX_BACKENDS] = {}; - bool replace_any = false; - for (int i = 0; i < sched->n_backends; ++i) { - if (required[i] == 0 || - (sched->batch_staging_bases[i] && - required[i] <= sched->batch_staging_sizes[i])) { - continue; - } - replace_any = true; - ggml_backend_dev_t device = ggml_backend_get_device(sched->backends[i]); - ggml_backend_buffer_type_t host_buft = device - ? ggml_backend_dev_host_buffer_type(device) : nullptr; - if (host_buft) { - new_buffers[i] = ggml_backend_buft_alloc_buffer( - host_buft, required[i]); - if (new_buffers[i]) { - new_bases[i] = static_cast( - ggml_backend_buffer_get_base(new_buffers[i])); - } - } - if (!new_bases[i]) { - if (new_buffers[i]) { - ggml_backend_buffer_free(new_buffers[i]); - new_buffers[i] = nullptr; - } - new_bases[i] = static_cast(malloc(required[i])); - } - if (!new_bases[i]) { - for (int j = 0; j < sched->n_backends; ++j) { - if (new_buffers[j]) { - ggml_backend_buffer_free(new_buffers[j]); - } else { - free(new_bases[j]); - } - } - return false; - } - } - if (!replace_any) { - return true; - } - - // A graph may be reallocated after an earlier execution. Quiesce every - // runtime before releasing host pages that an asynchronous H2D transfer - // could still reference. - for (int i = 0; i < sched->n_backends; ++i) { - ggml_backend_synchronize(sched->backends[i]); - } - sched->backends_synchronized = true; - for (int i = 0; i < sched->n_backends; ++i) { - if (!new_bases[i]) { - continue; - } - if (sched->batch_staging_buffers[i]) { - ggml_backend_buffer_free(sched->batch_staging_buffers[i]); - } else { - free(sched->batch_staging_bases[i]); - } - sched->batch_staging_buffers[i] = new_buffers[i]; - sched->batch_staging_bases[i] = new_bases[i]; - sched->batch_staging_sizes[i] = required[i]; - GGML_LOG_DEBUG( - "%s: backend=%s size=%zu memory=%s\n", __func__, - ggml_backend_name(sched->backends[i]), required[i], - sched->batch_staging_buffers[i] ? "pinned" : "pageable"); - } - return true; -} - -static void ggml_backend_sched_free_batch_staging( - ggml_backend_sched_t sched) { - if (!sched) { - return; - } - for (int i = 0; i < sched->n_backends; ++i) { - if (sched->batch_staging_buffers[i]) { - ggml_backend_buffer_free(sched->batch_staging_buffers[i]); - } else { - free(sched->batch_staging_bases[i]); - } - sched->batch_staging_buffers[i] = nullptr; - sched->batch_staging_bases[i] = nullptr; - sched->batch_staging_sizes[i] = 0; - } -} - -using ggml_backend_sched_profile_clock = std::chrono::steady_clock; - -static uint64_t ggml_backend_sched_elapsed_us( - ggml_backend_sched_profile_clock::time_point start, - ggml_backend_sched_profile_clock::time_point end) { - return (uint64_t) std::chrono::duration_cast( - end - start).count(); -} - -struct ggml_backend_sched_staged_copy { - ggml_backend_t source_backend; - int source_backend_id; - struct ggml_tensor * destination; - ggml_backend_sched_tensor_payload_layout layout; - size_t source_offset; - size_t destination_offset; -}; - -template -static void ggml_backend_sched_for_each_payload_run( - const struct ggml_tensor * tensor, - const ggml_backend_sched_tensor_payload_layout & layout, - CopyRun && copy_run) { - if (layout.bytes == 0) { - return; - } - if (layout.contiguous) { - copy_run(/*tensor_offset=*/0, /*packed_offset=*/0, layout.bytes); - return; - } - - const size_t tensor_span = ggml_nbytes(tensor); - size_t packed_offset = 0; - for (size_t i3 = 0; i3 < layout.ne3; ++i3) { - for (size_t i2 = 0; i2 < layout.ne2; ++i2) { - for (size_t i1 = 0; i1 < layout.ne1; ++i1) { - for (size_t i0 = 0; i0 < layout.runs_dim0; ++i0) { - const size_t tensor_offset = - i0 * tensor->nb[0] + i1 * tensor->nb[1] + - i2 * tensor->nb[2] + i3 * tensor->nb[3]; - GGML_ASSERT(tensor_offset <= tensor_span && - layout.run_bytes <= tensor_span - tensor_offset); - copy_run(tensor_offset, packed_offset, layout.run_bytes); - packed_offset += layout.run_bytes; - } - } - } - } - GGML_ASSERT(packed_offset == layout.bytes); -} - -static void ggml_backend_sched_get_payload_async( - ggml_backend_t backend, - const struct ggml_tensor * tensor, - uint8_t * packed, - const ggml_backend_sched_tensor_payload_layout & layout) { - GGML_ASSERT(backend && backend->iface.get_tensor_async && - tensor && tensor->data && packed); - ggml_backend_sched_for_each_payload_run( - tensor, layout, - [&](size_t tensor_offset, size_t packed_offset, size_t bytes) { - backend->iface.get_tensor_async( - backend, tensor, packed + packed_offset, - tensor_offset, bytes); - }); -} - -static void ggml_backend_sched_set_payload_async( - ggml_backend_t backend, - struct ggml_tensor * tensor, - const uint8_t * packed, - const ggml_backend_sched_tensor_payload_layout & layout) { - GGML_ASSERT(backend && backend->iface.set_tensor_async && - tensor && tensor->data && packed); - ggml_backend_sched_for_each_payload_run( - tensor, layout, - [&](size_t tensor_offset, size_t packed_offset, size_t bytes) { - backend->iface.set_tensor_async( - backend, tensor, packed + packed_offset, - tensor_offset, bytes); - }); -} - -// Owns copy synchronization and host-staging state for one scheduler split. -// Keeping this policy separate from graph submission makes the execution loop -// readable and gives every early/fallback path one synchronization contract. -struct ggml_backend_sched_split_copy_state { - ggml_backend_sched_t sched; - int destination_backend_id; - ggml_backend_t destination_backend; - bool staging_arena_reusable; - bool destination_generation_ready; - size_t * staging_cursors; - bool destination_host_ready = false; - ggml_backend_t synchronized_sources[GGML_SCHED_MAX_BACKENDS] = {}; - int n_synchronized_sources = 0; - ggml_backend_sched_staged_copy staged[GGML_SCHED_MAX_SPLIT_INPUTS] = {}; - int n_staged = 0; - - bool profiling() const { - return sched->profile.active; - } - - void mark_destination_ready() { - destination_generation_ready = true; - destination_host_ready = true; - } - - void wait_for_destination_generation() { - if (sched->batch_split_copies && destination_generation_ready) { - return; - } - const auto start = profiling() - ? ggml_backend_sched_profile_clock::now() - : ggml_backend_sched_profile_clock::time_point{}; - if (sched->events[destination_backend_id][sched->cur_copy]) { - ggml_backend_event_wait( - destination_backend, - sched->events[destination_backend_id][sched->cur_copy]); - } else { - ggml_backend_synchronize(destination_backend); - destination_host_ready = true; - } - if (profiling()) { - sched->profile.destination_wait_us += - ggml_backend_sched_elapsed_us( - start, ggml_backend_sched_profile_clock::now()); - } - destination_generation_ready = true; - } - - void synchronize_source(ggml_backend_t source_backend) { - const auto end = synchronized_sources + n_synchronized_sources; - if (sched->batch_split_copies && - std::find(synchronized_sources, end, source_backend) != end) { - return; - } - const auto start = profiling() - ? ggml_backend_sched_profile_clock::now() - : ggml_backend_sched_profile_clock::time_point{}; - ggml_backend_synchronize(source_backend); - if (profiling()) { - for (int backend_id = 0; backend_id < sched->n_backends; - ++backend_id) { - if (sched->backends[backend_id] == source_backend) { - sched->profile.source_wait_us[backend_id] += - ggml_backend_sched_elapsed_us( - start, ggml_backend_sched_profile_clock::now()); - sched->profile.source_waits[backend_id]++; - break; - } - } - } - if (sched->batch_split_copies) { - GGML_ASSERT(n_synchronized_sources < GGML_SCHED_MAX_BACKENDS); - synchronized_sources[n_synchronized_sources++] = source_backend; - } - } - - bool can_stage(int source_backend_id, - const struct ggml_tensor * source) const { - ggml_backend_sched_tensor_payload_layout layout; - return sched->batch_split_copies && - staging_arena_reusable && - ggml_backend_sched_batch_staging_candidate( - sched, source_backend_id, - destination_backend_id, source) && - ggml_backend_sched_make_tensor_payload_layout(source, layout) && - layout.bytes > 0 && - sched->batch_staging_bases[source_backend_id] && - sched->batch_staging_bases[destination_backend_id]; - } - - void stage(ggml_backend_t source_backend, - int source_backend_id, - const struct ggml_tensor * source, - struct ggml_tensor * destination) { - GGML_ASSERT(ggml_are_same_layout(source, destination)); - ggml_backend_sched_tensor_payload_layout layout; - GGML_ASSERT(ggml_backend_sched_make_tensor_payload_layout( - source, layout)); - size_t source_offset = 0; - size_t destination_offset = 0; - GGML_ASSERT(ggml_backend_sched_reserve_staging( - staging_cursors[source_backend_id], - layout.bytes, &source_offset)); - GGML_ASSERT(ggml_backend_sched_reserve_staging( - staging_cursors[destination_backend_id], - layout.bytes, &destination_offset)); - GGML_ASSERT(staging_cursors[source_backend_id] <= - sched->batch_staging_sizes[source_backend_id]); - GGML_ASSERT(staging_cursors[destination_backend_id] <= - sched->batch_staging_sizes[destination_backend_id]); - - const auto start = profiling() - ? ggml_backend_sched_profile_clock::now() - : ggml_backend_sched_profile_clock::time_point{}; - ggml_backend_sched_get_payload_async( - source_backend, source, - sched->batch_staging_bases[source_backend_id] + source_offset, - layout); - if (profiling()) { - sched->profile.d2h_submit_us += ggml_backend_sched_elapsed_us( - start, ggml_backend_sched_profile_clock::now()); - sched->profile.staged_bytes += layout.bytes; - sched->profile.staged_copies++; - } - - GGML_ASSERT(n_staged < GGML_SCHED_MAX_SPLIT_INPUTS); - staged[n_staged++] = { - source_backend, source_backend_id, destination, - layout, source_offset, destination_offset}; - } - - void prepare_blocking_copy() { - if (sched->events[destination_backend_id][sched->cur_copy]) { - ggml_backend_event_synchronize( - sched->events[destination_backend_id][sched->cur_copy]); - destination_host_ready = true; - } else if (!destination_host_ready) { - ggml_backend_synchronize(destination_backend); - destination_host_ready = true; - } - } - - void flush_staged() { - for (int copy_id = 0; copy_id < n_staged; ++copy_id) { - synchronize_source(staged[copy_id].source_backend); - } - for (int copy_id = 0; copy_id < n_staged; ++copy_id) { - const ggml_backend_sched_staged_copy & copy = staged[copy_id]; - auto start = profiling() - ? ggml_backend_sched_profile_clock::now() - : ggml_backend_sched_profile_clock::time_point{}; - memcpy( - sched->batch_staging_bases[destination_backend_id] + - copy.destination_offset, - sched->batch_staging_bases[copy.source_backend_id] + - copy.source_offset, - copy.layout.bytes); - if (profiling()) { - sched->profile.host_relay_us += - ggml_backend_sched_elapsed_us( - start, ggml_backend_sched_profile_clock::now()); - } - - start = profiling() - ? ggml_backend_sched_profile_clock::now() - : ggml_backend_sched_profile_clock::time_point{}; - ggml_backend_sched_set_payload_async( - destination_backend, copy.destination, - sched->batch_staging_bases[destination_backend_id] + - copy.destination_offset, - copy.layout); - if (profiling()) { - sched->profile.h2d_submit_us += - ggml_backend_sched_elapsed_us( - start, ggml_backend_sched_profile_clock::now()); - } - } - } -}; - -static ggml_backend_sched_profile_clock::time_point -ggml_backend_sched_begin_profile(ggml_backend_sched_t sched) { - const bool enabled = sched->profile_requested && - sched->n_splits >= sched->profile_min_splits; - sched->profile = {}; - sched->profile.active = enabled; - if (!enabled) { - return {}; - } - for (int split_id = 0; split_id < sched->n_splits; ++split_id) { - const int backend_id = sched->splits[split_id].backend_id; - if (backend_id >= 0 && backend_id < sched->n_backends) { - sched->profile.split_counts[backend_id]++; - } - } - return ggml_backend_sched_profile_clock::now(); -} - static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); struct ggml_backend_sched_split * splits = sched->splits; - const auto profile_loop_start = ggml_backend_sched_begin_profile(sched); + static const bool destination_stream_copies_enabled = [] { + const char * value = getenv("GGML_CUDA_DST_STREAM_PEER_COPIES"); + return value && *value && strcmp(value, "0") != 0; + }(); ggml_tensor * prev_ids_tensor = nullptr; std::vector ids; std::vector used_ids; - size_t batch_staging_cursors[GGML_SCHED_MAX_BACKENDS] = {}; - const bool copy_destinations_ready = sched->backends_synchronized; - // From this point onward an early return must remain conservative: a - // backend may have accepted work even if a later split fails. - sched->backends_synchronized = false; for (int split_id = 0; split_id < sched->n_splits; split_id++) { struct ggml_backend_sched_split * split = &splits[split_id]; @@ -2342,15 +1777,57 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } - ggml_backend_sched_split_copy_state copy_state{ - sched, split_backend_id, split_backend, - copy_destinations_ready, copy_destinations_ready, - batch_staging_cursors}; + // All copied inputs for this split use the same destination backend + // and scheduler copy generation. Waiting for that generation once is + // sufficient before overwriting any of its input buffers. The legacy + // loop waited again before every tensor; without scheduler events each + // wait synchronizes the whole destination stream and serializes a + // multi-input peer handoff. + bool split_copy_generation_ready = false; + bool split_copy_source_ready[GGML_SCHED_MAX_BACKENDS] = {}; + auto wait_for_split_copy_generation = [&](ggml_backend_t input_backend, + int input_backend_id) { + if (sched->batch_split_copies && split_copy_generation_ready) { + return; + } + + ggml_backend_event_t generation_event = + sched->events[split_backend_id][sched->cur_copy]; + if (sched->single_copy_event_fences && generation_event != NULL && + input_backend_id != split_backend_id && + input_backend->iface.event_wait != NULL) { + // Async copies run after backend_src, while backend_dst waits + // for their completion. Put the missing destination-buffer + // lifetime edge on the source stream: + // + // dst previous use -> src copy -> dst next use. + // + // A split can have inputs from multiple backends, so one wait + // is required per source stream rather than merely per split. + if (!sched->batch_split_copies || + !split_copy_source_ready[input_backend_id]) { + ggml_backend_event_wait(input_backend, generation_event); + split_copy_source_ready[input_backend_id] = true; + } + return; + } + + if (generation_event != NULL) { + ggml_backend_event_wait( + split_backend, generation_event); + } else { + ggml_backend_synchronize(split_backend); + } + split_copy_generation_ready = true; + }; // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { - ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); struct ggml_tensor * input = split->inputs[input_id]; + const int input_backend_id = tensor_backend_id(input); + GGML_ASSERT(input_backend_id >= 0 && + input_backend_id < sched->n_backends); + ggml_backend_t input_backend = sched->backends[input_backend_id]; struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy); if (input->flags & GGML_TENSOR_FLAG_INPUT) { @@ -2361,11 +1838,27 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } else { ggml_backend_synchronize(split_backend); } - copy_state.mark_destination_ready(); + split_copy_generation_ready = true; ggml_backend_tensor_copy(input, input_cpy); } else { - // wait for the split backend to finish using the input before overwriting it - copy_state.wait_for_destination_generation(); + // A destination-stream peer copy is naturally ordered after + // every prior use of its scheduler buffer on that stream. Do + // not add the ordinary source-stream lifetime fence: doing so + // makes the producer wait for the peer and destroys the fork + // overlap this explicitly marked staging path exists to gain. + // + // The flag is only attached when the HIP destination-stream + // implementation is enabled. Other copies retain the normal + // generation fence below. + const bool destination_stream_copy = + (input->flags & GGML_TENSOR_FLAG_DST_STREAM_COPY) != 0 && + destination_stream_copies_enabled; + if (!destination_stream_copy) { + // Wait for the split backend to finish using the input + // buffer before a source-stream copy overwrites it. + wait_for_split_copy_generation( + input_backend, input_backend_id); + } // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used ggml_tensor * node = split->graph.nodes[0]; @@ -2456,51 +1949,20 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { - const int input_backend_id = - ggml_backend_sched_backend_id(sched, input_backend); - if (copy_state.can_stage(input_backend_id, input)) { - copy_state.stage( - input_backend, input_backend_id, - input, input_cpy); - continue; + ggml_backend_synchronize(input_backend); + if (sched->events[split_backend_id][sched->cur_copy] != NULL) { + ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); + } else { + ggml_backend_synchronize(split_backend); } - - // All non-input copies are gathered before this split - // launches. Once a source backend is synchronized, its - // other inputs for the same split are ready as well. - // Cross-runtime MoE joins commonly carry activation, - // route-ID, and route-weight tensors together; waiting - // on the same source for each tensor serialized dozens - // of redundant stream round trips per verifier step. - copy_state.synchronize_source(input_backend); - - // With one scheduler copy there is no event object and - // wait_for_destination_generation() already established - // host-visible destination quiescence. Parallel-copy - // schedulers retain the explicit event synchronization - // before a blocking host-staged overwrite. - copy_state.prepare_blocking_copy(); ggml_backend_tensor_copy(input, input_cpy); } } } } - // Queue D2H for every source, wait once per source, then relay and - // enqueue all H2D transfers immediately before the consumer graph. - copy_state.flush_staged(); - if (!sched->callback_eval) { - const auto submit_start = sched->profile.active - ? ggml_backend_sched_profile_clock::now() - : ggml_backend_sched_profile_clock::time_point{}; enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); - if (sched->profile.active) { - sched->profile.compute_submit_us += - ggml_backend_sched_elapsed_us( - submit_start, - ggml_backend_sched_profile_clock::now()); - } if (ec != GGML_STATUS_SUCCESS) { return ec; } @@ -2538,7 +2000,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } - // Publish producer completion before a later consumer-backend graph + // Publish cold-owner completion before a later main-backend graph // reaches its in-graph event wait. Recording is asynchronous and does // not block the host from immediately enqueueing independent work. for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { @@ -2551,19 +2013,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } // Record the event of this copy generation. - if (split->n_inputs > 0) { + if (split->n_inputs > 0 || sched->single_copy_event_fences) { if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend); } } } - if (sched->profile.active) { - sched->profile.loop_us = ggml_backend_sched_elapsed_us( - profile_loop_start, - ggml_backend_sched_profile_clock::now()); - } - return GGML_STATUS_SUCCESS; } @@ -2590,13 +2046,6 @@ ggml_backend_sched_t ggml_backend_sched_new( const char * GGML_SCHED_DEBUG_REALLOC = getenv("GGML_SCHED_DEBUG_REALLOC"); sched->debug_realloc = GGML_SCHED_DEBUG_REALLOC ? atoi(GGML_SCHED_DEBUG_REALLOC) : sched->debug_realloc; - const char * profile_raw = getenv("GGML_SCHED_PROFILE"); - const char * profile_min_raw = getenv("GGML_SCHED_PROFILE_MIN_SPLITS"); - sched->profile_requested = - profile_raw && profile_raw[0] && strcmp(profile_raw, "0") != 0; - sched->profile_min_splits = profile_min_raw - ? std::max(1, atoi(profile_min_raw)) : 1; - sched->n_backends = n_backends; sched->n_copies = parallel ? GGML_SCHED_MAX_COPIES : 1; @@ -2647,13 +2096,13 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { if (sched == NULL) { return; } - // graph_compute_async callers may tear down without an explicit wait. - // Quiesce every backend before releasing graph allocations, native events, - // or pinned staging pages still referenced by queued transfers. + // graph_compute_async callers and native graph caches may still hold + // queued references to scheduler events or gallocr buffers. Teardown is + // rare (cache eviction/shutdown), so quiesce every backend before freeing + // any object that a device stream can still reference. for (int b = 0; b < sched->n_backends; ++b) { ggml_backend_synchronize(sched->backends[b]); } - ggml_backend_sched_free_batch_staging(sched); for (int b = 0; b < sched->n_backends; b++) { for (int c = 0; c < sched->n_copies; c++) { ggml_backend_event_free(sched->events[b][c]); @@ -2686,14 +2135,12 @@ void ggml_backend_sched_reset(ggml_backend_sched_t sched) { // reset explicitly discards that graph, so carrying either list into the // next allocation would dereference stale metadata (and leak the events). sched->n_late_cross_input_split_nodes = 0; - if (!sched->backends_synchronized) { - // A caller may reset after an asynchronous submission. Graph - // allocations, deferred events, and host-staging pages all remain - // reachable by queued work until every scheduler backend is idle. + if (sched->n_deferred_peer_copies > 0) { + // A caller may reset after an asynchronous submission. Do not destroy + // producer events while either backend can still reference them. for (int backend_id = 0; backend_id < sched->n_backends; ++backend_id) { ggml_backend_synchronize(sched->backends[backend_id]); } - sched->backends_synchronized = true; } for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { ggml_backend_event_free(sched->deferred_peer_copies[i].event); @@ -2751,10 +2198,6 @@ bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgra ggml_backend_sched_split_graph(sched, graph); - if (!ggml_backend_sched_prepare_batch_staging(sched)) { - return false; - } - if (!ggml_backend_sched_alloc_splits(sched)) { return false; } @@ -2782,6 +2225,12 @@ bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgra return true; } +enum ggml_status ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { + enum ggml_status err = ggml_backend_sched_graph_compute_async(sched, graph); + ggml_backend_sched_synchronize(sched); + return err; +} + enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { GGML_ASSERT(sched); if (!sched->is_reset && !sched->is_alloc) { @@ -2797,52 +2246,10 @@ enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sch return ggml_backend_sched_compute_splits(sched); } -enum ggml_status ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { - enum ggml_status err = ggml_backend_sched_graph_compute_async(sched, graph); - ggml_backend_sched_synchronize(sched); - return err; -} - void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) { GGML_ASSERT(sched); - using profile_clock = std::chrono::steady_clock; for (int i = 0; i < sched->n_backends; i++) { - const profile_clock::time_point wait_start = sched->profile.active - ? profile_clock::now() : profile_clock::time_point{}; ggml_backend_synchronize(sched->backends[i]); - if (sched->profile.active) { - sched->profile.final_wait_us[i] = - (uint64_t) std::chrono::duration_cast( - profile_clock::now() - wait_start).count(); - } - } - sched->backends_synchronized = true; - if (sched->profile.active) { - GGML_LOG_INFO( - "[ggml-sched-profile] splits=%d loop=%lluus dst_wait=%lluus " - "d2h_submit=%lluus host_relay=%lluus h2d_submit=%lluus " - "compute_submit=%lluus " - "staged_copies=%d staged_bytes=%llu\n", - sched->n_splits, - (unsigned long long) sched->profile.loop_us, - (unsigned long long) sched->profile.destination_wait_us, - (unsigned long long) sched->profile.d2h_submit_us, - (unsigned long long) sched->profile.host_relay_us, - (unsigned long long) sched->profile.h2d_submit_us, - (unsigned long long) sched->profile.compute_submit_us, - sched->profile.staged_copies, - (unsigned long long) sched->profile.staged_bytes); - for (int i = 0; i < sched->n_backends; ++i) { - GGML_LOG_INFO( - "[ggml-sched-profile] backend=%s splits=%d " - "source_waits=%d source_wait=%lluus final_wait=%lluus\n", - ggml_backend_name(sched->backends[i]), - sched->profile.split_counts[i], - sched->profile.source_waits[i], - (unsigned long long) sched->profile.source_wait_us[i], - (unsigned long long) sched->profile.final_wait_us[i]); - } - sched->profile.active = false; } if (!sched->is_alloc) { // if the graph is not already allocated, always use copy 0 after a synchronization @@ -2948,6 +2355,41 @@ void ggml_backend_sched_set_batch_split_copies( sched->batch_split_copies = enabled; } +void ggml_backend_sched_set_single_copy_event_fences( + ggml_backend_sched_t sched, bool enabled) { + GGML_ASSERT(sched); + GGML_ASSERT(!sched->is_alloc); + GGML_ASSERT(sched->n_copies == 1); + + if (enabled == sched->single_copy_event_fences) { + return; + } + + if (enabled) { + for (int b = 0; b < sched->n_backends; ++b) { + GGML_ASSERT(sched->events[b][0] == NULL); + sched->events[b][0] = + ggml_backend_event_new(sched->backends[b]->device); + if (sched->events[b][0] != NULL) { + // Give the first source-side wait a defined generation and + // include work queued before scheduler creation. + ggml_backend_event_record( + sched->events[b][0], sched->backends[b]); + } + } + } else { + // A reset can make is_alloc false while a prior submission is still + // running. Finish it before destroying events referenced by streams. + ggml_backend_sched_synchronize(sched); + for (int b = 0; b < sched->n_backends; ++b) { + ggml_backend_event_free(sched->events[b][0]); + sched->events[b][0] = NULL; + } + } + + sched->single_copy_event_fences = enabled; +} + int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); return sched->n_splits; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 5e169ee36..c07a93fa0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3428,6 +3428,14 @@ static bool ggml_cuda_batch_peer_copies_enabled() { return enabled; } +static bool ggml_cuda_destination_peer_copies_enabled() { + static const bool enabled = [] { + const char * value = getenv("GGML_CUDA_DST_STREAM_PEER_COPIES"); + return value && *value && strcmp(value, "0") != 0; + }(); + return enabled; +} + static void ggml_cuda_flush_peer_copy_batch(const char * reason) { auto & batch = ggml_cuda_pending_peer_copies; if (batch.copies == 0) { @@ -3504,6 +3512,34 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ return false; #else #if defined(GGML_USE_HIP) + if (ggml_cuda_destination_peer_copies_enabled() && + (src->flags & GGML_TENSOR_FLAG_DST_STREAM_COPY)) { + // The producer split has already been queued when the generic + // scheduler reaches this consumer input. Publish its stream, + // then perform the transfer on the destination stream. The + // producer can immediately continue with its independent + // branch instead of placing the peer copy in front of it. + // Only explicitly marked persistent staging tensors may take + // this path: ordinary graph temporaries can be recycled by the + // producer before a destination-stream copy has consumed them. + ggml_cuda_flush_peer_copy_batch("destination-stream-copy"); + ggml_cuda_set_device(cuda_ctx_src->device); + if (!cuda_ctx_src->copy_event) { + CUDA_CHECK(cudaEventCreateWithFlags( + &cuda_ctx_src->copy_event, cudaEventDisableTiming)); + } + CUDA_CHECK(cudaEventRecord( + cuda_ctx_src->copy_event, cuda_ctx_src->stream())); + + ggml_cuda_set_device(cuda_ctx_dst->device); + CUDA_CHECK(cudaStreamWaitEvent( + cuda_ctx_dst->stream(), cuda_ctx_src->copy_event, 0)); + CUDA_CHECK(cudaMemcpyPeerAsync( + dst->data, cuda_ctx_dst->device, + src->data, cuda_ctx_src->device, + ggml_nbytes(dst), cuda_ctx_dst->stream())); + return true; + } if (ggml_cuda_batch_peer_copies_enabled()) { auto & batch = ggml_cuda_pending_peer_copies; if (batch.copies != 0 && diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu index 5f1954e0c..a33f97411 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -396,6 +396,7 @@ static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna4(ggml_type // opt-in until qualified on each AMD target. #define MMID_GROUPED_MAX_PAIRS 256 #define MMID_GROUPED_MAX_TPG 8 +#define MMID_GROUP_REUSE_MAX_PAIRS 5 #define MMID_META_NG 0 #define MMID_META_GE 1 #define MMID_META_GS (MMID_META_GE + MMID_GROUPED_MAX_PAIRS) @@ -478,6 +479,21 @@ static bool mmid_grouped_device_ok() { return selected < 0 || selected == ggml_cuda_get_device(); } +// Decode each expert weight fragment once for every token routed to that +// expert. This is deliberately opt-in until it has been qualified on each +// target architecture; DFLASH_MMID_GROUPED remains the independent fallback. +static bool mmid_group_reuse_env() { + static const bool on = + mmvq_env_flag("DFLASH_CUDA_MMVQ_MOE_GROUP_REUSE", false); + return on; +} + +static bool mmid_group_reuse_ok(ggml_type type, int64_t ncols_dst) { + return mmid_group_reuse_env() && ncols_dst >= 2 && ncols_dst <= 5 && + (type == GGML_TYPE_Q2_0_ROCMFP2 || + type == GGML_TYPE_Q3_0_ROCMFPX); +} + static bool mmid_grouped_arch_ok(int cc) { return (GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING) || GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc); @@ -1406,7 +1422,8 @@ static __global__ void mul_mat_vec_q_moe( static __global__ void mmid_group_prep( int32_t * __restrict__ ids, int32_t * __restrict__ meta, const int n_slots, const int n_tok, const int ids_stride, - float * __restrict__ gate_w, const int gate_w_stride, const float gate_tau) { + float * __restrict__ gate_w, const int gate_w_stride, + const float gate_tau, const bool build_groups) { __shared__ int sh_expert[MMID_GROUPED_MAX_PAIRS]; const int np = n_slots*n_tok; const int i = threadIdx.x; @@ -1464,6 +1481,340 @@ static __global__ void mmid_group_prep( meta[MMID_META_PT + r] = i / n_slots; meta[MMID_META_PS + r] = i % n_slots; } + if (!build_groups) { + return; + } + __syncthreads(); + + // Build a compact expert-group index over the sorted pairs. The normal + // grouped kernel ignores this metadata; the weight-reuse variant maps one + // warp to one [start, end) group without a host readback or stream sync. + if (i == 0) { + int ng = 0; + int previous = 0x7FFFFFFF; + for (int p = 0; p < np; ++p) { + const int current = meta[MMID_META_GE + p]; + if (p == 0 || current != previous) { + meta[MMID_META_GS + ng++] = p; + previous = current; + } + } + meta[MMID_META_GS + ng] = np; + meta[MMID_META_NG] = ng; + } +} + +// Apply one ROCmFP2/ROCmFP3 weight fragment to every activation routed to the +// same expert. Weight unpacking and scale conversion happen once; each pair +// retains an independent integer accumulator and the original accumulation +// order, so its floating-point result matches the single-pair dot product. +template +static __device__ __forceinline__ void vec_dot_rocmfpx_q8_1_group( + const void * __restrict__ vx, + const block_q8_1 * const y[MMID_GROUP_REUSE_MAX_PAIRS], + const int group_n, const int kbx, const int kby, const int iqs, + float out[MMID_GROUP_REUSE_MAX_PAIRS]) { + static_assert(type == GGML_TYPE_Q2_0_ROCMFP2 || + type == GGML_TYPE_Q3_0_ROCMFPX, + "group reuse currently supports ROCmFP2/ROCmFP3 only"); + + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2) { + const block_rocmfp2 * bq2 = (const block_rocmfp2 *) vx + kbx; + int sumi[MMID_GROUP_REUSE_MAX_PAIRS] = {0}; + + if constexpr (c_fp2_packed32) { + uint32_t packed32; + memcpy(&packed32, bq2->qs + 4*iqs, sizeof(packed32)); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t bits8 = (packed32 >> (8*j)) & 0xFFu; + const int val = rocmfpx_pack4_fp2_bits8_vec_cuda(bits8); +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + const int u = get_int_b4(y[p][kby].qs, 4*iqs + j); + sumi[p] = ggml_cuda_dp4a(val, u, sumi[p]); + } + } + } + } else { +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int val = rocmfpx_pack4_fp2_bits8_vec_cuda( + (uint32_t) bq2->qs[4*iqs + j]); +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + const int u = get_int_b4(y[p][kby].qs, 4*iqs + j); + sumi[p] = ggml_cuda_dp4a(val, u, sumi[p]); + } + } + } + } + + const float scale = rocmfpx_ue4m3_to_fp32_finite(bq2->e[iqs]); +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + out[p] = __low2float(y[p][kby].ds) * scale * sumi[p]; + } + } + } else { + const block_rocmfp3 * bq3 = (const block_rocmfp3 *) vx + kbx; + int sumi0[MMID_GROUP_REUSE_MAX_PAIRS] = {0}; + int sumi1[MMID_GROUP_REUSE_MAX_PAIRS] = {0}; + + if constexpr (c_fp3_packed24) { + const int byte_offset = 3 * (iqs >> 1); + const uint32_t packed24 = + (uint32_t) bq3->qs[byte_offset + 0] | + ((uint32_t) bq3->qs[byte_offset + 1] << 8) | + ((uint32_t) bq3->qs[byte_offset + 2] << 16); + const int val0 = + rocmfpx_pack4_fp3_bits12_vec_cuda(packed24 & 0xFFFu); + const int val1 = + rocmfpx_pack4_fp3_bits12_vec_cuda(packed24 >> 12); +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + const int u0 = get_int_b4(y[p][kby].qs, iqs + 0); + const int u1 = get_int_b4(y[p][kby].qs, iqs + 1); + if (iqs < QK_ROCMFP3/8) { + sumi0[p] = ggml_cuda_dp4a(val0, u0, sumi0[p]); + sumi0[p] = ggml_cuda_dp4a(val1, u1, sumi0[p]); + } else { + sumi1[p] = ggml_cuda_dp4a(val0, u0, sumi1[p]); + sumi1[p] = ggml_cuda_dp4a(val1, u1, sumi1[p]); + } + } + } + } else { + uint32_t qs0, qs1, qs2; + memcpy(&qs0, bq3->qs + 0, 4); + memcpy(&qs1, bq3->qs + 4, 4); + memcpy(&qs2, bq3->qs + 8, 4); + const uint32_t qs[4] = {qs0, qs1, qs2, 0}; + const bool first = + iqs + VDR_ROCMFP3_Q8_1_MMVQ <= QK_ROCMFP3/8; + const bool second = iqs >= QK_ROCMFP3/8; +#pragma unroll + for (int j = 0; j < VDR_ROCMFP3_Q8_1_MMVQ; ++j) { + const int base = 4 * (iqs + j); + const int start_bit = 12 * (iqs + j); + const int reg_idx = start_bit >> 5; + const int reg_shift = start_bit & 31; + const uint32_t low = qs[reg_idx]; + const uint32_t high = qs[reg_idx + 1]; + const uint32_t bits12 = reg_shift == 0 ? low & 0xFFFu : + ((low >> reg_shift) | (high << (32 - reg_shift))) & + 0xFFFu; + const int val = rocmfpx_pack4_fp3_bits12_vec_cuda(bits12); +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + const int u = get_int_b4(y[p][kby].qs, iqs + j); + if (first || (!second && base < QK_ROCMFP3/2)) { + sumi0[p] = ggml_cuda_dp4a(val, u, sumi0[p]); + } else { + sumi1[p] = ggml_cuda_dp4a(val, u, sumi1[p]); + } + } + } + } + } + + const float scale0 = rocmfpx_ue4m3_to_fp32_finite(bq3->e[0]); + const float scale1 = rocmfpx_ue4m3_to_fp32_finite(bq3->e[1]); +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + const float db = __low2float(y[p][kby].ds); + out[p] = db * (scale0*sumi0[p] + scale1*sumi1[p]); + } + } + } +} + +// One warp owns one unique expert. Repeated routes are processed in chunks of +// five (the speculative-verify ceiling for this path), avoiding duplicate +// weight decode work while preserving an independent result for every slot. +template +__launch_bounds__(MMID_GROUPED_MAX_TPG*ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_q_moe_group_reuse( + const void * __restrict__ vx, const void * __restrict__ vy, + const int32_t * __restrict__ meta, + const ggml_cuda_mm_fusion_args_device fusion, + float * __restrict__ dst, + const uint32_t ncols_x, const uint3 nchannels_y, + const uint32_t nrows_x, const uint32_t stride_row_x, + const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, + const uint32_t stride_channel_dst) { + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int blocks_per_iter = vdr * warp_size / qi; + + const int group = + blockIdx.y*MMID_GROUPED_MAX_TPG + (int) threadIdx.y; + const int ng = meta[MMID_META_NG]; + if (group >= ng) { + return; + } + const int group_start = meta[MMID_META_GS + group]; + const int group_end = meta[MMID_META_GS + group + 1]; + const int channel_x = meta[MMID_META_GE + group_start]; + const int row0 = c_rows_per_block*blockIdx.x; + + if (channel_x < 0) { + if (threadIdx.x < c_rows_per_block && + (c_rows_per_block == 1 || + uint32_t(row0 + threadIdx.x) < nrows_x)) { + for (int p = group_start; p < group_end; ++p) { + const int tok = meta[MMID_META_PT + p]; + const int slot = meta[MMID_META_PS + p]; + dst[slot*stride_channel_dst + tok*stride_col_dst + row0 + + threadIdx.x] = 0.0f; + } + } + return; + } + + const int blocks_per_row_x = ncols_x/qk; + const int kbx_base = channel_x*stride_channel_x + row0*stride_row_x; + const bool use_gate = has_fusion && fusion.gate != nullptr; + + for (int chunk = group_start; chunk < group_end; + chunk += MMID_GROUP_REUSE_MAX_PAIRS) { + const int remaining = group_end - chunk; + const int group_n = remaining < MMID_GROUP_REUSE_MAX_PAIRS ? + remaining : MMID_GROUP_REUSE_MAX_PAIRS; + int tok[MMID_GROUP_REUSE_MAX_PAIRS] = {0}; + int slot[MMID_GROUP_REUSE_MAX_PAIRS] = {0}; + const block_q8_1 * y[MMID_GROUP_REUSE_MAX_PAIRS] = {nullptr}; +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + tok[p] = meta[MMID_META_PT + chunk + p]; + slot[p] = meta[MMID_META_PS + chunk + p]; + const uint32_t channel_y = + fastmodulo((uint32_t) slot[p], nchannels_y); + y[p] = ((const block_q8_1 *) vy) + + channel_y*stride_channel_y + tok[p]*stride_col_y; + } + } + + float tmp[c_rows_per_block][MMID_GROUP_REUSE_MAX_PAIRS] = {{0}}; + float tmp_gate[c_rows_per_block][MMID_GROUP_REUSE_MAX_PAIRS] = {{0}}; + for (int kbx = threadIdx.x/(qi/vdr); kbx < blocks_per_row_x; + kbx += blocks_per_iter) { + const int kby = kbx*(qk/QK8_1); + const int kqs = vdr*(threadIdx.x % (qi/vdr)); +#pragma unroll + for (int row = 0; row < c_rows_per_block; ++row) { + float dots[MMID_GROUP_REUSE_MAX_PAIRS] = {0}; + vec_dot_rocmfpx_q8_1_group< + type, c_fp3_packed24, c_fp2_packed32>( + vx, y, group_n, + kbx_base + row*stride_row_x + kbx, kby, kqs, dots); +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + tmp[row][p] += dots[p]; + } + } + if (use_gate) { + float gate_dots[MMID_GROUP_REUSE_MAX_PAIRS] = {0}; + vec_dot_rocmfpx_q8_1_group< + type, c_fp3_packed24, c_fp2_packed32>( + fusion.gate, y, group_n, + kbx_base + row*stride_row_x + kbx, kby, kqs, + gate_dots); +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + tmp_gate[row][p] += gate_dots[p]; + } + } + } + } + } + +#pragma unroll + for (int row = 0; row < c_rows_per_block; ++row) { +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + tmp[row][p] = warp_reduce_sum(tmp[row][p]); + if (use_gate) { + tmp_gate[row][p] = + warp_reduce_sum(tmp_gate[row][p]); + } + } + } + } + + if (threadIdx.x < c_rows_per_block && + (c_rows_per_block == 1 || + uint32_t(row0 + threadIdx.x) < nrows_x)) { + const int row = threadIdx.x; +#pragma unroll + for (int p = 0; p < MMID_GROUP_REUSE_MAX_PAIRS; ++p) { + if (p < group_n) { + float result = tmp[row][p]; + if constexpr (has_fusion) { + if (fusion.x_bias != nullptr) { + result += ((const float *) fusion.x_bias)[ + channel_x*stride_channel_dst + row0 + row]; + } + if (use_gate) { + float gate_value = tmp_gate[row][p]; + if (fusion.gate_bias != nullptr) { + gate_value += + ((const float *) fusion.gate_bias)[ + channel_x*stride_channel_dst + + row0 + row]; + } + switch (fusion.glu_op) { + case GGML_GLU_OP_SWIGLU: + result *= + ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= + ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single( + gate_value, result, + fusion.glu_param0, + fusion.glu_param1); + break; + case GGML_GLU_OP_SWIGLU_DS4: + if (fusion.gate_value_scale != 1.0f) { + gate_value *= fusion.gate_value_scale; + } + if (fusion.x_value_scale != 1.0f) { + result *= fusion.x_value_scale; + } + result = ggml_cuda_op_swiglu_ds4_single( + gate_value, result, + fusion.glu_param0); + break; + default: + result *= gate_value; + break; + } + } + } + dst[slot[p]*stride_channel_dst + + tok[p]*stride_col_dst + row0 + row] = result; + } + } + } + } } // [TAG_MMID_GROUPED] grouped MoE kernel. Identical launch shape and per-warp @@ -1682,13 +2033,49 @@ static void mul_mat_vec_q_moe_grouped_launch( } } +template +static void mul_mat_vec_q_moe_group_reuse_launch( + const void * vx, const void * vy, const int32_t * meta, + const ggml_cuda_mm_fusion_args_device & fusion, float * dst, + const uint32_t ncols_x, const uint3 nchannels_y, + const uint32_t nrows_x, const uint32_t stride_row_x, + const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, + const uint32_t stride_channel_dst, const int max_groups, + const int warp_size, cudaStream_t stream) { + constexpr int rows_per_block = 1; + const int64_t nblocks_rows = nrows_x; + const dim3 block_nums( + nblocks_rows, + (max_groups + MMID_GROUPED_MAX_TPG - 1)/MMID_GROUPED_MAX_TPG); + const dim3 block_dims(warp_size, MMID_GROUPED_MAX_TPG); + const bool has_fusion = fusion.gate != nullptr || + fusion.x_bias != nullptr || fusion.gate_bias != nullptr; + if (has_fusion) { + mul_mat_vec_q_moe_group_reuse< + type, rows_per_block, true, c_fp3_packed24, + c_fp2_packed32><<>>( + vx, vy, meta, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst); + } else { + mul_mat_vec_q_moe_group_reuse< + type, rows_per_block, false, c_fp3_packed24, + c_fp2_packed32><<>>( + vx, vy, meta, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst); + } +} + static bool mul_mat_vec_q_grouped_dispatch( const ggml_type type, const void * vx, const void * vy, const int32_t * meta, const ggml_cuda_mm_fusion_args_device & fusion, float * dst, const int ncols_x, const int nrows_x, const int nchannels_y, const int stride_row_x, const int stride_col_y, const int stride_col_dst, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, - const int max_groups, cudaStream_t stream) { + const int max_groups, const int ncols_dst, cudaStream_t stream) { const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; const uint3 nchannels_y_fd = init_fastdiv_values((uint32_t) nchannels_y); @@ -1702,6 +2089,7 @@ static bool mul_mat_vec_q_grouped_dispatch( std::getenv("DFLASH_CUDA_MMVQ_MOE_FP2_PACKED32"); return e && e[0] == '1' && e[1] == '\0'; }(); + const bool group_reuse = mmid_group_reuse_ok(type, ncols_dst); switch (type) { case GGML_TYPE_Q4_0: @@ -1725,6 +2113,26 @@ static bool mul_mat_vec_q_grouped_dispatch( stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst, max_groups, warp_size, stream); return true; case GGML_TYPE_Q2_0_ROCMFP2: + if (group_reuse) { + if (fp2_packed32) { + mul_mat_vec_q_moe_group_reuse_launch< + GGML_TYPE_Q2_0_ROCMFP2, false, true>( + vx, vy, meta, fusion, dst, ncols_x, + nchannels_y_fd, nrows_x, stride_row_x, + stride_col_y, stride_col_dst, stride_channel_x, + stride_channel_y, stride_channel_dst, max_groups, + warp_size, stream); + } else { + mul_mat_vec_q_moe_group_reuse_launch< + GGML_TYPE_Q2_0_ROCMFP2>( + vx, vy, meta, fusion, dst, ncols_x, + nchannels_y_fd, nrows_x, stride_row_x, + stride_col_y, stride_col_dst, stride_channel_x, + stride_channel_y, stride_channel_dst, max_groups, + warp_size, stream); + } + return true; + } if (fp2_packed32) { mul_mat_vec_q_moe_grouped_launch< GGML_TYPE_Q2_0_ROCMFP2, false, true>( @@ -1742,6 +2150,26 @@ static bool mul_mat_vec_q_grouped_dispatch( } return true; case GGML_TYPE_Q3_0_ROCMFPX: + if (group_reuse) { + if (fp3_packed24) { + mul_mat_vec_q_moe_group_reuse_launch< + GGML_TYPE_Q3_0_ROCMFPX, true, false>( + vx, vy, meta, fusion, dst, ncols_x, + nchannels_y_fd, nrows_x, stride_row_x, + stride_col_y, stride_col_dst, stride_channel_x, + stride_channel_y, stride_channel_dst, max_groups, + warp_size, stream); + } else { + mul_mat_vec_q_moe_group_reuse_launch< + GGML_TYPE_Q3_0_ROCMFPX>( + vx, vy, meta, fusion, dst, ncols_x, + nchannels_y_fd, nrows_x, stride_row_x, + stride_col_y, stride_col_dst, stride_channel_x, + stride_channel_y, stride_channel_dst, max_groups, + warp_size, stream); + } + return true; + } if (fp3_packed24) { mul_mat_vec_q_moe_grouped_launch< GGML_TYPE_Q3_0_ROCMFPX, true, false>( @@ -2794,18 +3222,21 @@ void ggml_cuda_mul_mat_vec_q( const int prep_threads = ((np + prep_warp - 1)/prep_warp)*prep_warp; mmid_group_prep<<<1, prep_threads, 0, stream>>>( prep_ids, mmid_meta.ptr, (int) nchannels_dst, (int) ncols_dst, prep_ids_stride, - gate_w, gate_w_stride, gate_tau); + gate_w, gate_w_stride, gate_tau, + mmid_group_reuse_ok(src0->type, ncols_dst)); CUDA_CHECK(cudaGetLastError()); if (mul_mat_vec_q_grouped_dispatch( src0->type, src0->data, src1_q8_d, mmid_meta.ptr, fusion_local, dst_d, (int) ne00, (int) ne01, (int) nchannels_y, (int) s01, (int) stride_col_y, (int) stride_col_dst, (int) s02, (int) stride_channel_y, (int) stride_channel_dst, - np, stream)) { + np, (int) ncols_dst, stream)) { if (mmid_telemetry) { std::fprintf(stderr, - "[dflash-mmid] event=mmvq type=%s width=%lld pairs=%d variant=grouped\n", - ggml_type_name(src0->type), (long long) ncols_dst, np); + "[dflash-mmid] event=mmvq type=%s width=%lld pairs=%d variant=%s\n", + ggml_type_name(src0->type), (long long) ncols_dst, np, + mmid_group_reuse_ok(src0->type, ncols_dst) ? + "group-reuse" : "grouped"); } return; } diff --git a/server/docs/DS4.md b/server/docs/DS4.md index d164a9b17..3e2669f8b 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -5,6 +5,10 @@ DFlash. DeepSeek4 supports a monolithic HIP backend, a layer-split backend, and an in-process heterogeneous expert-parallel backend for a discrete GPU paired with Strix Halo. +The qualified q=5 attention-head split for that heterogeneous pair, including +the exact benchmark recipe and rejected variants, is documented in +[DS4_HETEROGENEOUS_ATTENTION_TP.md](DS4_HETEROGENEOUS_ATTENTION_TP.md). + ## Model Architecture DeepSeek V4 Flash is a 43-layer MoE model with: @@ -124,56 +128,6 @@ decode and 415.52 tok/s median sparse prefill. Those numbers require the full qualified manifest, including the burn-in kernel switches; they are not a claim for the minimal activation example above. -#### CUDA 3090 + Strix Halo in one process - -The mixed-vendor build links the selected target runtime normally and loads -the other vendor as an isolated backend module. CUDA and HIP device indices -have separate namespaces, so `cuda:0` and `hip:0` are a valid pair. -Cross-vendor activations are staged in host memory inside the process; native -peer access is intentionally not attempted. - -```bash -cmake -S server -B server/build-cuda-hip \ - -DDFLASH27B_GPU_BACKEND=hip \ - -DDFLASH27B_ENABLE_MIXED_CUDA_HIP=ON \ - -DDFLASH27B_CUDA_ARCHITECTURES=86 \ - -DDFLASH27B_HIP_ARCHITECTURES=gfx1151 \ - -DCMAKE_BUILD_TYPE=Release -cmake --build server/build-cuda-hip -j -ctest --test-dir server/build-cuda-hip -R mixed_cuda_hip --output-on-failure -``` - -```bash -export DFLASH_DS4_MOE_TP=1 -export DFLASH_DS4_MOE_TP_INPROC=1 -export DFLASH_DS4_MOE_TP_BACKEND=cuda -export DFLASH_DS4_MOE_TP_GPU=0 # cuda:0 (RTX 3090) -export DFLASH_DS4_MOE_TP_CONCENTRATE_COLD=1 -export DFLASH_DS4_TP_SCHEDULE_BRANCHES=1 -export DFLASH_DS4_TP_TARGETED_JOIN_SPLIT=1 -export GGML_BATCH_PEER_COPIES=1 -# Start conservatively and tune from the startup placement and memory logs; -# the usable budget depends on the model, placement policy, and free VRAM. -export DFLASH_EXPERT_BUDGET_MB=85000 -export DFLASH_DS4_DRAFT=/path/to/dspark-draft.gguf -export DFLASH_DS4_DRAFT_BACKEND=cuda -export DFLASH_DS4_DRAFT_GPU=0 - -./server/build-cuda-hip/dflash_server /path/to/deepseek4-target.gguf \ - --target-device hip:0 \ - --ds4-expert-top-k 4 \ - --ds4-prefill sparse -``` - -The peer module is normally found beside the executable. Set -`DFLASH_CUDA_BACKEND_PATH` or `DFLASH_HIP_BACKEND_PATH` only when packaging it -elsewhere. Sparse/approximate DeepSeek4 prefill remains restricted to a HIP -target; CUDA-primary ROCmFP2 execution is not yet qualified. The mixed path is -burn-in functionality. On the qualified 3090 + Strix machine, the tuned top-4 -performance profile held 48.1 tok/s median on the deterministic 128-token -workload. The all-6-expert reference-exact mode is a correctness profile, not -a throughput profile. - ### Local single-shard If the adapter decides all 43 layers fit on one CUDA GPU, it loads a single shard locally and no IPC daemon is involved. @@ -232,30 +186,18 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | `DFLASH_DS4_CUDA_LAYERS` | Override the auto-split heuristic and pin the first `N` DeepSeek4 layers to CUDA. The remaining `43 - N` layers run on the Halo shard. | | `DFLASH_DS4_TIMING` | Enable DS4 timing logs for the layer-split parent and target-shard daemon. Useful for profiling prefill/decode breakdowns; leave unset for normal runs. | | `DFLASH_DS4_SPEC` / `DFLASH_DS4_DRAFT` | Enable DSpark and select its GGUF. | -| `DFLASH_DS4_DRAFT_BACKEND` / `DFLASH_DS4_DRAFT_GPU` | Backend and device for the in-process drafter. | +| `DFLASH_DS4_DRAFT_GPU` | HIP device for the in-process drafter. | | `DFLASH_DS4_MOE_TP` | Enable routed-expert partitioning. | -| `DFLASH_DS4_MOE_TP_INPROC` | Use two local GPU backends instead of an expert IPC worker. | -| `DFLASH_DS4_MOE_TP_BACKEND` | Cold expert backend (`cuda` or `hip`); mixed builds default to the peer runtime. | -| `DFLASH_DS4_MOE_TP_GPU` | Device index within the cold expert backend. | -| `DFLASH_DS4_MOE_TP_CONCENTRATE_COLD` | Cross-vendor burn-in mode: place complete cold expert layers on the peer to reduce joins. | -| `DFLASH_DS4_MOE_TP_PEER_HOT` | With a routing profile, place its hottest experts on the secondary owner. | -| `DFLASH_DS4_CROSS_VENDOR_OWNER_SUMS` | Reduce each owner's routed outputs locally before the final cross-vendor add. This changes floating-point association and is not the byte-identity mode. | -| `DFLASH_DS4_TP_SCHEDULE_BRANCHES` | Submit the two owner branches independently through the mixed scheduler. | -| `DFLASH_DS4_TP_TARGETED_JOIN_SPLIT` | Gather the peer result at the join without an extra peer fence per layer. | -| `DFLASH_DS4_COMP_PAD_STRIDE` | Exact compressed-KV padding bucket; wider buckets trade small masked work for fewer verifier graph captures. | -| `DFLASH_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION` | Diagnostic fallback for runtimes that cannot preserve grouped projection metadata across a scheduler copy. | -| `DFLASH_CUDA_BACKEND_PATH` / `DFLASH_HIP_BACKEND_PATH` | Optional explicit peer backend module path. | +| `DFLASH_DS4_MOE_TP_INPROC` | Use two local HIP backends instead of an expert IPC worker. | +| `DFLASH_DS4_MOE_TP_GPU` | HIP device that owns the cold expert stack. | | `DFLASH_EXPERT_BUDGET_MB` | Main-GPU memory budget for hot experts. | | `DFLASH_DS4_HOTNESS_CSV` | Optional per-layer routing profile for hot placement. | -| `GGML_BATCH_PEER_COPIES` | Batch peer-runtime copies and unlike-runtime pinned-host staging with one source wait per split. The old `GGML_CUDA_BATCH_PEER_COPIES` spelling remains an alias. | -| `DFLASH_DS4_TP_CRITICAL_PATH_PLACEMENT` | Use the routing profile and measured owner-rate ratio to minimize the predicted two-owner MoE critical path instead of maximizing aggregate hot-hit rate. Requires `DFLASH_DS4_HOTNESS_CSV`. | -| `DFLASH_DS4_TP_MAIN_TO_PEER_RATE` | Relative main/peer routed-expert rate used by critical-path placement. It must be finite and greater than zero; the default is `3.4`. | -| `DFLASH_DS4_TP_BALANCE_MIN_HOT` | Minimum hot experts retained on every routed layer by critical-path placement. Defaults to `0`. | | `DFLASH_DS4_Q5_VERIFY` | Opt in to the AMD q=5 fused verifier. This also selects the qualified MMVQ width and verifier-cache defaults when they are not explicitly overridden. | | `DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1` | Select the q=5 ROCmFP4 dense verifier kernel that reuses the existing x4 dot product for columns 0-3 and the exact scalar path for column 4. Defaults to `1` for q=5 on `gfx1201`; set `0` to force the generic five-column kernel. | | `DFLASH_DS4_TP_FUSED_CACHE_SLOTS` | Number of heterogeneous verifier graph slots. Defaults to `2` for q<=4 and `9` for the opt-in q=5 verifier; each slot retains scheduler scratch on both GPUs. | | `DFLASH_DS4_VERIFY_FORCE_GRAPH_REPLAY` | Skip the expensive property scan only for a warmed verifier graph. Rebuilt scheduler generations are always validated. Leave unset for the conservative production profile. | | `GGML_DS4_FA_SERIAL_INDEX_SCAN` | Restore the serial compressed-row mask scan for an indexed-attention A/B. By default, HIP scans contexts above 512 compressed rows in parallel. | +| `GGML_CUDA_BATCH_PEER_COPIES` | Batch ordered peer copies behind one dependency. | | `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` | Long-prefill arena kill switch; set `0` to restore per-layer owner allocation. | `DFLASH_DS4_TIMING` enables the existing timing banners: @@ -316,13 +258,11 @@ whole-model GPU graph uses stable padded reduction shapes, so near-tied greedy logits can select a different token than the normal causal verifier even at temperature 0. Leave it unset when comparing against the normal verifier, or set `DFLASH_DS4_SEQ_VERIFY=1` for the slower token-at-a-time verification -diagnostic. `DFLASH_DS4_SPEC_REFERENCE_EXACT=1` combines sequential target -verification with full rollback snapshots for byte-identity checks. Neither -fused verification nor the separate +diagnostic. Neither fused verification nor the separate `--ds4-expert-top-k 4` approximation should be presented as byte-identical AR. DSpark can verify against in-process heterogeneous expert placement. The -drafter remains local to its selected GPU backend; a failed draft load is +drafter remains local to its selected HIP backend; a failed draft load is reported and falls back to normal autoregressive decode. The target cache and sampler stay on the main backend while routed target experts execute on their configured owners. `--ds4-expert-top-k 4` remains a separate approximate @@ -353,50 +293,24 @@ the shape that crosses two ratio-4 compressor boundaries, preserves five raw SWA rows for rollback, and restores plus replays only the accepted prefix after a partial rejection. q<=4 behavior is unchanged when the flag is absent. -The heterogeneous verifier keeps all five lanes in one -`[n_embd * n_hc, q]` tensor. Attention HC-pre, attention HC-post, FFN HC-pre, -FFN HC-post, drafter-feature capture, and output HC merge are batched across -the verifier width. This removes the former per-lane HC controller paths and -progressive concatenations without changing the verifier result. The split -HC-post kernel also accepts a token dimension and joins the two owner outputs -inside that batched kernel. - -Critical-path placement models each routed layer as two concurrent branches: -the main branch includes its fixed shared-expert work and hot routed work, -while the peer branch executes the remaining routed work. The allocator adds -the next profiled expert only when its marginal reduction in -`max(main / main_to_peer_rate, peer)` is positive. The expert memory budget is -therefore an upper bound; leaving part of it unused is valid when another hot -expert would lengthen the predicted fork. - -On the qualified R9700 + Strix Halo profile, leaving the related q=5 controls +On the qualified R9700 + Strix Halo profile, leaving the related controls unset selects `LUCE_MMVQ_MAX_NCOLS=5`, nine heterogeneous verifier slots, and the ROCmFP4 x4+1 dense kernel on `gfx1201`. The wider MMVQ ceiling avoids the slow small-matrix crossover, while nine slots hold the recurring compressor phases without steady graph rebuilds. The x4+1 kernel decodes shared weights through the existing four-column vector path and retains the original scalar accumulation for the fifth verifier column. -Explicit environment values still take priority. - -Sparse heterogeneous prefill uses a reusable graph allocator with preferred -128 MiB backing chunks. This avoids depending on one large contiguous HIP -allocation on devices without virtual-memory-backed buffers. Individual -tensors remain unsplit and may exceed the preferred chunk size. Prompts ending -above 4K use a 1K-token prefill shape, and that cap remains sticky for later -requests in the process so a post-16K request cannot force a fragmented -1K-to-2K arena replacement. Reproducible decode graph caches are retired before -a necessary prefill-arena growth; persistent HC mirrors remain resident. +Explicit environment values still take priority. The hot-36 full sweep peaked +at 30.561 GiB on the reported 31.86 GiB R9700 and must be requalified on +smaller devices. The exact qualification launch used: ```bash export DFLASH_DS4_Q5_VERIFY=1 export DFLASH_DS4_SPEC_Q=5 -export DFLASH_EXPERT_BUDGET_MB=14350 +export DFLASH_EXPERT_BUDGET_MB=13200 # 36 hot experts/layer on this profile export DFLASH_DS4_HOTNESS_CSV=/path/to/ds4_moe_tp_hotness.csv -export DFLASH_DS4_TP_CRITICAL_PATH_PLACEMENT=1 -export DFLASH_DS4_TP_MAIN_TO_PEER_RATE=4.4 -export DFLASH_DS4_TP_BALANCE_MIN_HOT=0 ``` The checked-in wrapper reproduces the full exact-context protocol and records @@ -407,10 +321,6 @@ trace: TARGET_MODEL=/path/to/target.gguf \ DRAFT_MODEL=/path/to/dspark-draft.gguf \ HOTNESS_CSV=/path/to/ds4_moe_tp_hotness.csv \ -CRITICAL_PATH_PLACEMENT=1 \ -MAIN_TO_PEER_RATE=4.4 \ -BALANCE_MIN_HOT=0 \ -EXPERT_BUDGET_MB=14350 \ server/scripts/qualify_ds4_q5_amd.sh ``` @@ -421,44 +331,10 @@ compatible artifact. At temperature zero, all 25 requests in the 2K -> 4K -> 8K -> 16K -> 2K burn-in produced the same expected response hash. With the automatic q=5 -MMVQ/cache/kernel defaults and the critical-path profile above, measured client -decode medians were 75.818, 74.530, 69.898, 62.685, and 76.703 tok/s. The final -2K measurements after repeated 16K prefill were 76.685-76.727 tok/s, confirming -that the bounded sticky arena recovers steady decode rather than merely -surviving the request. The placement retained 1,688 profiled hot experts, 23-63 -per layer, and the full sweep peaked at 31.089 GiB on the reported 31.86 GiB -main GPU. Treat these as workload-specific burn-in measurements, not as a -portable default for unrelated memory layouts. - -For an overlap trace, run the same wrapper with the delayed profiler launcher: - -```bash -SERVER_BIN=server/scripts/rocprof_server_wrapper.sh \ -PROFILED_SERVER_BIN=/path/to/dflash_server \ -ROCPROF_OUTPUT_DIR=/path/to/trace-output \ -ROCPROF_START_SECONDS=180 \ -ROCPROF_DURATION_SECONDS=90 \ -server/scripts/qualify_ds4_q5_amd.sh - -server/scripts/analyze_rocprof_overlap.py \ - /path/to/trace-output/trace_kernel_trace.csv -``` - -The analyzer reports per-owner busy time, simultaneous kernel-busy time, -time-binned overlap, and the kernels dominating each owner. Use a steady decode -window rather than model load or prefill when comparing placement changes. - -In the post-batching trace, steady 2K decode windows placed only 16-22% of -either owner's kernel-busy time inside a simultaneously busy interval. The -unprofiled server timing attributed 63.7 ms of each 74.4 ms speculative step to -target verification; draft, head, snapshot, and apply work together accounted -for the remaining 10.7 ms. Changing the placement rate from 4.4 to 3.8 at the -same 14,350 MiB budget moved 116 experts to the peer but changed the measured -2K median by less than 0.1 tok/s. These measurements show that placement is -already near its local balance point. Further large gains require removing -split/copy dispatches or parallelizing work outside the routed-expert fork; -adding the two devices' headline bandwidths is not a valid throughput model -because attention, routing, HC boundaries, and every layer join remain ordered. +MMVQ/cache/kernel defaults and the explicit hot-36 hardware profile, measured +medians were 67.957, 65.927, 62.544, 56.335, and 67.458 tok/s. Treat these as +workload-specific burn-in measurements, not as a portable default for unrelated +AMD memory layouts. On HIP `gfx1151`, enabling DSpark defaults `LUCE_MMVQ_MAX_NCOLS` to `4` when the variable is unset. This keeps the four-row verifier on MMVQ. On a 128 GiB diff --git a/server/docs/DS4_HETEROGENEOUS_ATTENTION_TP.md b/server/docs/DS4_HETEROGENEOUS_ATTENTION_TP.md new file mode 100644 index 000000000..e7c9329ed --- /dev/null +++ b/server/docs/DS4_HETEROGENEOUS_ATTENTION_TP.md @@ -0,0 +1,98 @@ +# DS4 heterogeneous attention TP + +This experiment extends the existing in-process expert split with real +attention parallelism across an R9700 and Strix Halo. It is opt-in and limited +to q=5 fused verification. + +## Qualified layout + +- R9700 owns 6 of the 8 output groups: 48 of 64 attention heads. +- Strix owns 2 groups: 16 heads (25%). +- Only compression-ratio-4 layers are split (21 of 43 layers at this model + shape). Their longer KV span is large enough to repay the device fork. +- QR and the assembled KV matrix are packed into one persistent F32 staging + packet. The packet is copied on the Strix stream after a main-stream event, + so the R9700 immediately continues its independent 75% branch. +- Strix performs its Q projection, score/softmax/value attention, inverse RoPE, + output-A, and its partial output-B projection. +- The peer returns one small partial residual. The existing HC-post operation + adds the main and peer results; there is no intermediate head join. + +`GGML_HIP_GRAPHS=ON` is mandatory. Disabling it reduced the same control from +88.538 to 83.208 tok/s and invalidated comparisons with the qualified runtime. + +## Reproduce + +Configure and build: + +```bash +cmake -S server -B server/build-hip-dual \ + -DGGML_HIP=ON \ + -DGGML_HIP_GRAPHS=ON +cmake --build server/build-hip-dual --target dflash_server \ + test_deepseek4_unit -j 12 +``` + +Set `TARGET_MODEL`, `DRAFT_MODEL`, `HOTNESS_CSV`, and optionally +`DECODE_HOTNESS_CSV`, then run: + +```bash +server/scripts/qualify_ds4_q5_amd_attention_tp.sh +``` + +The wrapper checks the CMake cache before starting and defaults to two warmups, +seven measured 2K/128-token runs. Override `TARGETS` to run the full context +sweep. + +## 2026-08-06 result on lucebox5 + +All accepted runs generated the exact expected SHA-256: +`0f785a7ffa406498aafb14553966eaed0f52220fed0f7cc016b66921d104d194`. + +| Configuration | 2K median client decode | Result | +| --- | ---: | --- | +| Same binary, attention TP off, HIP graphs on | 88.538 tok/s | control | +| 25%, packed fork, explicit attention, output-B on peer | 88.704 tok/s | initial three-run screen | +| Same 25% configuration, final seven-run confirmation | 88.472 tok/s | retained and stable | +| 25%, same split, output-B on main | 88.531 tok/s | correct, neutral | +| 12.5%, packed fork | 87.968 tok/s | under-filled peer | +| 25%, fused peer attention | 83.797 tok/s | rejected | +| 25%, graph-native late fork | 75.231 tok/s | rejected: peer starts too late | + +The first three-run screen measured 88.704 tok/s. The final two-warmup, +seven-measurement confirmation measured **88.472 tok/s** (88.309–88.611), with +7/7 correct hashes and no failed requests. Against the 88.538 same-binary +control, this is a tie within run noise. Describe the result as “more attention +overlap without a decode regression,” not as a large throughput win. +Raw result directories are under: + +`/home/lucebox5/ds4-attention-full-tp-proven-20260806/results/attention_tp/` + +The final confirmation run ID is +`attention-full-25-outputb-final-confirm-r7-teardownfix-2k-20260806ak`. + +The first long run exposed a cache-eviction teardown race after four measured +requests. Multi-backend scheduler events and gallocr buffers were being freed +without first quiescing both GPU streams; the next graph build then crashed in +the attention node list. Scheduler teardown and native-graph invalidation now +synchronize both backends first. The identical nine-request rerun completed. + +The matched kernel-level profile and corrected component attribution are in +[`DS4_HETEROGENEOUS_PROFILE_20260806.md`](DS4_HETEROGENEOUS_PROFILE_20260806.md). +It proves that the split raises R9700 overlap from 21.20% to 30.67%, but adds +more Strix work than it removes from the R9700. Do not use the fused verifier's +small coarse `attention_us` counter as total attention time. + +## Important rejected paths + +- Copying ordinary graph temporaries on the destination stream can race the + allocator. Only explicitly marked persistent staging tensors are safe. +- A normal dependency-tree join reduces scheduler segments, but records the + fork event after the complete main branch has already been queued. That + serializes Strix behind the R9700 despite the smaller graph. +- Global single-copy generation fences are unnecessary for the retained path + and lowered its median. +- Moving 50% of the heads overloads Strix; moving 12.5% leaves useful Strix + bandwidth idle. Two output groups are the measured balance point. +- Peer flash attention is exact for this qualification but materially slower + than the explicit score/softmax/value kernels on gfx1151. diff --git a/server/docs/DS4_HETEROGENEOUS_PROFILE_20260806.md b/server/docs/DS4_HETEROGENEOUS_PROFILE_20260806.md new file mode 100644 index 000000000..f6d726e13 --- /dev/null +++ b/server/docs/DS4_HETEROGENEOUS_PROFILE_20260806.md @@ -0,0 +1,133 @@ +# DS4 heterogeneous decode profile — 2026-08-06 + +This note records the kernel-level profile of the qualified R9700 + Strix Halo +q=5 path. Its main conclusion is that attention head splitting does create real +concurrent work, but attention is not the throughput bottleneck. The routed +expert kernels are the largest optimization target. + +## Measurement method + +`rocprofv3` collected kernel and memory-copy traces from matched control and +25% attention-split runs. The publication client records Linux monotonic-clock +timestamps for request start, first token, and request completion. The analyzer +selects only the first-token-to-completion intervals of measured requests, so +model loading, prefill, warmups, and time between requests are excluded. + +Three measured 2K-context, 128-output-token requests were present in each +short trace. All six responses matched SHA-256 +`0f785a7ffa406498aafb14553966eaed0f52220fed0f7cc016b66921d104d194`. +Profiler throughput is lower than unprofiled throughput and is used only for +matched attribution. + +Agent 1 is the R9700 (`gfx1201`); Agent 2 is the Strix Halo (`gfx1151`). + +## Exact decode-window result + +| Measurement | Control | 25% attention split | +| --- | ---: | ---: | +| Selected decode span, three requests | 5.658 s | 5.780 s | +| R9700 busy | 3.313 s | 3.225 s | +| Strix busy | 0.897 s | 1.335 s | +| Both devices busy | 0.703 s | 0.989 s | +| R9700 work overlapped | 21.20% | 30.67% | +| Strix work overlapped | 78.34% | 74.07% | +| Summed R9700 dispatch work per request | 1104.782 ms | 1075.195 ms | +| Summed Strix dispatch work per request | 298.961 ms | 445.068 ms | + +The split increases simultaneous device work by about **95.3 ms per request**. +It removes only **29.6 ms per request** from the R9700 while adding +**146.1 ms per request** to Strix. The peer branch therefore extends past the +work it was intended to hide. This is why increased overlap does not produce a +throughput gain: useful overlap increased, but total work increased more. + +The packed direct-KV path also performs 117.055 ms of device-to-device copies +over the three split requests, or about **39.0 ms per request**. Removing that +copy alone is insufficient: the added Strix attention kernels still cost much +more than the R9700 work they replace. + +The main added Strix work per request was approximately: + +- attention GEMMs: 41.6 ms; +- ROCmFP4-fast matrix-vector work: 27.4 ms; +- Q projection: 16.3 ms; +- RoPE: 8.5 ms; +- remaining copies, softmax, normalization, and joins: the balance. + +The R9700's largest removed individual contribution was 37.9 ms per request, +followed by smaller matrix-vector, softmax, projection, and RoPE reductions. +Smaller 48-head R9700 shapes also select less efficient kernels, so removed +logical work does not translate linearly into saved dispatch time. + +## Telemetry warning + +Do not interpret the coarse `attention_us` counter as the complete fused +attention cost. In q=5 fused verification, the large scheduler graph contains +attention, expert work, and joins together. The small counter is populated by +fallback/replay bookkeeping and may also be divided across speculative steps. +It cannot support a claim that attention costs 0.7–0.9 ms. Use timestamped +kernel traces and matched request windows for component attribution. + +## Rejected experiments + +| Experiment | Exact 2K median | Decision | +| --- | ---: | --- | +| Persistent peer KV cache, ordinary replay policy | 60.32 tok/s | reject | +| Persistent peer KV cache, forced property-scan bypass | 60.246 tok/s | reject | +| Five-token same-expert weight-reuse kernel | 79.380 tok/s | reject | +| Two-token task-parallel weight-reuse kernel | 78.068 tok/s | reject | + +The peer-cache design removed the full-history transfer but changed scheduler +topology and lost native HIP-graph replay for the main verification graph. The +property-scan bypass could not recover it. The two weight-reuse kernels were +correct, but serializing multiple activation dot products per wave reduced +occupancy and parallelism more than shared weight decoding reduced memory +traffic. All rejected paths remain disabled; the two-token source experiment +was removed after qualification. + +## Real bottleneck and next target + +The dominant active-device work is the full-width routed-expert ROCmFP2 and +ROCmFP3 matrix-vector computation. In the control trace, those kernels consume +roughly 0.69 s on the R9700 and 0.80 s on Strix across three requests, before +their launch-preparation and combine work. The peer is already almost fully +hidden when it runs, while the R9700 has the long tail. + +The next useful optimization must make the existing full-width expert kernels +faster without reducing occupancy, or move a coarse independent stage whose +complete peer branch finishes before the R9700 tail. Small attention/shared +matrix shards and per-expert serialization are measured dead ends on this +hardware pair. + +## Reproduce the analysis + +Run the overlap analyzer on each trace: + +```bash +server/scripts/analyze_rocprof_overlap.py \ + rocprof/trace_kernel_trace.csv \ + --requests-json decode-client.json \ + --memory-copy-trace rocprof/trace_memory_copy_trace.csv \ + --top 30 +``` + +Compare matched kernel work: + +```bash +server/scripts/compare_rocprof_decode.py \ + CONTROL/rocprof/trace_kernel_trace.csv CONTROL/decode-client.json \ + CANDIDATE/rocprof/trace_kernel_trace.csv CANDIDATE/decode-client.json \ + --top 30 +``` + +Raw artifacts are retained on lucebox5 under: + +```text +/home/lucebox5/ds4-attention-full-tp-proven-20260806/results/attention_tp/ + profile-control-kernel-short-20260806/ + profile-attention25-kernel-short-20260806/ + attention25-incremental-coherent-v2-2k-20260806/ + attention25-incremental-force-replay-2k-20260806/ +/home/lucebox5/ds4-attention-full-tp-proven-20260806/results/expert_kernel/ + control-mmid-group-reuse-2k-20260806/ + control-mmid-pair-reuse-r2-2k-20260806/ +``` diff --git a/server/docs/HETEROGENEOUS_STAGE_PLANNER.md b/server/docs/HETEROGENEOUS_STAGE_PLANNER.md index 6194cba08..95efa5f0c 100644 --- a/server/docs/HETEROGENEOUS_STAGE_PLANNER.md +++ b/server/docs/HETEROGENEOUS_STAGE_PLANNER.md @@ -104,6 +104,10 @@ enough to reach 100 tok/s; the next optimization must reduce full-width expert kernel time or overlap another independent stage, not add synchronization or small matrix shards. +The timestamped two-device kernel profile, attention-split attribution, and +subsequent rejected expert weight-reuse experiments are recorded in +[`DS4_HETEROGENEOUS_PROFILE_20260806.md`](DS4_HETEROGENEOUS_PROFILE_20260806.md). + ## Reproduction Use `scripts/qualify_ds4_q5_amd.sh`. The runner records the source commit, diff --git a/server/scripts/analyze_rocprof_overlap.py b/server/scripts/analyze_rocprof_overlap.py index 4d5d9f11a..ed8a80175 100755 --- a/server/scripts/analyze_rocprof_overlap.py +++ b/server/scripts/analyze_rocprof_overlap.py @@ -5,6 +5,7 @@ import argparse import csv +import json from collections import defaultdict from pathlib import Path @@ -72,12 +73,71 @@ def clipped_duration( return total +def duration_in_windows( + intervals: list[tuple[int, int]], windows: list[tuple[int, int]] +) -> int: + return sum(clipped_duration(intervals, start, end) for start, end in windows) + + +def row_duration_in_windows( + start: int, end: int, windows: list[tuple[int, int]] +) -> int: + total = 0 + for window_start, window_end in windows: + if window_end <= start: + continue + if window_start >= end: + break + total += max(0, min(end, window_end) - max(start, window_start)) + return total + + +def request_decode_windows( + path: Path, include_warmup: bool +) -> list[tuple[int, int, str]]: + payload = json.loads(path.read_text(encoding="utf-8")) + windows: list[tuple[int, int, str]] = [] + groups = payload.get("groups") + if groups is None: + groups = [{"target_context": "publication", "records": payload.get("records", [])}] + for group in groups: + target = group.get("target_context", "unknown") + for record in group.get("records", []): + if not record.get("ok") or (not include_warmup and not record.get("measured")): + continue + start = record.get("first_token_monotonic_ns") + end = record.get("request_end_monotonic_ns") + if start is None or end is None or int(end) <= int(start): + continue + label = f"ctx={target},index={record.get('index', '?')}" + windows.append((int(start), int(end), label)) + if not windows: + raise SystemExit( + "request JSON has no usable decode timestamps; rerun with the " + "instrumented publication client" + ) + windows.sort() + return windows + + def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("kernel_trace", type=Path) parser.add_argument("--bin-ms", type=float, default=1000.0) parser.add_argument("--window-start-s", type=float, default=0.0) parser.add_argument("--window-end-s", type=float) + parser.add_argument( + "--requests-json", type=Path, + help="select first-token-to-request-end windows from benchmark JSON", + ) + parser.add_argument( + "--include-warmup", action="store_true", + help="include warmup request windows when --requests-json is used", + ) + parser.add_argument( + "--memory-copy-trace", type=Path, + help="optional rocprofv3 memory-copy CSV summarized over the same windows", + ) parser.add_argument("--top", type=int, default=12) parser.add_argument( "--timeline-max", type=int, default=0, @@ -122,12 +182,29 @@ def main() -> int: window_end = min(trace_end, requested_end) if window_end <= window_start: raise SystemExit("selected window is empty") - span = window_end - window_start + labeled_windows: list[tuple[int, int, str]] + if args.requests_json: + labeled_windows = [] + for start, end, label in request_decode_windows( + args.requests_json, args.include_warmup + ): + start = max(start, window_start) + end = min(end, window_end) + if start < end: + labeled_windows.append((start, end, label)) + if not labeled_windows: + raise SystemExit("no request decode window overlaps the kernel trace") + else: + labeled_windows = [(window_start, window_end, "trace")] + selected_windows = merge_intervals( + [(start, end) for start, end, _ in labeled_windows] + ) + span = sum(end - start for start, end in selected_windows) busy = { - agent: clipped_duration(merged[agent], window_start, window_end) + agent: duration_in_windows(merged[agent], selected_windows) for agent in agents } - overlap_ns = clipped_duration(overlap, window_start, window_end) + overlap_ns = duration_in_windows(overlap, selected_windows) duration_by_kernel: dict[str, dict[str, int]] = defaultdict( lambda: defaultdict(int) @@ -139,17 +216,19 @@ def main() -> int: with args.kernel_trace.open(newline="") as handle: for row in csv.DictReader(handle): agent = row["Agent_Id"] - start = max(window_start, int(row["Start_Timestamp"])) - end = min(window_end, int(row["End_Timestamp"])) - if end <= start: + start = int(row["Start_Timestamp"]) + end = int(row["End_Timestamp"]) + duration = row_duration_in_windows(start, end, selected_windows) + if duration <= 0: continue name = row["Kernel_Name"] - duration_by_kernel[agent][name] += end - start + duration_by_kernel[agent][name] += duration count_by_kernel[agent][name] += 1 print( - f"window_s={(window_start-trace_start)/1e9:.3f}:" - f"{(window_end-trace_start)/1e9:.3f} span_s={span/1e9:.3f}" + f"window_s={(selected_windows[0][0]-trace_start)/1e9:.3f}:" + f"{(selected_windows[-1][1]-trace_start)/1e9:.3f} " + f"selected_span_s={span/1e9:.3f} windows={len(selected_windows)}" ) for agent in agents: print( @@ -157,26 +236,47 @@ def main() -> int: f"utilization={100.0*busy[agent]/span:.2f}%" ) union_ns = busy[agents[0]] + busy[agents[1]] - overlap_ns + agent0_only = busy[agents[0]] - overlap_ns + agent1_only = busy[agents[1]] - overlap_ns + idle_ns = span - union_ns print( f"both_busy_s={overlap_ns/1e9:.3f} " f"overlap_of_{agents[0]}={100.0*overlap_ns/max(1,busy[agents[0]]):.2f}% " f"overlap_of_{agents[1]}={100.0*overlap_ns/max(1,busy[agents[1]]):.2f}% " f"either_busy_s={union_ns/1e9:.3f}" ) + print( + f"{agents[0]}_only_s={agent0_only/1e9:.3f} " + f"{agents[1]}_only_s={agent1_only/1e9:.3f} " + f"neither_busy_s={idle_ns/1e9:.3f}" + ) bin_ns = max(1, int(args.bin_ms * 1e6)) print("bin_start_s,agent1_busy_pct,agent2_busy_pct,both_busy_pct") - cursor = window_start - while cursor < window_end: - end = min(cursor + bin_ns, window_end) - width = end - cursor - print( - f"{(cursor-trace_start)/1e9:.3f}," - f"{100.0*clipped_duration(merged[agents[0]], cursor, end)/width:.2f}," - f"{100.0*clipped_duration(merged[agents[1]], cursor, end)/width:.2f}," - f"{100.0*clipped_duration(overlap, cursor, end)/width:.2f}" - ) - cursor = end + for selected_start, selected_end in selected_windows: + cursor = selected_start + while cursor < selected_end: + end = min(cursor + bin_ns, selected_end) + width = end - cursor + print( + f"{(cursor-trace_start)/1e9:.3f}," + f"{100.0*clipped_duration(merged[agents[0]], cursor, end)/width:.2f}," + f"{100.0*clipped_duration(merged[agents[1]], cursor, end)/width:.2f}," + f"{100.0*clipped_duration(overlap, cursor, end)/width:.2f}" + ) + cursor = end + + if args.requests_json: + print("request_decode,label,start_s,span_s,agent1_busy_pct,agent2_busy_pct,both_busy_pct") + for start, end, label in labeled_windows: + width = end - start + print( + f"request_decode,{label},{(start-trace_start)/1e9:.6f}," + f"{width/1e9:.6f}," + f"{100.0*clipped_duration(merged[agents[0]], start, end)/width:.2f}," + f"{100.0*clipped_duration(merged[agents[1]], start, end)/width:.2f}," + f"{100.0*clipped_duration(overlap, start, end)/width:.2f}" + ) for agent in agents: print(f"top_kernels_{agent}") @@ -188,6 +288,32 @@ def main() -> int: f"{duration/1e9:.6f}s count={count_by_kernel[agent][name]} {name}" ) + if args.memory_copy_trace: + copy_duration: dict[tuple[str, str, str], int] = defaultdict(int) + copy_count: dict[tuple[str, str, str], int] = defaultdict(int) + with args.memory_copy_trace.open(newline="") as handle: + for row in csv.DictReader(handle): + start = int(row["Start_Timestamp"]) + end = int(row["End_Timestamp"]) + duration = row_duration_in_windows(start, end, selected_windows) + if duration <= 0: + continue + key = ( + row["Direction"], row["Source_Agent_Id"], + row["Destination_Agent_Id"], + ) + copy_duration[key] += duration + copy_count[key] += 1 + print("memory_copies") + for key, duration in sorted( + copy_duration.items(), key=lambda item: item[1], reverse=True + ): + direction, source, destination = key + print( + f"{duration/1e6:.3f}ms count={copy_count[key]} " + f"{direction} {source}->{destination}" + ) + if args.timeline_max > 0: gap_ns = max(0, int(args.timeline_merge_gap_us * 1e3)) bursts: list[tuple[int, int, str]] = [] @@ -195,10 +321,11 @@ def main() -> int: for start, end in merge_nearby_intervals( intervals_by_agent[agent], gap_ns ): - start = max(start, window_start) - end = min(end, window_end) - if start < end: - bursts.append((start, end, agent)) + for window_start, window_end in selected_windows: + clipped_start = max(start, window_start) + clipped_end = min(end, window_end) + if clipped_start < clipped_end: + bursts.append((clipped_start, clipped_end, agent)) bursts.sort() print( "timeline_start_s,duration_us,agent," diff --git a/server/scripts/compare_rocprof_decode.py b/server/scripts/compare_rocprof_decode.py new file mode 100755 index 000000000..1d782dea6 --- /dev/null +++ b/server/scripts/compare_rocprof_decode.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python3 +"""Compare per-kernel decode work between two timestamped rocprof traces.""" + +from __future__ import annotations + +import argparse +import csv +from collections import defaultdict +from pathlib import Path + +from analyze_rocprof_overlap import request_decode_windows, row_duration_in_windows + + +def collect( + trace: Path, requests: Path +) -> tuple[dict[tuple[str, str], int], dict[tuple[str, str], int], int]: + labeled = request_decode_windows(requests, include_warmup=False) + windows = [(start, end) for start, end, _ in labeled] + durations: dict[tuple[str, str], int] = defaultdict(int) + counts: dict[tuple[str, str], int] = defaultdict(int) + trace_start: int | None = None + trace_end: int | None = None + with trace.open(newline="") as handle: + for row in csv.DictReader(handle): + start = int(row["Start_Timestamp"]) + end = int(row["End_Timestamp"]) + trace_start = start if trace_start is None else min(trace_start, start) + trace_end = end if trace_end is None else max(trace_end, end) + duration = row_duration_in_windows(start, end, windows) + if duration <= 0: + continue + key = (row["Agent_Id"], row["Kernel_Name"]) + durations[key] += duration + counts[key] += 1 + if trace_start is None or trace_end is None: + raise SystemExit(f"empty kernel trace: {trace}") + request_count = sum( + 1 for start, end in windows + if max(start, trace_start) < min(end, trace_end) + ) + if request_count == 0: + raise SystemExit(f"no measured request overlaps trace: {trace}") + return durations, counts, request_count + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("control_trace", type=Path) + parser.add_argument("control_requests", type=Path) + parser.add_argument("candidate_trace", type=Path) + parser.add_argument("candidate_requests", type=Path) + parser.add_argument("--top", type=int, default=20) + args = parser.parse_args() + + control_durations, control_counts, control_n = collect( + args.control_trace, args.control_requests + ) + candidate_durations, candidate_counts, candidate_n = collect( + args.candidate_trace, args.candidate_requests + ) + agents = sorted( + {key[0] for key in control_durations} | + {key[0] for key in candidate_durations} + ) + print(f"control_requests={control_n} candidate_requests={candidate_n}") + for agent in agents: + control_total = sum( + duration for (key_agent, _), duration in control_durations.items() + if key_agent == agent + ) / control_n + candidate_total = sum( + duration for (key_agent, _), duration in candidate_durations.items() + if key_agent == agent + ) / candidate_n + print( + f"{agent} summed_dispatch_ms_per_request " + f"control={control_total/1e6:.3f} " + f"candidate={candidate_total/1e6:.3f} " + f"delta={(candidate_total-control_total)/1e6:+.3f}" + ) + + differences: list[tuple[float, str, float, float, float, float]] = [] + names = { + name for key_agent, name in control_durations if key_agent == agent + } | { + name for key_agent, name in candidate_durations if key_agent == agent + } + for name in names: + key = (agent, name) + control_ms = control_durations.get(key, 0) / control_n / 1e6 + candidate_ms = candidate_durations.get(key, 0) / candidate_n / 1e6 + control_count = control_counts.get(key, 0) / control_n + candidate_count = candidate_counts.get(key, 0) / candidate_n + differences.append(( + candidate_ms - control_ms, name, control_ms, candidate_ms, + control_count, candidate_count, + )) + + print(f"largest_added_{agent}") + for delta, name, control_ms, candidate_ms, control_count, candidate_count in sorted( + differences, reverse=True + )[: args.top]: + print( + f"{delta:+.3f}ms/request count={control_count:.1f}->{candidate_count:.1f} " + f"time={control_ms:.3f}->{candidate_ms:.3f} {name}" + ) + print(f"largest_removed_{agent}") + for delta, name, control_ms, candidate_ms, control_count, candidate_count in sorted( + differences + )[: args.top]: + print( + f"{delta:+.3f}ms/request count={control_count:.1f}->{candidate_count:.1f} " + f"time={control_ms:.3f}->{candidate_ms:.3f} {name}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/server/scripts/ds4_publication_decode_client.py b/server/scripts/ds4_publication_decode_client.py index 76097d3d3..e002853d2 100755 --- a/server/scripts/ds4_publication_decode_client.py +++ b/server/scripts/ds4_publication_decode_client.py @@ -85,7 +85,9 @@ def stream_request( ) started = time.perf_counter() + started_ns = time.perf_counter_ns() first_token_at: float | None = None + first_token_at_ns: int | None = None text_parts: list[str] = [] usage: dict[str, Any] = {} finish_reason: str | None = None @@ -119,6 +121,7 @@ def stream_request( if piece: if first_token_at is None: first_token_at = time.perf_counter() + first_token_at_ns = time.perf_counter_ns() text_parts.append(str(piece)) except urllib.error.HTTPError as exc: return { @@ -130,6 +133,7 @@ def stream_request( return {"ok": False, "status": status, "error": repr(exc)} finished = time.perf_counter() + finished_ns = time.perf_counter_ns() text = "".join(text_parts) wall_s = finished - started ttft_s = first_token_at - started if first_token_at is not None else None @@ -144,6 +148,12 @@ def stream_request( "ok": status == 200 and bool(text) and saw_done, "status": status, "stream_done": saw_done, + # rocprofv3 dispatch timestamps use the same host monotonic clock. These + # boundaries let the overlap analyzer select decode only, excluding + # model loading, prefill, and collection-window transitions. + "request_start_monotonic_ns": started_ns, + "first_token_monotonic_ns": first_token_at_ns, + "request_end_monotonic_ns": finished_ns, "wall_s": round(wall_s, 6), "ttft_s": round(ttft_s, 6) if ttft_s is not None else None, "client_decode_s": round(decode_s, 6) if decode_s is not None else None, diff --git a/server/scripts/qualify_ds4_q5_amd.sh b/server/scripts/qualify_ds4_q5_amd.sh index c80addcb8..0cbebf78f 100755 --- a/server/scripts/qualify_ds4_q5_amd.sh +++ b/server/scripts/qualify_ds4_q5_amd.sh @@ -49,6 +49,28 @@ SHARED_FFN_PEER_FRACTION="${SHARED_FFN_PEER_FRACTION:-0}" FUSED_OWNER_RESIDUAL="${FUSED_OWNER_RESIDUAL:-0}" ALIGN_SHARED_IDS="${ALIGN_SHARED_IDS:-0}" EXPERT_TOP_K="${EXPERT_TOP_K:-4}" +ATTENTION_TP_GROUPS="${ATTENTION_TP_GROUPS:-0}" +ATTENTION_TP_RATIO="${ATTENTION_TP_RATIO:--1}" +ATTENTION_TP_VALUES="${ATTENTION_TP_VALUES:-0}" +ATTENTION_TP_DIRECT_KV="${ATTENTION_TP_DIRECT_KV:-0}" +ATTENTION_TP_FLASH="${ATTENTION_TP_FLASH:-0}" +ATTENTION_TP_OUTPUT_B="${ATTENTION_TP_OUTPUT_B:-0}" +ATTENTION_TP_PROJECTION_ONLY="${ATTENTION_TP_PROJECTION_ONLY:-0}" +ATTENTION_TP_CORE_ONLY="${ATTENTION_TP_CORE_ONLY:-0}" +ATTENTION_TP_DEBUG_AB="${ATTENTION_TP_DEBUG_AB:-0}" +SCHED_PROFILE="${SCHED_PROFILE:-0}" +DST_STREAM_PEER_COPIES="${DST_STREAM_PEER_COPIES:-0}" +ATTENTION_TP_PACKED_STAGE="${ATTENTION_TP_PACKED_STAGE:-0}" +SINGLE_COPY_EVENT_FENCES="${SINGLE_COPY_EVENT_FENCES:-0}" +DRAFT_GPU="${DRAFT_GPU:-0}" +DRAFT_SEPARATE_STREAM="${DRAFT_SEPARATE_STREAM:-0}" +DRAFT_LOW_PRIORITY="${DRAFT_LOW_PRIORITY:-0}" +DRAFT_OVERLAP_PROBE="${DRAFT_OVERLAP_PROBE:-0}" +DRAFT_OVERLAP_REUSE_CONTEXT="${DRAFT_OVERLAP_REUSE_CONTEXT:-0}" +DRAFT_DEVICE_CHAIN="${DRAFT_DEVICE_CHAIN:-0}" +MMID_GROUP_REUSE="${MMID_GROUP_REUSE:-0}" +DUPLICATE_HOT_ON_COLD="${DUPLICATE_HOT_ON_COLD:-1}" +FULL_COLD_PARALLEL="${FULL_COLD_PARALLEL:-1}" VERIFY_WIDTH=$((4 + Q5_VERIFY + 2 * Q6_VERIFY)) RUN_ID="${RUN_ID:-ds4-q${VERIFY_WIDTH}-fr${FORCE_GRAPH_REPLAY}-direct${DIRECT_INDEXER_TOPK}-radix${BLOCK_RADIX_TOPK}-x4p1${FP4_Q5_X4_PLUS1}-cp${CRITICAL_PATH_PLACEMENT}-r${MAIN_TO_PEER_RATE}-sf${SHARED_FFN_PEER_FRACTION}-or${FUSED_OWNER_RESIDUAL}-ai${ALIGN_SHARED_IDS}-$(date -u +%Y%m%dT%H%M%SZ)}" OUT_ROOT="${OUT_ROOT:-$CHECKOUT/results/ds4_q5_context_qualification}" @@ -91,6 +113,19 @@ case "$Q6_VERIFY" in 0|1) ;; *) echo "Q6_VERIFY must be 0 or 1" >&2; exit 2 ;; esac +for toggle in "$DRAFT_SEPARATE_STREAM" "$DRAFT_LOW_PRIORITY" \ + "$DRAFT_OVERLAP_PROBE" "$DRAFT_OVERLAP_REUSE_CONTEXT" \ + "$DRAFT_DEVICE_CHAIN" "$MMID_GROUP_REUSE" \ + "$DUPLICATE_HOT_ON_COLD" "$FULL_COLD_PARALLEL"; do + case "$toggle" in + 0|1) ;; + *) echo "draft stream/probe and MMID reuse controls must be 0 or 1" >&2; exit 2 ;; + esac +done +if ! [[ "$DRAFT_GPU" =~ ^[0-9]+$ ]]; then + echo "DRAFT_GPU must be a non-negative integer" >&2 + exit 2 +fi if ((Q5_VERIFY + Q6_VERIFY > 1)); then echo "Q5_VERIFY and Q6_VERIFY are mutually exclusive" >&2 exit 2 @@ -111,6 +146,62 @@ if [[ ! "$EXPERT_TOP_K" =~ ^[1-9][0-9]*$ ]] || ((EXPERT_TOP_K > 16)); then echo "EXPERT_TOP_K must be an integer from 1 through 16" >&2 exit 2 fi +if [[ ! "$ATTENTION_TP_GROUPS" =~ ^[0-9]+$ ]] || ((ATTENTION_TP_GROUPS > 7)); then + echo "ATTENTION_TP_GROUPS must be an integer from 0 through 7" >&2 + exit 2 +fi +if [[ ! "$ATTENTION_TP_RATIO" =~ ^(-1|[0-9]+)$ ]]; then + echo "ATTENTION_TP_RATIO must be -1 or a non-negative integer" >&2 + exit 2 +fi +case "$ATTENTION_TP_DEBUG_AB" in + 0|1) ;; + *) echo "ATTENTION_TP_DEBUG_AB must be 0 or 1" >&2; exit 2 ;; +esac +case "$ATTENTION_TP_VALUES" in + 0|1) ;; + *) echo "ATTENTION_TP_VALUES must be 0 or 1" >&2; exit 2 ;; +esac +case "$ATTENTION_TP_DIRECT_KV" in + 0|1) ;; + *) echo "ATTENTION_TP_DIRECT_KV must be 0 or 1" >&2; exit 2 ;; +esac +case "$ATTENTION_TP_FLASH" in + 0|1) ;; + *) echo "ATTENTION_TP_FLASH must be 0 or 1" >&2; exit 2 ;; +esac +case "$ATTENTION_TP_OUTPUT_B" in + 0|1) ;; + *) echo "ATTENTION_TP_OUTPUT_B must be 0 or 1" >&2; exit 2 ;; +esac +case "$ATTENTION_TP_PROJECTION_ONLY" in + 0|1) ;; + *) echo "ATTENTION_TP_PROJECTION_ONLY must be 0 or 1" >&2; exit 2 ;; +esac +case "$ATTENTION_TP_CORE_ONLY" in + 0|1) ;; + *) echo "ATTENTION_TP_CORE_ONLY must be 0 or 1" >&2; exit 2 ;; +esac +if ((ATTENTION_TP_PROJECTION_ONLY + ATTENTION_TP_CORE_ONLY > 1)); then + echo "ATTENTION_TP_PROJECTION_ONLY and ATTENTION_TP_CORE_ONLY are mutually exclusive" >&2 + exit 2 +fi +case "$SCHED_PROFILE" in + 0|1) ;; + *) echo "SCHED_PROFILE must be 0 or 1" >&2; exit 2 ;; +esac +case "$DST_STREAM_PEER_COPIES" in + 0|1) ;; + *) echo "DST_STREAM_PEER_COPIES must be 0 or 1" >&2; exit 2 ;; +esac +case "$ATTENTION_TP_PACKED_STAGE" in + 0|1) ;; + *) echo "ATTENTION_TP_PACKED_STAGE must be 0 or 1" >&2; exit 2 ;; +esac +case "$SINGLE_COPY_EVENT_FENCES" in + 0|1) ;; + *) echo "SINGLE_COPY_EVENT_FENCES must be 0 or 1" >&2; exit 2 ;; +esac if [[ ! "$DYNAMIC_MAIN_SLOTS" =~ ^[1-9][0-9]*$ ]] || ((DYNAMIC_MAIN_SLOTS > EXPERT_TOP_K)); then echo "DYNAMIC_MAIN_SLOTS must be an integer from 1 through EXPERT_TOP_K ($EXPERT_TOP_K)" >&2 @@ -224,6 +315,21 @@ server_env=( "DFLASH_DS4_MOE_TP=1" "DFLASH_DS4_MOE_TP_INPROC=1" "DFLASH_DS4_MOE_TP_GPU=1" + "DFLASH_DS4_ATTENTION_TP_GROUPS=$ATTENTION_TP_GROUPS" + "DFLASH_DS4_ATTENTION_TP_RATIO=$ATTENTION_TP_RATIO" + "DFLASH_DS4_ATTENTION_TP_VALUES=$ATTENTION_TP_VALUES" + "DFLASH_DS4_ATTENTION_TP_DIRECT_KV=$ATTENTION_TP_DIRECT_KV" + "DFLASH_DS4_ATTENTION_TP_FLASH=$ATTENTION_TP_FLASH" + "DFLASH_DS4_ATTENTION_TP_OUTPUT_B=$ATTENTION_TP_OUTPUT_B" + "DFLASH_DS4_ATTENTION_TP_PROJECTION_ONLY=$ATTENTION_TP_PROJECTION_ONLY" + "DFLASH_DS4_ATTENTION_TP_CORE_ONLY=$ATTENTION_TP_CORE_ONLY" + "DFLASH_DS4_ATTENTION_TP_DEBUG_AB=$ATTENTION_TP_DEBUG_AB" + "DFLASH_DS4_ATTENTION_TP_DST_STREAM_STAGE=$DST_STREAM_PEER_COPIES" + "DFLASH_DS4_ATTENTION_TP_PACKED_STAGE=$ATTENTION_TP_PACKED_STAGE" + "DFLASH_DS4_TP_SCHED_TRACE=$ATTENTION_TP_DEBUG_AB" + "GGML_SCHED_PROFILE=$SCHED_PROFILE" + "GGML_SCHED_PROFILE_MIN_SPLITS=100" + "GGML_CUDA_BATCH_PEER_COPY_TRACE=0" "DFLASH_EXPERT_BUDGET_MB=$EXPERT_BUDGET_MB" "DFLASH_DS4_HOTNESS_CSV=$HOTNESS_CSV" "DFLASH_DS4_TP_CAPTURE_CACHE_SLOTS=4" @@ -239,10 +345,12 @@ server_env=( "DFLASH_DS4_TP_COARSE_OWNER_SPLIT=0" "DFLASH_DS4_TP_NATIVE_ROUTE_WIDTH=1" "GGML_CUDA_BATCH_PEER_COPIES=1" - "DFLASH_MOE_DUPLICATE_HOT_ON_COLD=1" + "GGML_CUDA_DST_STREAM_PEER_COPIES=$DST_STREAM_PEER_COPIES" + "GGML_SCHED_SINGLE_COPY_EVENT_FENCES=$SINGLE_COPY_EVENT_FENCES" + "DFLASH_MOE_DUPLICATE_HOT_ON_COLD=$DUPLICATE_HOT_ON_COLD" "DFLASH_DS4_HYBRID_PREFILL_GPU_HC=1" "DFLASH_DS4_HYBRID_PREFILL_EAGER=1" - "DFLASH_MOE_FULL_COLD_PARALLEL=1" + "DFLASH_MOE_FULL_COLD_PARALLEL=$FULL_COLD_PARALLEL" "DFLASH_DS4_PREFILL_TRACE=0" "DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC=1" "DFLASH_DS4_PINNED_ROLLBACK=1" @@ -251,7 +359,13 @@ server_env=( "DFLASH_DS4_SPEC_Q=$VERIFY_WIDTH" "DFLASH_DS4_ADAPTIVE_WIDTH=0" "DFLASH_DS4_DRAFT=$DRAFT_MODEL" - "DFLASH_DS4_DRAFT_GPU=0" + "DFLASH_DS4_DRAFT_GPU=$DRAFT_GPU" + "DFLASH_DS4_DRAFT_SEPARATE_STREAM=$DRAFT_SEPARATE_STREAM" + "DFLASH_DS4_DRAFT_LOW_PRIORITY=$DRAFT_LOW_PRIORITY" + "DFLASH_DS4_DRAFT_OVERLAP_PROBE=$DRAFT_OVERLAP_PROBE" + "DFLASH_DS4_DRAFT_OVERLAP_REUSE_CONTEXT=$DRAFT_OVERLAP_REUSE_CONTEXT" + "DFLASH_DS4_DRAFT_DEVICE_CHAIN=$DRAFT_DEVICE_CHAIN" + "DFLASH_CUDA_MMVQ_MOE_GROUP_REUSE=$MMID_GROUP_REUSE" "DFLASH_DS4_DRAFT_CONTEXT_KV_CACHE=1" "DFLASH_MOE_FUSED_COMBINE=0" ) @@ -382,6 +496,28 @@ server_args=( echo "fused_owner_residual=$FUSED_OWNER_RESIDUAL" echo "align_shared_ids=$ALIGN_SHARED_IDS" echo "expert_top_k=$EXPERT_TOP_K" + echo "attention_tp_groups=$ATTENTION_TP_GROUPS" + echo "attention_tp_ratio=$ATTENTION_TP_RATIO" + echo "attention_tp_values=$ATTENTION_TP_VALUES" + echo "attention_tp_direct_kv=$ATTENTION_TP_DIRECT_KV" + echo "attention_tp_flash=$ATTENTION_TP_FLASH" + echo "attention_tp_output_b=$ATTENTION_TP_OUTPUT_B" + echo "attention_tp_projection_only=$ATTENTION_TP_PROJECTION_ONLY" + echo "attention_tp_core_only=$ATTENTION_TP_CORE_ONLY" + echo "attention_tp_debug_ab=$ATTENTION_TP_DEBUG_AB" + echo "dst_stream_peer_copies=$DST_STREAM_PEER_COPIES" + echo "attention_tp_packed_stage=$ATTENTION_TP_PACKED_STAGE" + echo "single_copy_event_fences=$SINGLE_COPY_EVENT_FENCES" + echo "sched_profile=$SCHED_PROFILE" + echo "draft_gpu=$DRAFT_GPU" + echo "draft_separate_stream=$DRAFT_SEPARATE_STREAM" + echo "draft_low_priority=$DRAFT_LOW_PRIORITY" + echo "draft_overlap_probe=$DRAFT_OVERLAP_PROBE" + echo "draft_overlap_reuse_context=$DRAFT_OVERLAP_REUSE_CONTEXT" + echo "draft_device_chain=$DRAFT_DEVICE_CHAIN" + echo "mmid_group_reuse=$MMID_GROUP_REUSE" + echo "duplicate_hot_on_cold=$DUPLICATE_HOT_ON_COLD" + echo "full_cold_parallel=$FULL_COLD_PARALLEL" echo "cache_slots=$CACHE_SLOTS" echo "mmvq_max_ncols=$MMVQ_MAX_NCOLS" echo "targets=$TARGETS" diff --git a/server/scripts/qualify_ds4_q5_amd_attention_tp.sh b/server/scripts/qualify_ds4_q5_amd_attention_tp.sh new file mode 100755 index 000000000..ad2d8141b --- /dev/null +++ b/server/scripts/qualify_ds4_q5_amd_attention_tp.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Reproduce the qualified R9700 + Strix Halo q=5 attention split. +# Caller must provide the same model/profile paths required by +# qualify_ds4_q5_amd.sh. Every setting remains overrideable for explicit A/Bs. + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +CHECKOUT="${CHECKOUT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" +BUILD_DIR="${BUILD_DIR:-$CHECKOUT/server/build-hip-dual}" + +cache="$BUILD_DIR/CMakeCache.txt" +if [[ ! -f "$cache" ]]; then + echo "missing build cache: $cache" >&2 + echo "configure the HIP build with -DGGML_HIP_GRAPHS=ON" >&2 + exit 2 +fi +if ! grep -qx 'GGML_HIP_GRAPHS:BOOL=ON' "$cache"; then + echo "attention TP qualification requires GGML_HIP_GRAPHS=ON" >&2 + echo "reconfigure with: cmake -S $CHECKOUT/server -B $BUILD_DIR -DGGML_HIP_GRAPHS=ON" >&2 + exit 2 +fi + +export CHECKOUT BUILD_DIR +export TARGETS="${TARGETS:-2048}" +export WARMUP="${WARMUP:-2}" +export RUNS="${RUNS:-7}" +export MAX_TOKENS="${MAX_TOKENS:-128}" + +export CRITICAL_PATH_PLACEMENT="${CRITICAL_PATH_PLACEMENT:-1}" +export MAIN_TO_PEER_RATE="${MAIN_TO_PEER_RATE:-4.4}" +export EXPERT_BUDGET_MB="${EXPERT_BUDGET_MB:-14350}" +export DYNAMIC_ROUTE_BALANCE="${DYNAMIC_ROUTE_BALANCE:-1}" +export DYNAMIC_MAIN_SLOTS="${DYNAMIC_MAIN_SLOTS:-3}" +export DYNAMIC_MAIN_SLOTS_X4="${DYNAMIC_MAIN_SLOTS_X4:-13}" +export FUSED_OWNER_RESIDUAL="${FUSED_OWNER_RESIDUAL:-1}" + +# Two of eight output groups = 16 of 64 heads (25%). Ratio-4 layers are the +# only layers whose longer attention span amortized the heterogeneous fork. +export ATTENTION_TP_GROUPS="${ATTENTION_TP_GROUPS:-2}" +export ATTENTION_TP_RATIO="${ATTENTION_TP_RATIO:-4}" +export ATTENTION_TP_VALUES="${ATTENTION_TP_VALUES:-1}" +export ATTENTION_TP_DIRECT_KV="${ATTENTION_TP_DIRECT_KV:-1}" +export ATTENTION_TP_FLASH="${ATTENTION_TP_FLASH:-0}" +export ATTENTION_TP_OUTPUT_B="${ATTENTION_TP_OUTPUT_B:-1}" +export ATTENTION_TP_PROJECTION_ONLY="${ATTENTION_TP_PROJECTION_ONLY:-0}" +export ATTENTION_TP_CORE_ONLY="${ATTENTION_TP_CORE_ONLY:-0}" +export DST_STREAM_PEER_COPIES="${DST_STREAM_PEER_COPIES:-1}" +export ATTENTION_TP_PACKED_STAGE="${ATTENTION_TP_PACKED_STAGE:-1}" + +# The packed destination-stream fork and deferred result join already provide +# the two required lifetime edges. Global generation fences only add overhead. +export SINGLE_COPY_EVENT_FENCES="${SINGLE_COPY_EVENT_FENCES:-0}" +export FORCE_GRAPH_REPLAY="${FORCE_GRAPH_REPLAY:-0}" + +exec bash "$SCRIPT_DIR/qualify_ds4_q5_amd.sh" diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index f0df52c10..8455d0ddc 100644 --- a/server/src/common/dspark_head.cpp +++ b/server/src/common/dspark_head.cpp @@ -186,12 +186,35 @@ struct MarkovChainGraph { std::vector confidence; // optional sigmoid score per depth }; +// Make a leaf tensor in the head graph that borrows an already allocated +// device buffer. A regular ggml_view_* keeps the producer tensor as a graph +// dependency; for a cached draft output that would recursively splice the +// entire draft graph into the Markov graph. The producer has already been +// submitted on the same backend stream, so the head only needs a non-owning +// leaf alias with matching storage metadata. +ggml_tensor * borrow_device_prefix_2d(ggml_context * ctx, + ggml_tensor * src, + int64_t ne0, + int64_t ne1) { + if (!ctx || !src || !src->buffer || !src->data || + src->ne[0] != ne0 || src->ne[1] < ne1) { + return nullptr; + } + ggml_tensor * borrowed = ggml_new_tensor_2d(ctx, src->type, ne0, ne1); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + borrowed->nb[i] = src->nb[i]; + } + borrowed->buffer = src->buffer; + borrowed->data = src->data; + borrowed->extra = src->extra; + return borrowed; +} + // Guards shared by the fused Markov paths: head present, usable inputs, and // the target lm_head vocab matching the head's training vocab. bool dspark_fused_usable(const DraftWeights & dw, ggml_backend_t backend, - ggml_tensor * lm_head, const float * hidden, - const char * who) { - if (!dw.dspark.enabled || !hidden || !backend || !lm_head) return false; + ggml_tensor * lm_head, const char * who) { + if (!dw.dspark.enabled || !backend || !lm_head) return false; if (!dw.dspark.markov_w1 || !dw.dspark.markov_w2) return false; if (dw.n_embd <= 0 || dw.dspark.markov_rank <= 0) return false; const int vocab = (int)lm_head->ne[1]; @@ -222,7 +245,9 @@ bool build_markov_chain_graph(const DraftWeights & dw, bool corrected_are_outputs, bool confidence_are_outputs, std::vector & arena, - MarkovChainGraph & out) { + MarkovChainGraph & out, + ggml_tensor * device_hidden = nullptr, + ggml_tensor * device_confidence_hidden = nullptr) { const int hdim = dw.n_embd; const int vocab = (int)lm_head->ne[1]; const int n_corr = n_positions - first_corrected; @@ -245,14 +270,51 @@ bool build_markov_chain_graph(const DraftWeights & dw, if (!out.ctx) return false; out.gf = ggml_new_graph_custom(out.ctx, 512, false); - out.inp_hidden = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hdim, n_positions); - if (have_confidence) { - out.inp_confidence_hidden = + if (device_hidden) { + if (device_hidden->type != GGML_TYPE_F32 || + device_hidden->ne[0] != hdim || + device_hidden->ne[1] < n_positions || + device_hidden->nb[0] != sizeof(float)) { + ggml_free(out.ctx); + out.ctx = nullptr; + return false; + } + out.inp_hidden = borrow_device_prefix_2d( + out.ctx, device_hidden, hdim, n_positions); + if (!out.inp_hidden) { + ggml_free(out.ctx); + out.ctx = nullptr; + return false; + } + } else { + out.inp_hidden = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, hdim, n_positions); - ggml_set_input(out.inp_confidence_hidden); + ggml_set_input(out.inp_hidden); + } + if (have_confidence) { + if (device_confidence_hidden) { + if (device_confidence_hidden->type != GGML_TYPE_F32 || + device_confidence_hidden->ne[0] != hdim || + device_confidence_hidden->ne[1] < n_positions || + device_confidence_hidden->nb[0] != sizeof(float)) { + ggml_free(out.ctx); + out.ctx = nullptr; + return false; + } + out.inp_confidence_hidden = borrow_device_prefix_2d( + out.ctx, device_confidence_hidden, hdim, n_positions); + if (!out.inp_confidence_hidden) { + ggml_free(out.ctx); + out.ctx = nullptr; + return false; + } + } else { + out.inp_confidence_hidden = ggml_new_tensor_2d( + out.ctx, GGML_TYPE_F32, hdim, n_positions); + ggml_set_input(out.inp_confidence_hidden); + } } out.inp_seed = ggml_new_tensor_1d(out.ctx, GGML_TYPE_I32, 1); - ggml_set_input(out.inp_hidden); ggml_set_input(out.inp_seed); out.base = ggml_mul_mat(out.ctx, lm_head, out.inp_hidden); @@ -317,7 +379,10 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, std::vector * confidence_out, const float * confidence_hidden) { if (q_len <= 1) return false; - if (!dspark_fused_usable(dw, backend, lm_head, local_hidden, "dspark_fused")) return false; + if (!local_hidden || + !dspark_fused_usable(dw, backend, lm_head, "dspark_fused")) { + return false; + } const int hdim = dw.n_embd; const int n_cand = q_len - 1; @@ -382,6 +447,80 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, return true; } +bool dspark_markov_correct_greedy_chain_fused_device( + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + ggml_tensor * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok, + std::vector * confidence_out, + ggml_tensor * confidence_hidden) { + if (q_len <= 1 || !local_hidden || + (confidence_out && !confidence_hidden) || + !dspark_fused_usable( + dw, backend, lm_head, "dspark_fused_device")) { + return false; + } + const int n_cand = q_len - 1; + + static thread_local std::vector g_arena_chain_device; + MarkovChainGraph g; + const bool want_confidence = confidence_out != nullptr; + if (confidence_out) confidence_out->clear(); + if (!build_markov_chain_graph( + dw, lm_head, n_cand, /*first_corrected=*/0, + /*corrected_are_outputs=*/false, + /*confidence_are_outputs=*/want_confidence, + g_arena_chain_device, g, local_hidden, confidence_hidden)) { + return false; + } + + static thread_local ggml_gallocr_t galloc_chain_device = nullptr; + if (!galloc_chain_device) { + galloc_chain_device = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + if (!ggml_gallocr_alloc_graph(galloc_chain_device, g.gf)) { + std::fprintf(stderr, + "dspark_fused_device: gallocr_alloc_graph failed\n"); + ggml_free(g.ctx); + return false; + } + + ggml_backend_tensor_set(g.inp_seed, &last_tok, 0, sizeof(int32_t)); + if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "dspark_fused_device: graph_compute failed\n"); + ggml_free(g.ctx); + return false; + } + + draft_tok.assign((size_t) q_len, 0); + draft_tok[0] = last_tok; + int32_t t_out[16]; + float c_out[16] = {}; + const int n_get = n_cand < 16 ? n_cand : 16; + for (int i = 0; i < n_get; ++i) { + ggml_backend_tensor_get_async( + backend, g.toks[(size_t) i], &t_out[i], 0, sizeof(int32_t)); + if (want_confidence && g.confidence[(size_t) i]) { + ggml_backend_tensor_get_async( + backend, g.confidence[(size_t) i], &c_out[i], 0, + sizeof(float)); + } + } + ggml_backend_synchronize(backend); + for (int i = 0; i < n_get; ++i) { + draft_tok[(size_t) i + 1] = t_out[i]; + } + if (want_confidence && !g.confidence.empty() && g.confidence[0]) { + confidence_out->assign(c_out, c_out + n_get); + } + ggml_free(g.ctx); + return true; +} + bool dspark_markov_project_topk(const DraftWeights & dw, ggml_backend_t backend, ggml_tensor * lm_head, @@ -391,7 +530,10 @@ bool dspark_markov_project_topk(const DraftWeights & dw, std::vector & top_log_probs, std::vector & top_token_ids) { if (n_tokens <= 1 || K <= 0) return false; - if (!dspark_fused_usable(dw, backend, lm_head, hidden, "dspark_topk")) return false; + if (!hidden || + !dspark_fused_usable(dw, backend, lm_head, "dspark_topk")) { + return false; + } const int hdim = dw.n_embd; const int vocab = (int)lm_head->ne[1]; diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h index 9b97b261d..ee9d310fd 100644 --- a/server/src/common/dspark_head.h +++ b/server/src/common/dspark_head.h @@ -36,6 +36,20 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, std::vector * confidence_out = nullptr, const float * confidence_hidden = nullptr); +// Device-resident sibling of the fused chain. The hidden tensors are outputs +// of a previously submitted graph on the same backend/stream, so stream order +// replaces the intermediate device-to-host-to-device round trip. +bool dspark_markov_correct_greedy_chain_fused_device( + const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + ggml_tensor * local_hidden, + int q_len, + int32_t last_tok, + std::vector & draft_tok, + std::vector * confidence_out = nullptr, + ggml_tensor * confidence_hidden = nullptr); + // DDTree candidate generation with the Markov correction: base logits for // all n_tokens positions in ONE lm_head matmul; rows 1..n-1 get the low-rank // previous-token bias chained along the main (argmax) path; top-K extracted diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 3562d8e2f..cdd8af51d 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -2,7 +2,6 @@ #include "deepseek4_backend.h" #include "deepseek4_internal.h" -#include "common/dynamic_backend.h" #include "common/peer_access.h" #include "common/sampler.h" @@ -22,7 +21,6 @@ #include #include #include -#include namespace dflash::common { @@ -42,6 +40,18 @@ static bool env_flag_enabled(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +static int attention_tp_peer_groups() { + const char * value = std::getenv("DFLASH_DS4_ATTENTION_TP_GROUPS"); + if (!value || !*value) return 0; + char * end = nullptr; + const long groups = std::strtol(value, &end, 10); + if (end == value || *end != '\0' || groups < 0 || + groups > std::numeric_limits::max()) { + return -1; + } + return (int) groups; +} + static bool positive_env_double(const char * name, double fallback, double & out, std::string * err) { out = fallback; @@ -138,71 +148,17 @@ static void configure_gfx1201_hybrid_sub_batch_default(int gpu) { #endif } -struct Ds4MoeTpConfig { - bool requested = false; - bool in_process = false; - bool backend_valid = true; - PlacementBackend secondary_backend = PlacementBackend::Auto; - int secondary_gpu = 0; - bool all_on_secondary = false; - bool concentrate_secondary = false; - bool profile_hot_on_secondary = false; -}; - -static Ds4MoeTpConfig ds4_moe_tp_config(int local_gpu) { - Ds4MoeTpConfig result; - result.requested = env_flag_enabled("DFLASH_DS4_MOE_TP"); - result.in_process = result.requested && - env_flag_enabled("DFLASH_DS4_MOE_TP_INPROC"); - result.all_on_secondary = result.requested && - env_flag_enabled("DFLASH_DS4_MOE_TP_ALL_COLD"); - result.concentrate_secondary = result.requested && - env_flag_enabled("DFLASH_DS4_MOE_TP_CONCENTRATE_COLD"); - result.profile_hot_on_secondary = result.in_process && - env_flag_enabled("DFLASH_DS4_MOE_TP_PEER_HOT"); - - const char * raw = std::getenv("DFLASH_DS4_MOE_TP_BACKEND"); - if (!raw || !*raw) raw = std::getenv("DFLASH_MOE_TP_BACKEND"); - if (!raw || !*raw) { -#if defined(DFLASH27B_BACKEND_MIXED) - result.secondary_backend = - compiled_placement_backend() == PlacementBackend::Cuda - ? PlacementBackend::Hip : PlacementBackend::Cuda; -#else - result.secondary_backend = compiled_placement_backend(); -#endif - } else { - result.backend_valid = parse_placement_backend( - raw, result.secondary_backend) && - result.secondary_backend != PlacementBackend::Auto; - } - - const char * gpu_raw = std::getenv("DFLASH_DS4_MOE_TP_GPU"); - if (!gpu_raw || !*gpu_raw) { - gpu_raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU"); - } - if (gpu_raw && *gpu_raw) { - result.secondary_gpu = std::max(0, std::atoi(gpu_raw)); - } else if (result.backend_valid && - result.secondary_backend != compiled_placement_backend()) { - // CUDA and HIP have independent device namespaces. The first device - // in the peer runtime is therefore backend:0 even when the target is - // also device zero in its own runtime. - result.secondary_gpu = 0; - } else { - result.secondary_gpu = local_gpu == 0 ? 1 : 0; - } - return result; +static bool ds4_inprocess_moe_tp_enabled() { + return env_flag_enabled("DFLASH_DS4_MOE_TP_INPROC"); } -static bool ds4_draft_backend(PlacementBackend & out) { - const char * raw = std::getenv("DFLASH_DS4_DRAFT_BACKEND"); +static int ds4_moe_tp_gpu(int local_gpu) { + const char * raw = std::getenv("DFLASH_DS4_MOE_TP_GPU"); if (!raw || !*raw) { - out = compiled_placement_backend(); - return true; + raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU"); } - return parse_placement_backend(raw, out) && - out != PlacementBackend::Auto; + if (raw && *raw) return std::max(0, std::atoi(raw)); + return local_gpu == 0 ? 1 : 0; } static double gib(uint64_t bytes) { @@ -421,45 +377,9 @@ static void fill_prefix_hot_placement(const DeepSeek4Weights & w, } } -// Cross-runtime joins are much more expensive than native peer handoffs. Keep -// approximately the same expert residency as the uniform placement, but -// concentrate the cold owner into complete layers. A partial cold layer costs -// another cross-runtime join and some CUDA prefill paths require a complete -// expert stack, so retain the small remainder on the target backend. -static int fill_concentrated_cold_placement(const DeepSeek4Weights & w, - int hot_per_layer, - MoeHybridPlacement & out) { - out = {}; - out.n_layer = w.n_layer; - out.n_expert = w.n_expert; - out.n_expert_used = w.n_expert_used; - out.hot_counts.assign((size_t) w.n_layer, w.n_expert); - out.hot_expert_ids.resize((size_t) w.n_layer); - - const int requested_cold = - w.n_layer * std::max(0, w.n_expert - hot_per_layer); - int cold_remaining = w.n_expert > 0 - ? requested_cold / w.n_expert * w.n_expert : 0; - const int retained_local = requested_cold - cold_remaining; - for (int il = w.n_layer - 1; il >= 0; --il) { - const int cold = std::min(w.n_expert, cold_remaining); - const int hot = w.n_expert - cold; - out.hot_counts[(size_t) il] = hot; - auto & ids = out.hot_expert_ids[(size_t) il]; - ids.reserve((size_t) hot); - for (int ie = 0; ie < hot; ++ie) { - ids.push_back((int32_t) ie); - } - out.total_hot += hot; - cold_remaining -= cold; - } - return retained_local; -} - static bool fill_profiled_hot_placement(const DeepSeek4Weights & w, int hot_per_layer, const char * profile_path, - bool profile_hot_on_secondary, MoeHybridPlacement & out, std::string * err) { MoeHybridRoutingStats stats; @@ -481,41 +401,9 @@ static bool fill_profiled_hot_placement(const DeepSeek4Weights & w, out.hot_expert_ids.resize((size_t)w.n_layer); out.total_hot = hot_per_layer * w.n_layer; for (int il = 0; il < w.n_layer; ++il) { + std::vector ranked = stats.hot_experts(il, hot_per_layer); auto & ids = out.hot_expert_ids[(size_t)il]; - if (!profile_hot_on_secondary) { - std::vector ranked = stats.hot_experts(il, hot_per_layer); - ids.assign(ranked.begin(), ranked.end()); - continue; - } - - // `hot` is the primary-backend side of MoeHybridPlacement. On a - // memory-rich iGPU paired with a smaller, faster dGPU, filling that - // primary side with the most frequently routed experts starves the - // dGPU of useful work. Reserve the peer-sized complement for the - // hottest experts and keep every other expert on the primary. This - // changes ownership only; route order and reduction semantics stay - // unchanged. - const int peer_count = w.n_expert - hot_per_layer; - const std::vector ranked_peer = - stats.hot_experts(il, peer_count); - std::vector on_peer((size_t)w.n_expert, 0); - for (int expert : ranked_peer) { - if (expert >= 0 && expert < w.n_expert) { - on_peer[(size_t)expert] = 1; - } - } - ids.reserve((size_t)hot_per_layer); - for (int expert = 0; expert < w.n_expert; ++expert) { - if (!on_peer[(size_t)expert]) { - ids.push_back((int32_t)expert); - } - } - if ((int)ids.size() != hot_per_layer) { - if (err) { - *err = "routing profile did not yield a complete expert ranking"; - } - return false; - } + ids.assign(ranked.begin(), ranked.end()); } return true; } @@ -524,23 +412,18 @@ static bool fill_profiled_hot_placement(const DeepSeek4Weights & w, // but distribute those slots across layers to minimize the predicted owner // critical path. Uniform expert counts are a poor fit for heterogeneous EP: // routing skew varies substantially by layer, while every layer joins on the -// slower of its primary/shared and secondary expert branches. +// slower of its R9700 hot/shared and Strix cold branches. // // The cost model intentionally uses measured bandwidth rather than advertised // peak bandwidth. It is only an allocation objective; actual placement still // uses authoritative router statistics and evaluates every selected expert. static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, - ggml_backend_t backend, + int gpu, int max_ctx, Ds4HybridBudgetInfo & out, std::string * err) { out = {}; - if (!backend || !ggml_backend_get_device(backend)) { - if (err) *err = "target backend has no device"; - return false; - } - ggml_backend_dev_memory( - ggml_backend_get_device(backend), &out.gpu_free, &out.gpu_total); + 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; @@ -550,8 +433,7 @@ static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, return false; } - out.core_bytes = moe_hybrid_core_bytes_from_memory( - "deepseek4", out.gpu_free, out.gpu_total); + out.core_bytes = out.gpu_total - out.gpu_free; out.kv_bytes = estimate_ds4_cache_bytes(w, max_ctx); if (out.gpu_total > out.core_bytes + out.kv_bytes + out.warm_bytes + out.safety_bytes) { @@ -657,7 +539,7 @@ bool DeepSeek4Backend::load_model() { // Fused decode and layer-major prefill normally require monolithic expert // residency. Heterogeneous TP is the exception: its fused graph owns the - // routed experts across two local GPU backends, so forcing a full load would + // routed experts across two HIP backends, so forcing a full load would // disable the requested split before the TP runtime can initialize. const bool force_full = env_flag_enabled("DFLASH_DS4_FORCE_FULL_LOAD"); const bool heterogeneous_tp = env_flag_enabled("DFLASH_DS4_MOE_TP"); @@ -683,9 +565,9 @@ bool DeepSeek4Backend::load_model() { cfg_.model_path); return false; } - } else if (target_backend == PlacementBackend::Hip || heterogeneous_tp) { + } else if (target_backend == PlacementBackend::Hip) { std::fprintf(stderr, - "[deepseek4] heterogeneous target detected; using hybrid expert load path\n"); + "[deepseek4] HIP target detected; using hybrid expert load path\n"); if (!init_hybrid_model()) { std::fprintf(stderr, "[deepseek4] hybrid mode failed: %s\n", cfg_.model_path); return false; @@ -729,37 +611,23 @@ bool DeepSeek4Backend::load_spec_drafter() { } const bool separate_draft_stream = env_flag_enabled("DFLASH_DS4_DRAFT_SEPARATE_STREAM"); - PlacementBackend draft_kind = PlacementBackend::Auto; - if (!ds4_draft_backend(draft_kind)) { - std::fprintf(stderr, - "[deepseek4] invalid DFLASH_DS4_DRAFT_BACKEND; " - "expected cuda or hip\n"); - return false; - } - const PlacementBackend target_kind = placement_backend_of(backend_); - if (draft_kind != target_kind || draft_gpu != cfg_.device.gpu || - separate_draft_stream) { - std::string backend_error; - spec_backend_ = init_placement_backend( - draft_kind, draft_gpu, &backend_error); + if (draft_gpu != cfg_.device.gpu || separate_draft_stream) { + spec_backend_ = ggml_backend_cuda_init(draft_gpu); if (!spec_backend_) { std::fprintf(stderr, - "[deepseek4] failed to initialize DSpark %s:%d: %s\n", - placement_backend_name(draft_kind), draft_gpu, - backend_error.c_str()); + "[deepseek4] failed to initialize DSpark GPU %d\n", + draft_gpu); return false; } draft_backend = spec_backend_; const bool low_priority = separate_draft_stream && env_flag_enabled("DFLASH_DS4_DRAFT_LOW_PRIORITY"); const bool priority_configured = low_priority && - backend_pair_capabilities(backend_, spec_backend_).same_runtime && ggml_backend_cuda_set_low_priority_stream(spec_backend_); std::fprintf(stderr, - "[deepseek4] DSpark backend=%s:%d target=%s:%d " + "[deepseek4] DSpark backend gpu=%d target_gpu=%d " "separate_stream=%d low_priority=%d\n", - placement_backend_name(draft_kind), draft_gpu, - placement_backend_name(target_kind), cfg_.device.gpu, + draft_gpu, cfg_.device.gpu, (int) separate_draft_stream, (int) priority_configured); } @@ -830,54 +698,6 @@ void DeepSeek4Backend::release_spec_drafter(bool mark_parked) { spec_drafter_parked_ = mark_parked && !spec_draft_path_.empty(); } -void DeepSeek4Backend::keep_spec_feature_tail( - std::vector & features, size_t max_rows) const { - if (!spec_drafter_) return; - const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; - if (feat_row <= 0 || features.size() % (size_t) feat_row != 0) { - features.clear(); - return; - } - const size_t rows = features.size() / (size_t) feat_row; - const size_t keep_rows = std::min(rows, max_rows); - if (rows == keep_rows) return; - const size_t keep_floats = keep_rows * (size_t) feat_row; - const size_t drop_floats = features.size() - keep_floats; - if (keep_floats > 0) { - std::memmove(features.data(), features.data() + drop_floats, - keep_floats * sizeof(float)); - } - features.resize(keep_floats); -} - -int DeepSeek4Backend::capture_safe_prefill_tokens( - int token_offset, - int requested_tokens, - int final_capture_from, - bool batch_final_capture, - bool snapshot_pending, - int snapshot_capture_from, - int snapshot_capture_to) { - if (requested_tokens <= 0) return 0; - - int safe_tokens = requested_tokens; - const auto split_at = [&](int boundary) { - const int distance = boundary - token_offset; - if (distance > 0 && distance < safe_tokens) { - safe_tokens = distance; - } - }; - - if (!batch_final_capture) { - split_at(final_capture_from); - } - if (snapshot_pending) { - split_at(snapshot_capture_from); - split_at(snapshot_capture_to); - } - return safe_tokens; -} - bool DeepSeek4Backend::init() { // The shared MMVQ/MMQ crossover defaults to q=3 for NVIDIA. On gfx1151, // DSpark q=4 is faster through MMVQ. Keep AR and other devices unchanged, @@ -974,8 +794,7 @@ bool DeepSeek4Backend::init_moe_tensor_parallel() { return false; } - const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); - if (tp.in_process) { + if (ds4_inprocess_moe_tp_enabled()) { if (!expert_backend_ || !moe_hybrid_->materialized_cold_experts || moe_hybrid_->cold_backend != expert_backend_) { std::fprintf(stderr, @@ -983,16 +802,10 @@ bool DeepSeek4Backend::init_moe_tensor_parallel() { return false; } expert_runtime_.reset(); - const PlacementBackend local_kind = - cfg_.device.backend == PlacementBackend::Auto - ? compiled_placement_backend() : cfg_.device.backend; std::fprintf(stderr, - "[deepseek4-moe-tp] enabled mode=in-process local=%s:%d " - "secondary=%s:%d primary_experts=%d " - "secondary_experts=%d\n", - placement_backend_name(local_kind), cfg_.device.gpu, - placement_backend_name(tp.secondary_backend), - tp.secondary_gpu, + "[deepseek4-moe-tp] enabled mode=in-process local_gpu=%d " + "expert_gpu=%d local_experts=%d remote_experts=%d\n", + cfg_.device.gpu, ds4_moe_tp_gpu(cfg_.device.gpu), moe_placement_.total_hot, w_.n_layer * w_.n_expert - moe_placement_.total_hot); return true; @@ -1036,40 +849,24 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & std::string * err) const { if (decode_out) *decode_out = {}; Ds4HybridBudgetInfo budget; - if (!compute_ds4_hybrid_budget_info(w, backend_, max_ctx, budget, err)) { + if (!compute_ds4_hybrid_budget_info(w, cfg_.device.gpu, max_ctx, budget, err)) { return false; } - const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); - int hot_per_layer = tp.all_on_secondary ? 0 : budget.max_hot_per_layer; - if (tp.all_on_secondary) { + const bool all_cold = env_flag_enabled("DFLASH_DS4_MOE_TP_ALL_COLD"); + int hot_per_layer = all_cold ? 0 : budget.max_hot_per_layer; + if (all_cold) { std::fprintf(stderr, - "[deepseek4-moe-tp] all routed experts assigned to the " - "secondary backend\n"); + "[deepseek4-moe-tp] all routed experts assigned to the cold backend\n"); } - const bool concentrate_requested = tp.concentrate_secondary; - bool concentrated = false; - int retained_local = 0; const char * profile_path = std::getenv("DFLASH_DS4_HOTNESS_CSV"); const char * decode_profile_path = std::getenv("DFLASH_DS4_DECODE_HOTNESS_CSV"); const bool phase_aware_placement = decode_profile_path && *decode_profile_path; - const bool critical_path_placement = - !tp.all_on_secondary && !concentrate_requested && + const bool critical_path_placement = !all_cold && env_flag_enabled("DFLASH_DS4_TP_CRITICAL_PATH_PLACEMENT"); - const int requested_cold = - w.n_layer * std::max(0, w.n_expert - hot_per_layer); - if (concentrate_requested && requested_cold >= w.n_expert) { - retained_local = - fill_concentrated_cold_placement(w, hot_per_layer, out); - concentrated = true; - } else if (concentrate_requested) { - std::fprintf(stderr, - "[deepseek4] concentrated secondary placement needs at least " - "one complete layer; using uniform placement\n"); - fill_prefix_hot_placement(w, hot_per_layer, out); - } else if (critical_path_placement) { + if (critical_path_placement) { if (!profile_path || !*profile_path) { if (err) { *err = "critical-path placement requires DFLASH_DS4_HOTNESS_CSV"; @@ -1179,19 +976,13 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & std::fprintf(stderr, "\n"); } else if (profile_path) { if (*profile_path) { - const bool profile_hot_on_secondary = - tp.in_process && tp.profile_hot_on_secondary; if (!fill_profiled_hot_placement( - w, hot_per_layer, profile_path, - profile_hot_on_secondary, - out, err)) { + w, hot_per_layer, profile_path, out, err)) { return false; } std::fprintf(stderr, - "[deepseek4] hybrid placement profile=%s%s\n", - profile_path, - profile_hot_on_secondary - ? " profile-hot-owner=secondary" : ""); + "[deepseek4] hybrid placement profile=%s\n", + profile_path); } else { fill_prefix_hot_placement(w, hot_per_layer, out); } @@ -1203,28 +994,6 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & if (!compute_ds4_expert_memory_info(w, &out, placed_mem, err)) { return false; } - if (concentrated && placed_mem.hot_bytes > budget.expert_budget) { - std::fprintf(stderr, - "[deepseek4] concentrated secondary placement exceeds the " - "primary expert budget; using uniform placement\n"); - fill_prefix_hot_placement(w, hot_per_layer, out); - if (!compute_ds4_expert_memory_info(w, &out, placed_mem, err)) { - return false; - } - concentrated = false; - } - if (concentrated) { - const int cold_layers = - w.n_expert > 0 - ? (w.n_layer * w.n_expert - out.total_hot) / w.n_expert : 0; - std::fprintf(stderr, - "[deepseek4] concentrated secondary placement: " - "cross-owner layers=%d primary_experts=%d " - "secondary_experts=%d retained_primary=%d\n", - cold_layers, out.total_hot, - w.n_layer * w.n_expert - out.total_hot, - retained_local); - } const std::string hot_label = critical_path_placement ? "balanced" : std::to_string(hot_per_layer); @@ -1271,47 +1040,26 @@ bool DeepSeek4Backend::init_hybrid_model() { auto hybrid = std::make_shared(); MoeHybridConfig hybrid_cfg = make_ds4_parent_worker_cfg(w_); - const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); - const bool inprocess_tp = tp.requested && tp.in_process; + const bool inprocess_tp = + env_flag_enabled("DFLASH_DS4_MOE_TP") && ds4_inprocess_moe_tp_enabled(); if (inprocess_tp) { - const int expert_gpu = tp.secondary_gpu; - const PlacementBackend expert_kind = tp.secondary_backend; - if (!tp.backend_valid) { - std::fprintf(stderr, - "[deepseek4-moe-tp] invalid DFLASH_DS4_MOE_TP_BACKEND; " - "expected cuda or hip\n"); - return false; - } - const PlacementBackend local_kind = - cfg_.device.backend == PlacementBackend::Auto - ? compiled_placement_backend() : cfg_.device.backend; - if (expert_kind == local_kind && expert_gpu == cfg_.device.gpu) { + const int expert_gpu = ds4_moe_tp_gpu(cfg_.device.gpu); + if (expert_gpu == cfg_.device.gpu) { std::fprintf(stderr, - "[deepseek4-moe-tp] in-process secondary device must " - "differ from the primary device\n"); + "[deepseek4-moe-tp] in-process expert GPU must differ from local GPU\n"); return false; } - if (expert_kind == local_kind && g_peer_access_opt_in) { + if (g_peer_access_opt_in) { const bool peer_ok = enable_peer_access_pair(cfg_.device.gpu, expert_gpu); std::fprintf(stderr, - "[deepseek4-moe-tp] peer access %s:%d <-> %s:%d: %s\n", - placement_backend_name(local_kind), cfg_.device.gpu, - placement_backend_name(expert_kind), expert_gpu, - peer_ok ? "enabled" : "unavailable"); - } else if (expert_kind != local_kind) { - std::fprintf(stderr, - "[deepseek4-moe-tp] cross-vendor owner join %s:%d <-> %s:%d " - "uses in-process host staging\n", - placement_backend_name(local_kind), cfg_.device.gpu, - placement_backend_name(expert_kind), expert_gpu); + "[deepseek4-moe-tp] peer access GPU %d <-> GPU %d: %s\n", + cfg_.device.gpu, expert_gpu, peer_ok ? "enabled" : "unavailable"); } - expert_backend_ = init_placement_backend(expert_kind, expert_gpu, &err); + expert_backend_ = ggml_backend_cuda_init(expert_gpu); if (!expert_backend_) { std::fprintf(stderr, - "[deepseek4-moe-tp] failed to initialize in-process " - "secondary backend %s:%d: %s\n", - placement_backend_name(expert_kind), expert_gpu, - err.c_str()); + "[deepseek4-moe-tp] failed to initialize in-process expert GPU %d\n", + expert_gpu); return false; } hybrid_cfg.materialize_cold_experts = true; @@ -1328,6 +1076,30 @@ bool DeepSeek4Backend::init_hybrid_model() { return false; } + const int attention_peer_groups = attention_tp_peer_groups(); + if (attention_peer_groups < 0) { + std::fprintf(stderr, + "[deepseek4-attention-tp] invalid " + "DFLASH_DS4_ATTENTION_TP_GROUPS\n"); + return false; + } + if (attention_peer_groups > 0) { + if (!inprocess_tp || !expert_backend_ || + !env_flag_enabled("DFLASH_DS4_FUSED_VERIFY")) { + std::fprintf(stderr, + "[deepseek4-attention-tp] requires in-process " + "heterogeneous fused verification\n"); + return false; + } + if (!init_deepseek4_attention_tp( + w_, expert_backend_, attention_peer_groups, &err)) { + std::fprintf(stderr, + "[deepseek4-attention-tp] initialization failed: %s\n", + err.c_str()); + return false; + } + } + // The physical placement is shared by both phases. Decode may own only a // subset so its fast main branch does not outrun and then wait on the peer; // prefill continues to consume every resident expert. @@ -1410,21 +1182,22 @@ bool DeepSeek4Backend::park(ParkTarget target) { maybe_save_routing_stats(); for (int i = 0; i < PREFIX_SLOTS; ++i) { - snapshot_free(i); + free_deepseek4_snapshot(snapshots_[i]); } last_logits_.clear(); - last_logits_pos_ = -1; free_deepseek4_cache(cache_); expert_runtime_.reset(); stream_engine_.destroy(); moe_hybrid_.reset(); + // Attention-TP weights are allocated by expert_backend_. Release them + // before destroying that borrowed backend. + free_deepseek4_weights(w_); if (expert_backend_) { ggml_backend_free(expert_backend_); expert_backend_ = nullptr; } moe_placement_ = {}; moe_decode_placement_ = {}; - free_deepseek4_weights(w_); parked_ = true; if (spec_drafter_) { std::printf("[deepseek4] target parked (target VRAM released; " @@ -1534,9 +1307,7 @@ int deepseek4_hybrid_prefill_chunk_tokens( int DeepSeek4Backend::do_prefill(const std::vector & tokens, const DaemonIO & io, - int kv_offset, - int snap_slot, - int snap_pos) { + int kv_offset) { // The all-hot layer-range path supports causal chunked prefill. The // optimized graph snapshots the previous raw SWA window, attends over // that snapshot plus the current ubatch, and commits only the final SWA @@ -1584,9 +1355,6 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, base_chunk, chunk, kv_offset + n_total); } int pos = kv_offset; - const bool save_snapshot = - snap_slot >= 0 && snap_slot < PREFIX_SLOTS && - snap_pos > kv_offset && snap_pos <= kv_offset + n_total; // New sequence: clear the cache buffer so compressor state double-buffers // and compressed-KV rows start from zeros, exactly like a fresh server. // Without this, the first flush windows of a request pool over the @@ -1596,34 +1364,19 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, reset_deepseek4_cache(cache_); } last_logits_.clear(); - last_logits_pos_ = -1; - int spec_final_from = n_total; - int spec_snap_from = n_total; - int spec_snap_to = 0; - int spec_old_rows_for_final = 0; + int spec_capture_from = n_total; if (spec_enabled_ && spec_drafter_) { const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; - const int snap_tokens = save_snapshot ? snap_pos - kv_offset : n_total; - spec_final_from = std::max(0, n_total - w_.n_swa); - if (save_snapshot) { - spec_snap_from = std::max(0, snap_tokens - w_.n_swa); - spec_snap_to = snap_tokens; - } - if (kv_offset == 0 || feat_row <= 0 || + if (kv_offset == 0 || n_total >= w_.n_swa || feat_row <= 0 || spec_feat_window_.size() % (size_t) feat_row != 0) { spec_feat_window_.clear(); + spec_capture_from = std::max(0, n_total - w_.n_swa); } else { - // Preserve enough restored rows for both the requested checkpoint - // and the final prompt tail. The live vector is trimmed after - // prefill; the snapshot copy is independently trimmed at save. + // Keep enough prior rows for the new prompt suffix, then append all + // new rows. This bounds host capture storage at n_swa without + // shifting a multi-megabyte feature window after every token. const size_t old_rows = spec_feat_window_.size() / (size_t) feat_row; - spec_old_rows_for_final = std::max(0, w_.n_swa - n_total); - const int old_rows_for_snap = save_snapshot - ? std::max(0, w_.n_swa - snap_tokens) : 0; - const size_t keep_rows = std::min( - old_rows, - (size_t) std::max(spec_old_rows_for_final, - old_rows_for_snap)); + const size_t keep_rows = (size_t) std::max(0, w_.n_swa - n_total); if (old_rows > keep_rows) { const size_t drop_floats = (old_rows - keep_rows) * (size_t) feat_row; const size_t keep_floats = keep_rows * (size_t) feat_row; @@ -1632,6 +1385,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, keep_floats * sizeof(float)); spec_feat_window_.resize(keep_floats); } + spec_capture_from = 0; } } const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); @@ -1639,28 +1393,10 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, DeepSeek4StepTelemetry tel_acc; int steps = 0; - bool snapshot_saved = false; for (int i = 0; i < n_total;) { if (io.cancelled) return pos; - int n_tok = std::min(chunk, n_total - i); - // A snapshot must represent an exact token boundary. Split a batched - // prefill chunk when the requested boundary falls inside it. - if (save_snapshot && !snapshot_saved && - snap_pos > pos && snap_pos < pos + n_tok) { - n_tok = snap_pos - pos; - } - if (spec_enabled_ && spec_drafter_) { - const bool batch_final_capture = - !w_.moe_hybrid && - cache_.prefill_mode != PrefillAttentionMode::Exact && - n_tok > 4 && - n_tok <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS; - n_tok = capture_safe_prefill_tokens( - i, n_tok, spec_final_from, batch_final_capture, - save_snapshot && !snapshot_saved, - spec_snap_from, spec_snap_to); - } + const int n_tok = std::min(chunk, n_total - i); // Embed tokens std::vector embed(w_.n_embd * n_tok); @@ -1675,29 +1411,9 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, Ds4VerifyHooks spec_hooks; std::vector spec_cap; Ds4VerifyHooks * hp = nullptr; - const bool capture_final = i + n_tok > spec_final_from; - const bool capture_snapshot = - !snapshot_saved && i < spec_snap_to && - i + n_tok > spec_snap_from; - if (spec_enabled_ && spec_drafter_ && - (capture_final || capture_snapshot)) { + if (spec_enabled_ && spec_drafter_ && i + n_tok > spec_capture_from) { spec_hooks.capture_layer_ids = &spec_drafter_->capture_layer_ids; spec_hooks.capture_out = &spec_cap; - int capture_begin = n_tok; - int capture_end = 0; - if (capture_final) { - capture_begin = std::min( - capture_begin, std::max(0, spec_final_from - i)); - capture_end = n_tok; - } - if (capture_snapshot) { - capture_begin = std::min( - capture_begin, std::max(0, spec_snap_from - i)); - capture_end = std::max( - capture_end, std::min(n_tok, spec_snap_to - i)); - } - spec_hooks.capture_token_begin = capture_begin; - spec_hooks.capture_token_end = capture_end; hp = &spec_hooks; } if (moe_hybrid_ && (expert_runtime_.compute || expert_backend_)) { @@ -1729,13 +1445,8 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, } if (ok && hp && !spec_cap.empty()) { const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; - for (int t = 0; t < n_tok; ++t) { - const int token_index = i + t; - const bool keep_for_final = token_index >= spec_final_from; - const bool keep_for_snapshot = - !snapshot_saved && token_index >= spec_snap_from && - token_index < spec_snap_to; - if (!keep_for_final && !keep_for_snapshot) continue; + const int first_capture = std::max(0, spec_capture_from - i); + for (int t = first_capture; t < n_tok; ++t) { spec_feat_window_.insert(spec_feat_window_.end(), spec_cap.begin() + (size_t) t * feat_row, spec_cap.begin() + (size_t) (t + 1) * feat_row); @@ -1751,31 +1462,8 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, } last_logits_ = std::move(logits); pos += n_tok; - last_logits_pos_ = cache_.cur_pos; i += n_tok; - if (save_snapshot && !snapshot_saved && pos == snap_pos) { - snapshot_saved = snapshot_save(snap_slot); - if (!snapshot_saved) { - std::fprintf(stderr, - "[deepseek4] failed to save snapshot slot=%d pos=%d\n", - snap_slot, snap_pos); - } else if (spec_enabled_ && spec_drafter_) { - // Discard checkpoint-only rows once their snapshot is saved. - // Retain just the already-captured prefix of the final SWA - // window, so distant checkpoints do not bridge a huge gap in - // host feature memory. - const int processed = i; - const int final_new_rows = - std::max(0, processed - spec_final_from); - keep_spec_feature_tail( - spec_feat_window_, - (size_t) spec_old_rows_for_final + - (size_t) final_new_rows); - } - } } - keep_spec_feature_tail(spec_feat_window_, - (size_t) std::max(0, w_.n_swa)); if (timing) { log_step_tel("prefill", n_total, steps, elapsed_s(phase_t0), tel_acc); } @@ -1889,13 +1577,6 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, if (process_logits) { history.push_back(next_token); } - if (generated > 0) { - // The forward above advanced cache_ through the previously - // emitted token. Retain its logits so a later manual/continued - // snapshot can resume at exactly cache_.cur_pos. - last_logits_ = std::move(logits); - last_logits_pos_ = cache_.cur_pos; - } out_tokens.push_back(next_token); const auto emit_t0 = Clock::now(); io.emit(next_token); @@ -1913,11 +1594,6 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, const DaemonIO & io) { - return generate_from_state(req, io, 0); -} - -GenerateResult DeepSeek4Backend::generate_from_state( - const GenerateRequest & req, const DaemonIO & io, int kv_offset) { GenerateResult result; DaemonIO out_io = io.with_token_callback(req.on_token); auto t0 = Clock::now(); @@ -1926,25 +1602,8 @@ GenerateResult DeepSeek4Backend::generate_from_state( sampler_rng_.seed(sampler_.seed); } - if (kv_offset < 0 || kv_offset > (int) req.prompt.size()) { - result.fail(GenerateErrorCode::PrefillFailed, - "restored prefix exceeds prompt length"); - return result; - } - - // Prefill only the suffix that is not already represented by a restored - // snapshot. An exact full-prompt hit can decode immediately from the - // logits and speculative feature window saved with the cache state. - int committed = kv_offset; - if (kv_offset == 0) { - committed = do_prefill(req.prompt, out_io, 0, - req.snap_slot, req.snap_pos); - } else if (kv_offset < (int) req.prompt.size()) { - std::vector suffix(req.prompt.begin() + kv_offset, - req.prompt.end()); - committed = do_prefill(suffix, out_io, kv_offset, - req.snap_slot, req.snap_pos); - } + // Prefill + int committed = do_prefill(req.prompt, out_io); if (committed < 0) { result.fail(GenerateErrorCode::PrefillFailed); return result; @@ -1988,10 +1647,6 @@ GenerateResult DeepSeek4Backend::generate_from_state( const int win_len = feat_row > 0 ? (int) (spec_feat_window_.size() / feat_row) : 0; std::vector spec_toks; spec_ran = true; - // The DSpark API does not return the final target logits. Once it - // advances the target cache, reject post-decode snapshots rather - // than pairing that state with stale prefill logits. - last_logits_pos_ = -1; if (!run_deepseek4_dspark_spec_decode( backend_, cfg_.device.gpu, w_, cache_, *spec_drafter_, committed, seed, req.n_gen - 1, @@ -2044,55 +1699,19 @@ GenerateResult DeepSeek4Backend::generate_from_state( // ── Snapshots ─────────────────────────────────────────────────────────── bool DeepSeek4Backend::snapshot_save(int slot) { - if (slot < 0 || slot >= PREFIX_SLOTS || !snap_backend_ || - cache_.cur_pos <= 0 || last_logits_pos_ != cache_.cur_pos || - w_.n_vocab <= 0 || - last_logits_.size() != (size_t) w_.n_vocab) { - return false; - } - - snapshot_free(slot); - if (!deepseek4_snapshot_save(cache_, snap_backend_, snapshots_[slot])) { - return false; - } - - try { - auto & aux = snapshot_aux_[slot]; - aux.last_logits = last_logits_; - aux.spec_feat_window = spec_feat_window_; - keep_spec_feature_tail(aux.spec_feat_window, - (size_t) std::max(0, w_.n_swa)); - aux.used = true; - } catch (const std::bad_alloc &) { - snapshot_free(slot); - return false; - } - - const size_t core_bytes = snapshots_[slot].buf - ? ggml_backend_buffer_get_size(snapshots_[slot].buf) : 0; - const size_t aux_bytes = - (snapshot_aux_[slot].last_logits.size() + - snapshot_aux_[slot].spec_feat_window.size()) * sizeof(float); - std::fprintf(stderr, - "[deepseek4] snapshot saved slot=%d pos=%d size=%.1f MiB\n", - slot, snapshots_[slot].cur_pos, - (double) (core_bytes + aux_bytes) / (1024.0 * 1024.0)); - return true; + if (slot < 0 || slot >= PREFIX_SLOTS) return false; + // TODO: Implement snapshot save (copy KV cache + HC state to CPU) + return false; } void DeepSeek4Backend::snapshot_free(int slot) { if (slot < 0 || slot >= PREFIX_SLOTS) return; free_deepseek4_snapshot(snapshots_[slot]); - snapshot_aux_[slot] = SnapshotAux{}; } bool DeepSeek4Backend::snapshot_used(int slot) const { if (slot < 0 || slot >= PREFIX_SLOTS) return false; - const auto & snap = snapshots_[slot]; - const auto & aux = snapshot_aux_[slot]; - return snap.ctx != nullptr && snap.buf != nullptr && snap.cur_pos > 0 && - aux.used && w_.n_vocab > 0 && - aux.last_logits.size() == (size_t) w_.n_vocab; + return snapshots_[slot].ctx != nullptr; } int DeepSeek4Backend::snapshot_cur_pos(int slot) const { @@ -2100,48 +1719,11 @@ int DeepSeek4Backend::snapshot_cur_pos(int slot) const { return snapshots_[slot].cur_pos; } -bool DeepSeek4Backend::snapshot_restore(int slot) { - if (!snapshot_used(slot)) return false; - - std::vector restored_logits; - std::vector restored_features; - try { - restored_logits = snapshot_aux_[slot].last_logits; - restored_features = snapshot_aux_[slot].spec_feat_window; - } catch (const std::bad_alloc &) { - return false; - } - - if (!deepseek4_snapshot_restore(snapshots_[slot], cache_)) { - return false; - } - last_logits_ = std::move(restored_logits); - spec_feat_window_ = std::move(restored_features); - last_logits_pos_ = cache_.cur_pos; - return true; -} - GenerateResult DeepSeek4Backend::restore_and_generate_impl( int slot, const GenerateRequest & req, const DaemonIO & io) { - GenerateResult result; - if (!snapshot_used(slot)) { - result.fail(GenerateErrorCode::InvalidSnapshotSlot); - return result; - } - - const int snap_pos = snapshot_cur_pos(slot); - if (snap_pos > (int) req.prompt.size()) { - std::fprintf(stderr, - "[pc] DeepSeek snapshot longer than prompt " - "(snap=%d > prompt=%zu) -- fresh prefill fallback\n", - snap_pos, req.prompt.size()); - return generate_impl(req, io); - } - if (!snapshot_restore(slot)) { - result.fail(GenerateErrorCode::BackendSpecific, "snapshot restore"); - return result; - } - return generate_from_state(req, io, snap_pos); + // TODO: Implement snapshot restore + generate + (void)slot; + return generate_impl(req, io); } bool DeepSeek4Backend::handle_compress(const std::string & line, @@ -2170,12 +1752,13 @@ void DeepSeek4Backend::shutdown() { maybe_save_routing_stats(); free_drafter(); for (int i = 0; i < PREFIX_SLOTS; i++) { - snapshot_free(i); + free_deepseek4_snapshot(snapshots_[i]); } free_deepseek4_cache(cache_); expert_runtime_.reset(); stream_engine_.destroy(); moe_hybrid_.reset(); + free_deepseek4_weights(w_); if (expert_backend_) { ggml_backend_free(expert_backend_); expert_backend_ = nullptr; @@ -2184,7 +1767,6 @@ void DeepSeek4Backend::shutdown() { routing_stats_out_path_.clear(); moe_placement_ = {}; moe_decode_placement_ = {}; - free_deepseek4_weights(w_); if (snap_backend_) { ggml_backend_free(snap_backend_); snap_backend_ = nullptr; } if (backend_) { ggml_backend_free(backend_); backend_ = nullptr; } } diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index 628a464ee..5d7591fa3 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -144,6 +144,13 @@ bool deepseek4_dspark_draft_read_async_output( ggml_backend_t backend, std::vector & out_hidden, std::vector * confidence_hidden = nullptr); +// Borrow the cached graph's device-resident outputs after async submission. +// The draft runtime retains ownership; callers must consume them before the +// graph cache is rebuilt, reset, or submitted again. +bool deepseek4_dspark_draft_device_outputs( + ggml_backend_t backend, + ggml_tensor ** out_hidden, + ggml_tensor ** confidence_hidden = nullptr); void deepseek4_dspark_draft_wait(ggml_backend_t backend); // Batched target verify forward WITH feature capture (defined in diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index d4e409fda..cf88c3ca9 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -67,18 +67,15 @@ class DeepSeek4DFlashTarget : public DFlashTarget { n, n > 0 ? tokens[0] : -1, n > 1 ? tokens[1] : -1, w_.n_vocab); return false; } - // Sequential verify: q single-token forwards through the same cached - // graph as ordinary AR decode. This preserves target arithmetic; exact - // rollback still requires a full snapshot and replay after rejection. - // DFLASH_DS4_SEQ_VERIFY is a diagnostic. The supported reference mode, - // DFLASH_DS4_SPEC_REFERENCE_EXACT, enables both requirements together. + // Sequential verify (measurement mode): q single-token forwards through + // the legacy AR decode path. Causal by construction, compressor fed every + // token. Slow; used to measure the drafter's token-at-a-time accept rate. + // It is not a bit-exact oracle: graph shape can change floating-point + // reduction order around near-tied logits. Enable: DFLASH_DS4_SEQ_VERIFY=1 + // (pair with DFLASH_DS4_FULL_SNAP=1 so rollback/replay stay exact). static const bool seq_verify = [] { - const char * exact = - std::getenv("DFLASH_DS4_SPEC_REFERENCE_EXACT"); - const char * sequential = - std::getenv("DFLASH_DS4_SEQ_VERIFY"); - return (exact && *exact && *exact != '0') || - (sequential && *sequential && *sequential != '0'); + const char * v = std::getenv("DFLASH_DS4_SEQ_VERIFY"); + return v && *v && *v != '0'; }(); if (seq_verify) { std::vector am_all; @@ -94,7 +91,7 @@ class DeepSeek4DFlashTarget : public DFlashTarget { tokens.data() + t, 1, base_pos + t, am1, keep_logits_ ? &logits1 : nullptr, feat1, telemetry_, - /*allow_graph_reuse=*/true, + /*allow_graph_reuse=*/false, moe_hybrid_, expert_runtime_, routing_stats_)) { return false; @@ -112,13 +109,13 @@ class DeepSeek4DFlashTarget : public DFlashTarget { return true; } std::vector am; - // Reuse the normal cached graph for q==1 so reference verification has - // exactly the same target arithmetic as ordinary AR decode. + // n==1 must take the dynamic (non-reuse) path: the reused decode graph + // skips the capture/all-logits hooks (backend HC), which this needs. if (!deepseek4_dspark_verify_forward(backend_, device_, w_, cache_, capture_ids_, embed_buf_.data(), tokens.data(), n, base_pos, am, keep_logits_ ? &verify_logits_ : nullptr, verify_features_, telemetry_, - /*allow_graph_reuse=*/true, + /*allow_graph_reuse=*/n > 1, moe_hybrid_, expert_runtime_, routing_stats_)) { return false; @@ -509,6 +506,21 @@ void spec_rollback_apply(const DeepSeek4SpecRollback & rb, const DeepSeek4Weight lc.raw_kv, src, (size_t) row * lc.raw_kv->nb[1], s.raw_row_bytes); } + if (il < cache.attention_tp_layers.size()) { + ggml_tensor * peer_raw = + cache.attention_tp_layers[il].raw_kv; + if (peer_raw && peer_raw->ne[1] == lc.raw_kv->ne[1] && + ggml_row_size(peer_raw->type, peer_raw->ne[0]) == + s.raw_row_bytes) { + // The saved main row is byte-identical to the peer + // mirror. Rollback is rare, so a direct synchronous + // write keeps the two persistent rings coherent. + ggml_backend_tensor_set( + peer_raw, src, + (size_t) row * peer_raw->nb[1], + s.raw_row_bytes); + } + } } } } @@ -528,6 +540,9 @@ void spec_rollback_apply(const DeepSeek4SpecRollback & rb, const DeepSeek4Weight } } } + if (cache.attention_tp_buf) { + cache.attention_tp_cur_pos = commit_pos; + } } using SpecClock = std::chrono::steady_clock; @@ -612,13 +627,6 @@ bool deepseek4_dspark_verify_forward(ggml_backend_t backend, argmax_out = std::move(gpu_argmax); return true; } - if (n_tokens == 1 && (int) all_logits.size() < w.n_vocab && - (int) last_logits.size() >= w.n_vocab) { - // The reference-exact q1 path reuses the normal AR graph. That graph - // returns its logits through the regular output vector rather than - // the verifier's q-wide hook. - all_logits = last_logits; - } if ((int) all_logits.size() < w.n_vocab * n_tokens) { std::fprintf(stderr, "[ds4-verify] all_logits too small: got=%zu need=%d (cap=%zu)\n", all_logits.size(), w.n_vocab * n_tokens, capture_out.size()); @@ -660,12 +668,8 @@ bool run_deepseek4_dspark_spec_decode( const bool debug = spec_env_flag("DFLASH_DS4_DSPARK_DEBUG"); const bool timing = spec_env_flag("DFLASH_DS4_TIMING"); - const bool reference_exact = - spec_env_flag("DFLASH_DS4_SPEC_REFERENCE_EXACT"); - const bool full_snap = reference_exact || - spec_env_flag("DFLASH_DS4_FULL_SNAP"); - const bool seq_verify_mode = reference_exact || - spec_env_flag("DFLASH_DS4_SEQ_VERIFY"); + const bool full_snap = spec_env_flag("DFLASH_DS4_FULL_SNAP"); + const bool seq_verify_mode = spec_env_flag("DFLASH_DS4_SEQ_VERIFY"); const bool async_rollback = spec_env_flag("DFLASH_DS4_ASYNC_ROLLBACK"); const bool pinned_rollback = spec_env_flag("DFLASH_DS4_PINNED_ROLLBACK"); const bool q6_verify = @@ -677,13 +681,17 @@ bool run_deepseek4_dspark_spec_decode( spec_env_flag("DFLASH_DS4_DRAFT_OVERLAP_PROBE"); const bool draft_overlap_reuse_context = spec_env_flag("DFLASH_DS4_DRAFT_OVERLAP_REUSE_CONTEXT"); - if (reference_exact) { - std::fprintf(stderr, - "[ds4-spec] reference-exact verifier: sequential target replay " - "with full rollback snapshots\n"); - } ggml_backend_t drafter_backend = drafter.core.backend ? drafter.core.backend : backend; + const bool draft_device_chain_requested = + spec_env_flag("DFLASH_DS4_DRAFT_DEVICE_CHAIN"); + const bool draft_device_chain = draft_device_chain_requested && + drafter_backend == backend && !debug; + if (draft_device_chain_requested && !draft_device_chain) { + std::fprintf(stderr, + "[ds4-spec] device draft/head chain requires the target backend " + "and debug mode off; falling back to host staging\n"); + } const bool draft_overlap_probe_active = draft_overlap_probe && drafter_backend != backend; bool draft_overlap_probe_enabled = draft_overlap_probe_active; @@ -808,6 +816,9 @@ bool run_deepseek4_dspark_spec_decode( // Noise block = [seed] + [MASK]*(block-1). SpecClock::time_point t0 = SpecClock::now(); + bool draft_device_pending = false; + ggml_tensor * draft_hidden_device = nullptr; + ggml_tensor * confidence_hidden_device = nullptr; if (q_cap >= 2) { noise_ids[0] = lt; for (int i = 1; i < block; i++) { @@ -821,13 +832,30 @@ bool run_deepseek4_dspark_spec_decode( break; } - // Drafter forward -> block normed hidden states. - const bool draft_ok = deepseek4_dspark_draft_forward( - drafter_backend, - drafter, noise_embed.data(), - ctx_len > 0 ? feat_win.data() : nullptr, - ctx_len, pos, local_hidden, - use_confidence_width ? &confidence_hidden : nullptr); + // Keep draft outputs on the device when the Markov head shares + // this backend. Stream order makes the head wait for the draft + // graph without a host round trip or an intermediate synchronize. + bool draft_ok = false; + if (draft_device_chain) { + draft_ok = deepseek4_dspark_draft_forward_async( + drafter_backend, drafter, noise_embed.data(), + ctx_len > 0 ? feat_win.data() : nullptr, + ctx_len, pos); + if (draft_ok) { + draft_ok = deepseek4_dspark_draft_device_outputs( + drafter_backend, &draft_hidden_device, + use_confidence_width + ? &confidence_hidden_device : nullptr); + draft_device_pending = draft_ok; + } + } else { + draft_ok = deepseek4_dspark_draft_forward( + drafter_backend, + drafter, noise_embed.data(), + ctx_len > 0 ? feat_win.data() : nullptr, + ctx_len, pos, local_hidden, + use_confidence_width ? &confidence_hidden : nullptr); + } if (!draft_ok) { std::fprintf(stderr, "[ds4-spec] drafter forward failed\n"); ok = false; @@ -882,24 +910,64 @@ bool run_deepseek4_dspark_spec_decode( if (w_cap < q_step_cap) q_step_cap = w_cap; } if (q_step_cap >= 2) { - std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), - sizeof(float) * (size_t) n_embd * block); - if (use_confidence_width) { - std::memcpy(padded_confidence_hidden.data() + n_embd, - confidence_hidden.data(), - sizeof(float) * (size_t) n_embd * block); + if (draft_device_pending) { + ds_ok = dspark_markov_correct_greedy_chain_fused_device( + dw, backend, target.lm_head_tensor(), + draft_hidden_device, q_step_cap, lt, draft_tok, + use_confidence_width ? &draft_confidence : nullptr, + use_confidence_width + ? confidence_hidden_device : nullptr); + if (!ds_ok) { + deepseek4_dspark_draft_wait(drafter_backend); + const bool read_ok = + deepseek4_dspark_draft_read_async_output( + drafter_backend, local_hidden, + use_confidence_width ? &confidence_hidden : nullptr); + draft_device_pending = false; + if (!read_ok) { + std::fprintf(stderr, + "[ds4-spec] device draft output read failed\n"); + ok = false; + break; + } + } } - ds_ok = dspark_markov_correct_greedy_chain_fused( + if (!draft_device_pending) { + std::memcpy( + padded_hidden.data() + n_embd, local_hidden.data(), + sizeof(float) * (size_t) n_embd * block); + if (use_confidence_width) { + std::memcpy( + padded_confidence_hidden.data() + n_embd, + confidence_hidden.data(), + sizeof(float) * (size_t) n_embd * block); + } + ds_ok = dspark_markov_correct_greedy_chain_fused( dw, backend, target.lm_head_tensor(), padded_hidden.data(), q_step_cap, lt, draft_tok, use_confidence_width ? &draft_confidence : nullptr, use_confidence_width ? padded_confidence_hidden.data() : nullptr); - if (!ds_ok) { - ds_ok = dspark_markov_correct_greedy_chain(dw, backend, target, - padded_hidden.data(), q_step_cap, lt, 0.0f, draft_tok); + if (!ds_ok) { + ds_ok = dspark_markov_correct_greedy_chain( + dw, backend, target, padded_hidden.data(), + q_step_cap, lt, 0.0f, draft_tok); + } } if (!ds_ok || (int) draft_tok.size() < 2) { + if (draft_device_pending) { + deepseek4_dspark_draft_wait(drafter_backend); + if (!deepseek4_dspark_draft_read_async_output( + drafter_backend, local_hidden, + use_confidence_width + ? &confidence_hidden : nullptr)) { + std::fprintf(stderr, + "[ds4-spec] device draft output fallback failed\n"); + ok = false; + break; + } + draft_device_pending = false; + } // Fallback: plain projection of the block hiddens. std::vector pj; if (!target.project_hidden_to_tokens( @@ -1048,7 +1116,7 @@ bool run_deepseek4_dspark_spec_decode( // The bonus token is DEFERRED: it becomes the next step's seed, whose // KV is written then. t0 = SpecClock::now(); - if (full_snap && accept < q) { + if (full_snap) { // Legacy: full restore + replay the committed tokens through the // target so ring/compressor/n_comp advance exactly. std::vector kv_toks; @@ -1066,7 +1134,7 @@ bool run_deepseek4_dspark_spec_decode( ok = false; break; } - } else if (!full_snap && accept < q && q > 4) { + } else if (accept < q && q > 4) { // A rejected wide verify may have crossed two ratio-4 boundaries. // Restore the compact pre-verify state and replay only the // accepted prefix (at most q5), which is exact and rare at high @@ -1086,7 +1154,7 @@ bool run_deepseek4_dspark_spec_decode( ok = false; break; } - } else if (!full_snap && accept < q) { + } else if (accept < q) { // The prev-half flush is bad only if the boundary sits at-or-past // the commit point (its chunk then contains rejected tokens). const bool restore_prev = boundary_crossed && first_boundary >= commit_pos; diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index ec16a459b..3a7d00098 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -13,71 +13,6 @@ static void ds4_fv_set(ggml_tensor * t, const void * data, size_t nbytes) { if (t && t->buffer) ggml_backend_tensor_set(t, data, 0, nbytes); } - -struct Ds4MixedMoePolicy { - bool owner_local_reduction = false; - bool direct_device_join = false; - bool schedule_branches = false; - bool native_route_width = false; - bool fused_hc_join = false; - bool late_join_split = false; - bool deferred_join_split = false; - bool batch_peer_copies = false; - bool report_split_count = false; - bool scheduler_trace = false; - bool pin_route_weights = false; - bool preserve_routes_for_diagnostics = false; -}; - -static const Ds4MixedMoePolicy & ds4_mixed_moe_policy() { - static const Ds4MixedMoePolicy policy = [] { - Ds4MixedMoePolicy result; - result.owner_local_reduction = - ds4_env_flag("DFLASH_DS4_CROSS_VENDOR_OWNER_SUMS"); - result.direct_device_join = - ds4_env_flag("DFLASH_DS4_TP_PEER_FENCE") || - moe_hybrid_graph_policy().device_join; - result.schedule_branches = - ds4_env_flag("DFLASH_DS4_TP_SCHEDULE_BRANCHES"); - result.native_route_width = - ds4_env_flag("DFLASH_DS4_TP_NATIVE_ROUTE_WIDTH"); - result.fused_hc_join = - ds4_env_flag("DFLASH_DS4_TP_FUSED_HC_JOIN"); - result.late_join_split = - ds4_env_flag("DFLASH_DS4_TP_LATE_JOIN_SPLIT"); - result.deferred_join_split = - ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN_SPLIT"); - result.batch_peer_copies = - ds4_env_flag("GGML_BATCH_PEER_COPIES") || - ds4_env_flag("GGML_CUDA_BATCH_PEER_COPIES"); - result.report_split_count = - ds4_env_flag("DFLASH_DS4_TP_SPLIT_COUNT"); - result.scheduler_trace = - ds4_env_flag("DFLASH_DS4_TP_SCHED_TRACE"); - result.pin_route_weights = - ds4_env_flag("DFLASH_DS4_TP_MAIN_ROUTE_WEIGHTS"); - result.preserve_routes_for_diagnostics = - ds4_env_flag("DFLASH_DS4_TP_ROUTE_STATS") || - ds4_env_flag("DFLASH_DS4_ROUTING_STATS_OUT") || - ds4_env_flag("DFLASH_DS4_TP_CACHE_AUDIT"); - return result; - }(); - return policy; -} - -template -static void ds4_fv_set_repeated_rows( - ggml_tensor * tensor, - const std::vector & row) { - if (!tensor || row.empty()) return; - const size_t elements = (size_t) ggml_nelements(tensor); - GGML_ASSERT(elements % row.size() == 0); - std::vector repeated(elements); - for (size_t offset = 0; offset < elements; offset += row.size()) { - std::copy(row.begin(), row.end(), repeated.begin() + offset); - } - ds4_fv_set(tensor, repeated.data(), repeated.size() * sizeof(T)); -} static bool ds4_fused_verify_enabled() { static int enabled = -1; if (enabled < 0) { @@ -93,7 +28,9 @@ static bool ds4_fused_verify_trace_node( const bool routed = tensor && (tensor->op == GGML_OP_MUL_MAT_ID || tensor->op == GGML_OP_GET_ROWS); - if (!tensor || (!trace_all && !routed)) return false; + const bool inverse_rope = tensor && tensor->op == GGML_OP_ROPE_BACK && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_DEBUG_AB"); + if (!tensor || (!trace_all && !routed && !inverse_rope)) return false; const ggml_tensor * src0 = tensor->src[0]; const ggml_tensor * ids = routed ? (tensor->op == GGML_OP_MUL_MAT_ID ? tensor->src[2] : tensor->src[1]) @@ -117,6 +54,19 @@ static bool ds4_fused_verify_trace_node( : "unallocated", ids ? ids->data : nullptr, tensor->src[1] ? tensor->src[1]->data : nullptr); + if (ask && inverse_rope && tensor->src[1] && + tensor->src[1]->buffer && tensor->src[1]->type == GGML_TYPE_I32) { + const size_t count = std::min( + (size_t) ggml_nelements(tensor->src[1]), 8u); + std::array positions{}; + ggml_backend_tensor_get(tensor->src[1], positions.data(), 0, + count * sizeof(positions[0])); + std::fprintf(stderr, "[ds4-tp-trace] inverse_rope_pos="); + for (size_t i = 0; i < count; ++i) { + std::fprintf(stderr, "%s%d", i ? "," : "", positions[i]); + } + std::fprintf(stderr, "\n"); + } if (((ask && tensor->op == GGML_OP_MUL_MAT_ID) || (!ask && tensor->op == GGML_OP_GET_ROWS)) && ids && ids->buffer && @@ -196,10 +146,14 @@ static void ds4_fused_verify_refresh_hybrid_luts( : (hot < 0 ? 1.0f : 0.0f); } } - ds4_fv_set_repeated_rows(inputs.hot_local_lut, hot_lut); - ds4_fv_set_repeated_rows(inputs.hot_valid_lut, hot_valid); - ds4_fv_set_repeated_rows(inputs.cold_local_lut, cold_lut); - ds4_fv_set_repeated_rows(inputs.cold_valid_lut, cold_valid); + ds4_fv_set(inputs.hot_local_lut, hot_lut.data(), + sizeof(int32_t) * hot_lut.size()); + ds4_fv_set(inputs.hot_valid_lut, hot_valid.data(), + sizeof(float) * hot_valid.size()); + ds4_fv_set(inputs.cold_local_lut, cold_lut.data(), + sizeof(int32_t) * cold_lut.size()); + ds4_fv_set(inputs.cold_valid_lut, cold_valid.data(), + sizeof(float) * cold_valid.size()); } } @@ -382,13 +336,14 @@ static bool ds4_build_fused_verify_graph( fg.hash_ids.assign((size_t) w.n_layer, nullptr); if (hybrid) { fg.hybrid_inputs.assign((size_t) w.n_layer, {}); + if (w.attention_tp_backend && + w.attention_tp_peer_groups > 0 && + cache.attention_tp_layers.size() == (size_t) w.n_layer) { + fg.attention_tp_inputs.assign((size_t) w.n_layer, {}); + } } const int n_embd = w.n_embd; const int n_hc = w.n_hc; - const BackendPairCapabilities pair_capabilities = hybrid - ? backend_pair_capabilities(backend, hybrid->cold_backend) - : BackendPairCapabilities{true, true}; - const bool same_gpu_runtime = pair_capabilities.same_runtime; const size_t arena_size = 256u * 1024 * 1024; if (fg.sg.meta_arena.size() < arena_size) fg.sg.meta_arena.resize(arena_size); @@ -419,6 +374,20 @@ static bool ds4_build_fused_verify_graph( ggml_set_input(fg.i32_bundle); fg.i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 2 * (int64_t) w.n_layer); ggml_set_input(fg.i64_bundle); + if (!fg.attention_tp_inputs.empty()) { + fg.attention_tp_pos_q = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, q); + fg.attention_tp_neg_q = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, q); + fg.attention_tp_rawrows = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, q); + fg.attention_tp_i64_bundle = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, 2 * (int64_t) w.n_layer); + ggml_set_input(fg.attention_tp_pos_q); + ggml_set_input(fg.attention_tp_neg_q); + ggml_set_input(fg.attention_tp_rawrows); + ggml_set_input(fg.attention_tp_i64_bundle); + } // Mask bundle: q>1 preserves the q ring rows that the batch overwrites, // so its attention span is [n_swa + padded + q]. q=1 reads the current @@ -441,6 +410,11 @@ static bool ds4_build_fused_verify_graph( } fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mask_total); ggml_set_input(fg.mask_bundle); + if (!fg.attention_tp_inputs.empty()) { + fg.attention_tp_mask_bundle = ggml_new_tensor_1d( + ctx, GGML_TYPE_F32, mask_total); + ggml_set_input(fg.attention_tp_mask_bundle); + } int64_t mask_off = 0; // Keep every speculative lane in one HC tensor. The previous graph built @@ -455,6 +429,7 @@ static bool ds4_build_fused_verify_graph( ctx, hc_repeated, (int64_t) n_embd * n_hc, q); std::vector capture_layers(capture_ids.size(), nullptr); + DeepSeek4AttentionTpForkState attention_tp_fork_state{}; for (int il = 0; il < w.n_layer; ++il) { const DeepSeek4Layer & L = w.layers[(size_t) il]; @@ -524,8 +499,10 @@ static bool ds4_build_fused_verify_graph( } const int64_t n_attn = (int64_t) w.n_swa + padded + (lane_q > 1 ? lane_q : 0); + const int64_t layer_mask_off = mask_off; ain.attn_row_mask = ggml_view_2d(ctx, fg.mask_bundle, n_attn, lane_q, - n_attn * sizeof(float), (size_t) mask_off * sizeof(float)); + n_attn * sizeof(float), + (size_t) layer_mask_off * sizeof(float)); ain.padded_comp = padded; mask_off += n_attn * lane_q; @@ -533,9 +510,47 @@ static bool ds4_build_fused_verify_graph( std::vector i32ab; std::vector i64ab; ggml_tensor * normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); - ggml_tensor * attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, - lane_kv_start, lane_q, &ain, - i32b, i32ab, i64ab); + DeepSeek4AttentionTpGraphInputs * attention_tp_inputs = + fg.attention_tp_inputs.empty() + ? nullptr + : &fg.attention_tp_inputs[(size_t) il]; + const DeepSeek4AttentionTpCacheLayer * attention_tp_cache = + attention_tp_inputs + ? &cache.attention_tp_layers[(size_t) il] + : nullptr; + DeepSeek4AttentionGraphInputs attention_tp_peer_ain{}; + const DeepSeek4AttentionGraphInputs * attention_tp_peer_inputs = + nullptr; + if (attention_tp_inputs) { + attention_tp_peer_ain.rope_pos = fg.attention_tp_pos_q; + attention_tp_peer_ain.neg_pos = fg.attention_tp_neg_q; + attention_tp_peer_ain.raw_kv_rows = + fg.attention_tp_rawrows; + if (ratio > 0) { + int n_flush = 0; + for (int t = 0; t < lane_q; ++t) { + if (((lane_kv_start + t + 1) % ratio) == 0) ++n_flush; + } + const int n_comp_inputs = std::max(1, n_flush); + attention_tp_peer_ain.attn_comp_rows = ggml_view_1d( + ctx, fg.attention_tp_i64_bundle, n_comp_inputs, + (size_t) il * 2 * sizeof(int64_t)); + } + attention_tp_peer_ain.attn_row_mask = ggml_view_2d( + ctx, fg.attention_tp_mask_bundle, n_attn, lane_q, + n_attn * sizeof(float), + (size_t) layer_mask_off * sizeof(float)); + attention_tp_peer_ain.padded_comp = padded; + attention_tp_peer_inputs = &attention_tp_peer_ain; + } + ggml_tensor * attn_out = build_mla_attention( + ctx, gf, normed, w, L, lc, il, + lane_kv_start, lane_q, &ain, + i32b, i32ab, i64ab, nullptr, + DeepSeek4AttentionImpl::Explicit, + attention_tp_cache, attention_tp_inputs, + attention_tp_peer_inputs, + attention_tp_inputs ? &attention_tp_fork_state : nullptr); if (!attn_out) return false; if (!i32b.empty() || !i32ab.empty() || !i64ab.empty()) { std::fprintf(stderr, "[ds4-fused-verify] layer %d dynamic bindings; cannot fuse\n", il); @@ -543,10 +558,17 @@ static bool ds4_build_fused_verify_graph( } // ── Batched HC post (attention) + HC pre (FFN) ── - ggml_tensor * attn_batch = ggml_is_contiguous(attn_out) - ? attn_out : ggml_cont(ctx, attn_out); - hc_cur = ggml_ds4_hc_post( - ctx, hc_flat, attn_batch, split_attn, n_hc); + if (attention_tp_inputs && attention_tp_inputs->main_output && + attention_tp_inputs->peer_output) { + hc_cur = ggml_ds4_hc_post_split( + ctx, hc_flat, attention_tp_inputs->main_output, + attention_tp_inputs->peer_output, split_attn, n_hc); + } else { + ggml_tensor * attn_batch = ggml_is_contiguous(attn_out) + ? attn_out : ggml_cont(ctx, attn_out); + hc_cur = ggml_ds4_hc_post( + ctx, hc_flat, attn_batch, split_attn, n_hc); + } hc_flat = hc_cur; ggml_tensor * split_ffn = nullptr; ggml_tensor * ffn_working = ds4_build_fused_hc_pre( @@ -614,34 +636,9 @@ static bool ds4_build_fused_verify_graph( } if (hybrid) { - const Ds4MixedMoePolicy & mixed_policy = - ds4_mixed_moe_policy(); - const bool cross_vendor_owner_sums = - !same_gpu_runtime && - mixed_policy.owner_local_reduction; const bool allow_fused_combine = - same_gpu_runtime || cross_vendor_owner_sums; - const bool allow_direct_device_join = - pair_capabilities.native_gpu_handoff && - mixed_policy.direct_device_join; - const bool schedule_host_staged_branches = - !same_gpu_runtime && - mixed_policy.schedule_branches; - ggml_cgraph * hybrid_schedule_graph = - (allow_direct_device_join || schedule_host_staged_branches) - ? gf : nullptr; - // Cross-vendor staging normally preserves canonical route order - // so it matches the single-owner reduction as closely as - // possible. For performance qualification, allow each owner to - // reduce its routed experts locally and stage one [n_embd, q] - // partial instead of the full [n_embd, n_used, q] route tensor. - // This is the same join shape used by the qualified same-runtime - // heterogeneous path, but remains opt-in until exact-output and - // mixed-vendor burn-in have passed. - const MoeHybridJoinMode join_mode = - same_gpu_runtime || cross_vendor_owner_sums - ? MoeHybridJoinMode::OwnerPartialSums - : MoeHybridJoinMode::CanonicalRouteOrder; + ggml_backend_is_cuda(backend) && + ggml_backend_is_cuda(hybrid->cold_backend); MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); hybrid_cfg.n_expert_used = (int) selected->ne[0]; MoeLayerDesc desc = make_ds4_moe_layer_desc(L); @@ -662,8 +659,7 @@ static bool ds4_build_fused_verify_graph( // six routes in one owner graph: the legacy 4 + padded-2 split // computes two duplicate expert paths whose weights are zero. const bool native_route_width = - !same_gpu_runtime || - mixed_policy.native_route_width; + ds4_env_flag("DFLASH_DS4_TP_NATIVE_ROUTE_WIDTH"); if (lane_q > 1 && hybrid_cfg.n_expert_used > 4 && has_hot_routed && !native_route_width) { const int route_width = hybrid_cfg.n_expert_used; @@ -707,39 +703,41 @@ static bool ds4_build_fused_verify_graph( } if (!build_moe_hybrid_ffn_graph( ctx, - hybrid_schedule_graph, + (ds4_env_flag("DFLASH_DS4_TP_PEER_FENCE") || + ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN")) ? gf : nullptr, hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, first_ids, first_weights, lane_q, inputs, true, - allow_fused_combine, join_mode)) { + allow_fused_combine)) { return false; } ffn_out = inputs.output; if (!build_moe_hybrid_ffn_graph( ctx, - hybrid_schedule_graph, + (ds4_env_flag("DFLASH_DS4_TP_PEER_FENCE") || + ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN")) ? gf : nullptr, hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, padded_ids, padded_weights, lane_q, inputs, false, - allow_fused_combine, join_mode)) { + allow_fused_combine)) { return false; } ffn_out = ggml_add(ctx, ffn_out, inputs.output); } else { if (!build_moe_hybrid_ffn_graph( ctx, - hybrid_schedule_graph, + (ds4_env_flag("DFLASH_DS4_TP_PEER_FENCE") || + ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN")) ? gf : nullptr, hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, selected, router_weights, - lane_q, inputs, true, allow_fused_combine, - join_mode)) { + lane_q, inputs, true, allow_fused_combine)) { return false; } ffn_out = inputs.output; if (lane_q > 1 && native_route_width && - mixed_policy.fused_hc_join && + ds4_env_flag("DFLASH_DS4_TP_FUSED_HC_JOIN") && inputs.main_output && inputs.peer_output) { fused_hc_join_inputs = &inputs; } @@ -851,17 +849,6 @@ static bool ds4_build_fused_verify_graph( "[ds4-fused-verify] scheduler CPU fallback unavailable\n"); return false; } - const Ds4MixedMoePolicy & mixed_policy = - ds4_mixed_moe_policy(); - if (!same_gpu_runtime && mixed_policy.direct_device_join) { - static bool warned_cross_vendor_join = false; - if (!warned_cross_vendor_join) { - std::fprintf(stderr, - "[ds4-fused-verify] direct peer join disabled across GPU " - "vendors; using scheduler host staging\n"); - warned_cross_vendor_join = true; - } - } ggml_backend_t backends[3] = { backend, peer, hybrid->cpu_backend}; fg.sched = ggml_backend_sched_new( @@ -871,21 +858,30 @@ static bool ds4_build_fused_verify_graph( "[ds4-fused-verify] scheduler creation failed\n"); return false; } - const bool late_join_split = mixed_policy.late_join_split; - const MoeHybridGraphPolicy & moe_policy = - moe_hybrid_graph_policy(); + const bool late_join_split = + ds4_env_flag("DFLASH_DS4_TP_LATE_JOIN_SPLIT"); const bool targeted_join_split = - moe_policy.targeted_join_split; - const bool device_join_split = mixed_policy.deferred_join_split; - const bool batch_split_copies = mixed_policy.batch_peer_copies; - const bool report_split_count = mixed_policy.report_split_count; + ds4_env_flag("DFLASH_DS4_TP_TARGETED_JOIN_SPLIT"); + const bool device_join_split = + ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN_SPLIT") || + !fg.attention_tp_inputs.empty(); + const bool batch_split_copies = + ds4_env_flag("GGML_CUDA_BATCH_PEER_COPIES"); + const bool single_copy_event_fences = + ds4_env_flag("GGML_SCHED_SINGLE_COPY_EVENT_FENCES"); + const bool gpu_i32_repeat = + ds4_env_flag("LUCE_CUDA_I32_REPEAT"); + const bool report_split_count = + ds4_env_flag("DFLASH_DS4_TP_SPLIT_COUNT"); ggml_backend_sched_set_late_cross_input_split( fg.sched, late_join_split); ggml_backend_sched_set_deferred_peer_copy_split( fg.sched, device_join_split); ggml_backend_sched_set_batch_split_copies( fg.sched, batch_split_copies); - if (mixed_policy.scheduler_trace) { + ggml_backend_sched_set_single_copy_event_fences( + fg.sched, single_copy_event_fences); + if (ds4_env_flag("DFLASH_DS4_TP_SCHED_TRACE")) { ggml_backend_sched_set_eval_callback( fg.sched, ds4_fused_verify_trace_node, nullptr); } @@ -910,24 +906,45 @@ static bool ds4_build_fused_verify_graph( pin_main(fg.i32_bundle); pin_main(fg.i64_bundle); pin_main(fg.mask_bundle); + pin_peer(fg.attention_tp_pos_q); + pin_peer(fg.attention_tp_neg_q); + pin_peer(fg.attention_tp_rawrows); + pin_peer(fg.attention_tp_i64_bundle); + pin_peer(fg.attention_tp_mask_bundle); for (ggml_tensor * hids : fg.hash_ids) pin_main(hids); + for (const DeepSeek4AttentionTpGraphInputs & inputs : + fg.attention_tp_inputs) { + for (ggml_tensor * node : inputs.main_nodes) pin_main(node); + for (ggml_tensor * node : inputs.peer_nodes) pin_peer(node); + for (ggml_tensor * node : inputs.deferred_peer_copy_nodes) { + ggml_backend_sched_add_deferred_peer_copy_node( + fg.sched, node); + } + pin_main(inputs.main_output); + pin_peer(inputs.peer_partial); + pin_main(inputs.peer_output); + } for (const MoeHybridGraphInputs & inputs : fg.hybrid_inputs) { - if (mixed_policy.pin_route_weights) { + if (ds4_env_flag("DFLASH_DS4_TP_MAIN_ROUTE_WEIGHTS")) { for (ggml_tensor * node : inputs.router_nodes) { pin_main(node); } pin_main(inputs.router_weights); } - if (moe_policy.route_prefork) { + if (ds4_env_flag("DFLASH_DS4_TP_ROUTE_PREFORK")) { for (ggml_tensor * node : inputs.route_prefork_nodes) { pin_main(node); } } pin_main(inputs.hot_local_lut); pin_main(inputs.hot_valid_lut); - // The LUT is already q-batched and consumed only by the cold - // remap, so keep it with that owner for every verifier width. - pin_peer(inputs.cold_local_lut); + // q1 has no repeat. The opt-in exact I32 GPU repeat lets q>1 keep + // both the LUT and its repeated output on the cold owner as well. + if (q == 1 || gpu_i32_repeat) { + pin_peer(inputs.cold_local_lut); + } else { + pin_main(inputs.cold_local_lut); + } pin_peer(inputs.cold_valid_lut); for (ggml_tensor * node : inputs.hot_remap_nodes) { pin_main(node); @@ -961,7 +978,9 @@ static bool ds4_build_fused_verify_graph( pin_main(route.selected); pin_main(route.weights); } - if (mixed_policy.preserve_routes_for_diagnostics) { + if (ds4_env_flag("DFLASH_DS4_TP_ROUTE_STATS") || + ds4_env_flag("DFLASH_DS4_ROUTING_STATS_OUT") || + ds4_env_flag("DFLASH_DS4_TP_CACHE_AUDIT")) { // Preserve the unsplit authoritative route and weight matrices. // Reading the pre-fork list is insufficient in the legacy 4+2 // lowering because its second build overwrites the per-layer @@ -1160,10 +1179,13 @@ static int ds4_try_fused_verify_step( std::vector lv(q); for (int i = 0; i < q; ++i) iv[(size_t) i] = kv_start + i; ds4_fv_set(ex->pos_q, iv.data(), sizeof(int32_t) * q); + ds4_fv_set(fg->attention_tp_pos_q, iv.data(), sizeof(int32_t) * q); for (int i = 0; i < q; ++i) iv[(size_t) i] = -(kv_start + i); ds4_fv_set(ex->neg_q, iv.data(), sizeof(int32_t) * q); + ds4_fv_set(fg->attention_tp_neg_q, iv.data(), sizeof(int32_t) * q); for (int i = 0; i < q; ++i) lv[(size_t) i] = (kv_start + i) % w.n_swa; ds4_fv_set(ex->rawrows, lv.data(), sizeof(int64_t) * q); + ds4_fv_set(fg->attention_tp_rawrows, lv.data(), sizeof(int64_t) * q); for (int i = 0; i < q; ++i) iv[(size_t) i] = (kv_start + i) % 4; ds4_fv_set(ex->ape4, iv.data(), sizeof(int32_t) * q); for (int i = 0; i < q; ++i) lv[(size_t) i] = 4 + (kv_start + i) % 4; @@ -1199,6 +1221,8 @@ static int ds4_try_fused_verify_step( } ds4_fv_set(fg->i32_bundle, i32v.data(), sizeof(int32_t) * i32v.size()); ds4_fv_set(fg->i64_bundle, i64v.data(), sizeof(int64_t) * i64v.size()); + ds4_fv_set(fg->attention_tp_i64_bundle, + i64v.data(), sizeof(int64_t) * i64v.size()); // causal mask values { @@ -1254,6 +1278,8 @@ static int ds4_try_fused_verify_step( } GGML_ASSERT(off == maskv.size()); ds4_fv_set(fg->mask_bundle, maskv.data(), sizeof(float) * maskv.size()); + ds4_fv_set(fg->attention_tp_mask_bundle, + maskv.data(), sizeof(float) * maskv.size()); } if (token_ids) { @@ -1317,6 +1343,86 @@ static int ds4_try_fused_verify_step( std::fprintf(stderr, "[ds4-fused-verify] compute failed\n"); return -1; } + if (ds4_env_flag("DFLASH_DS4_ATTENTION_TP_DEBUG_AB")) { + static int reports = 0; + if (reports++ == 0 && !fg->attention_tp_inputs.empty()) { + if (ex->pos_q && ex->neg_q && fg->attention_tp_pos_q && + ex->pos_q->buffer && ex->neg_q->buffer && + fg->attention_tp_pos_q->buffer) { + std::vector main_pos((size_t) q); + std::vector main_neg((size_t) q); + std::vector peer_pos((size_t) q); + ggml_backend_tensor_get(ex->pos_q, main_pos.data(), 0, + main_pos.size() * sizeof(int32_t)); + ggml_backend_tensor_get(ex->neg_q, main_neg.data(), 0, + main_neg.size() * sizeof(int32_t)); + ggml_backend_tensor_get(fg->attention_tp_pos_q, + peer_pos.data(), 0, + peer_pos.size() * sizeof(int32_t)); + for (int i = 0; i < q; ++i) { + std::fprintf( + stderr, + "[deepseek4-attention-tp-ab] positions i=%d " + "main_pos=%d peer_pos=%d main_neg=%d\n", + i, main_pos[(size_t) i], peer_pos[(size_t) i], + main_neg[(size_t) i]); + } + } + for (size_t il = 0; il < fg->attention_tp_inputs.size(); ++il) { + for (const auto & pair : + fg->attention_tp_inputs[il].debug_pairs) { + if (!pair.main || !pair.peer || + pair.main->type != GGML_TYPE_F32 || + pair.peer->type != GGML_TYPE_F32 || + ggml_nelements(pair.main) != + ggml_nelements(pair.peer)) { + continue; + } + const size_t count = + (size_t) ggml_nelements(pair.main); + std::vector main_values(count); + std::vector peer_values(count); + ggml_backend_tensor_get( + pair.main, main_values.data(), 0, + count * sizeof(float)); + ggml_backend_tensor_get( + pair.peer, peer_values.data(), 0, + count * sizeof(float)); + double sum_abs = 0.0; + float max_abs = 0.0f; + size_t max_index = 0; + size_t non_finite = 0; + for (size_t i = 0; i < count; ++i) { + const float a = main_values[i]; + const float b = peer_values[i]; + if (!std::isfinite(a) || !std::isfinite(b)) { + ++non_finite; + continue; + } + const float delta = std::fabs(a - b); + sum_abs += delta; + if (delta > max_abs) { + max_abs = delta; + max_index = i; + } + } + std::fprintf( + stderr, + "[deepseek4-attention-tp-ab] layer=%zu stage=%s " + "n=%zu mean_abs=%.9g max_abs=%.9g max_i=%zu " + "main=%.9g peer=%.9g non_finite=%zu\n", + il, pair.name ? pair.name : "unknown", count, + count > non_finite + ? sum_abs / (double) (count - non_finite) + : INFINITY, + max_abs, max_index, + count ? main_values[max_index] : 0.0f, + count ? peer_values[max_index] : 0.0f, + non_finite); + } + } + } + } if (telemetry) { telemetry->full_graph_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); } diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 19ce98f6c..dab1df492 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -13,7 +13,6 @@ #include "internal.h" #include "../common/step_graph.h" #include "../common/cuda_graph_overrides.h" -#include "../common/dynamic_backend.h" #include "../common/moe_expert_compute.h" #include "../common/moe_hybrid_ffn_eval.h" #include "../common/moe_hybrid_routing_stats.h" @@ -589,24 +588,70 @@ static ggml_tensor * build_tail_rope_3d(ggml_context * ctx, float attn_factor, float beta_fast, float beta_slow, - int n_ctx_orig) { + int n_ctx_orig, + std::vector * backend_nodes = nullptr) { + auto track = [backend_nodes](ggml_tensor * tensor) { + if (backend_nodes && tensor) backend_nodes->push_back(tensor); + return tensor; + }; const int n_nope = head_dim - n_rot; // Split: nope [n_nope, n_heads, n_tokens], tail [n_rot, n_heads, n_tokens] - ggml_tensor * nope = ggml_view_3d(ctx, x, n_nope, n_heads, n_tokens, - x->nb[1], x->nb[2], 0); - ggml_tensor * tail = ggml_view_3d(ctx, x, n_rot, n_heads, n_tokens, - x->nb[1], x->nb[2], - (size_t)n_nope * ggml_type_size(x->type)); + ggml_tensor * nope = track(ggml_view_3d( + ctx, x, n_nope, n_heads, n_tokens, x->nb[1], x->nb[2], 0)); + ggml_tensor * tail = track(ggml_view_3d( + ctx, x, n_rot, n_heads, n_tokens, x->nb[1], x->nb[2], + (size_t)n_nope * ggml_type_size(x->type))); // tail is non-contiguous (stride between heads = head_dim, not n_rot) - tail = ggml_cont(ctx, tail); + tail = track(ggml_cont(ctx, tail)); // Apply rope to the contiguous tail: [n_rot, n_heads, n_tokens] // DS4 uses standard sequential pairs (i, i+1), which is GGML_ROPE_TYPE_NORMAL - tail = ggml_rope_ext(ctx, tail, pos, nullptr, - n_rot, GGML_ROPE_TYPE_NORMAL, n_ctx_orig, - freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow); + tail = track(ggml_rope_ext( + ctx, tail, pos, nullptr, n_rot, GGML_ROPE_TYPE_NORMAL, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, + beta_slow)); // Concat nope + tail along dim 0 → [head_dim, n_heads, n_tokens] - return ggml_concat(ctx, ggml_cont(ctx, nope), tail, 0); + nope = track(ggml_cont(ctx, nope)); + return track(ggml_concat(ctx, nope, tail, 0)); +} + +// Apply the inverse tail rotation using ROPE_BACK and positive positions. +// For a rotary transform this is the same operation as forward RoPE at the +// negated position. Keeping the position positive is useful on a scheduled +// heterogeneous branch: the peer already consumes that proven input for Q, +// so inverse rotation cannot depend on a second dynamic position tensor. +static ggml_tensor * build_tail_rope_back_3d( + ggml_context * ctx, + ggml_tensor * x, + ggml_tensor * pos, + int n_rot, + int head_dim, + int n_heads, + int n_tokens, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + int n_ctx_orig, + std::vector * backend_nodes = nullptr) { + auto track = [backend_nodes](ggml_tensor * tensor) { + if (backend_nodes && tensor) backend_nodes->push_back(tensor); + return tensor; + }; + const int n_nope = head_dim - n_rot; + ggml_tensor * nope = track(ggml_view_3d( + ctx, x, n_nope, n_heads, n_tokens, x->nb[1], x->nb[2], 0)); + ggml_tensor * tail = track(ggml_view_3d( + ctx, x, n_rot, n_heads, n_tokens, x->nb[1], x->nb[2], + (size_t) n_nope * ggml_type_size(x->type))); + tail = track(ggml_cont(ctx, tail)); + tail = track(ggml_rope_ext_back( + ctx, tail, pos, nullptr, n_rot, GGML_ROPE_TYPE_NORMAL, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, + beta_slow)); + nope = track(ggml_cont(ctx, nope)); + return track(ggml_concat(ctx, nope, tail, 0)); } // For KV (single head): x is [head_dim, n_tokens] @@ -993,7 +1038,9 @@ static void build_compressor_step( ggml_tensor * cur_all = nullptr, int n_tokens_all = 1, int kv_start_all = -1, - bool indexer_qat = false) { + bool indexer_qat = false, + ggml_tensor ** comp_values_out = nullptr) { + if (comp_values_out) *comp_values_out = nullptr; if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; @@ -1204,6 +1251,9 @@ static void build_compressor_step( if (indexer_qat) { pooled = ggml_ds4_indexer_qat(ctx, ggml_cont(ctx, pooled)); } + if (comp_values_out) { + *comp_values_out = pooled; + } ggml_tensor * pooled_f16 = ggml_cast(ctx, pooled, GGML_TYPE_F16); const int comp_row = token_pos / ratio; @@ -1331,6 +1381,10 @@ static void build_compressor_step( second_pooled = ggml_ds4_indexer_qat( ctx, ggml_cont(ctx, second_pooled)); } + if (comp_values_out) { + *comp_values_out = ggml_cont(ctx, ggml_concat( + ctx, *comp_values_out, second_pooled, 1)); + } ggml_tensor * second_comp_row = ggml_view_1d( ctx, comp_rows_inp, 1, comp_rows_inp->nb[0]); comp_cache_source = ggml_set_rows( @@ -1474,33 +1528,10 @@ static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int // rows in [n_comp, padded) are masked to -1e30 in the score matrix, which // underflows to exactly 0 in softmax, so a padded read is bit-identical to an // unpadded read of the first n_comp rows. -static int ds4_comp_pad_stride() { - static const int stride = [] { - constexpr int default_stride = 16; - const char * raw = std::getenv("DFLASH_DS4_COMP_PAD_STRIDE"); - if (!raw || !*raw) return default_stride; - const int requested = std::atoi(raw); - switch (requested) { - case 16: - case 32: - case 64: - case 128: - return requested; - default: - std::fprintf(stderr, - "[deepseek4] invalid DFLASH_DS4_COMP_PAD_STRIDE=%s; " - "using %d\n", - raw, default_stride); - return default_stride; - } - }(); - return stride; -} - +static constexpr int DS4_COMP_PAD_STRIDE = 16; static int ds4_padded_comp_rows(int n_comp, int cap) { if (n_comp <= 0) return 0; - const int stride = ds4_comp_pad_stride(); - const int padded = ((n_comp + stride - 1) / stride) * stride; + const int padded = ((n_comp + DS4_COMP_PAD_STRIDE - 1) / DS4_COMP_PAD_STRIDE) * DS4_COMP_PAD_STRIDE; return padded < cap ? padded : cap; } @@ -1603,6 +1634,45 @@ static ggml_tensor * build_indexer_topk( // ─── MLA Attention Block ──────────────────────────────────────────────── +struct DeepSeek4AttentionTpGraphInputs { + struct DebugPair { + const char * name = nullptr; + ggml_tensor * main = nullptr; + ggml_tensor * peer = nullptr; + }; + + std::vector main_nodes; + std::vector peer_nodes; + std::vector deferred_peer_copy_nodes; + std::vector debug_pairs; + ggml_tensor * main_output = nullptr; + ggml_tensor * peer_partial = nullptr; + ggml_tensor * peer_output = nullptr; +}; + +// Reusable R9700-owned storage for the two main-to-peer attention payloads. +// A fused verifier graph processes layers in order, so one QR buffer and one +// assembled-KV buffer are sufficient for every selected layer. Chaining each +// write through the previous state also makes the reuse dependency explicit to +// ggml's graph scheduler. +struct DeepSeek4AttentionTpForkState { + ggml_tensor * qr_storage = nullptr; + ggml_tensor * qr_state = nullptr; + int64_t qr_elements = 0; + + ggml_tensor * kv_storage = nullptr; + ggml_tensor * kv_state = nullptr; + int64_t kv_elements = 0; + + ggml_tensor * packed_storage = nullptr; + ggml_tensor * packed_state = nullptr; + int64_t packed_elements = 0; + + ggml_tensor * incremental_storage = nullptr; + ggml_tensor * incremental_state = nullptr; + int64_t incremental_elements = 0; +}; + static ggml_tensor * build_mla_attention( ggml_context * ctx, ggml_cgraph * gf, @@ -1618,25 +1688,149 @@ static ggml_tensor * build_mla_attention( std::vector & i32_array_inputs, std::vector & i64_array_inputs, std::vector * f32_array_inputs = nullptr, - DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit) { + DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit, + const DeepSeek4AttentionTpCacheLayer * attention_tp_cache = nullptr, + DeepSeek4AttentionTpGraphInputs * attention_tp = nullptr, + const DeepSeek4AttentionGraphInputs * attention_tp_peer_inputs = nullptr, + DeepSeek4AttentionTpForkState * attention_tp_fork_state = nullptr) { + + if (attention_tp) *attention_tp = {}; const int n_embd = w.n_embd; const int head_dim = w.head_dim; - const int n_head = w.n_head; + const int total_head = w.n_head; const int n_rot = w.n_rot; const int n_out_group = w.n_out_group; const int n_lora_o = w.n_lora_o; const int ratio = w.compress_ratios[layer_idx]; + // Attention cost differs sharply between compressed-layer families. A + // ratio filter lets the heterogeneous planner offload only layers whose + // longer score/value work can amortize a device fork. The default keeps + // the original all-layer behavior; a non-negative value selects an exact + // compression ratio (for example 4). + const char * attention_tp_ratio_raw = + std::getenv("DFLASH_DS4_ATTENTION_TP_RATIO"); + const int attention_tp_ratio = attention_tp_ratio_raw + ? std::atoi(attention_tp_ratio_raw) : -1; + const bool attention_tp_layer_selected = + attention_tp_ratio < 0 || ratio == attention_tp_ratio; + + const bool mirror_attention_cache = + attention_tp_layer_selected && + attention_tp && attention_tp_cache && cached_inputs && + attention_tp_peer_inputs && cached_inputs->raw_kv_rows && + attention_tp_peer_inputs->raw_kv_rows && + w.attention_tp_backend && w.attention_tp_peer_groups > 0 && + w.attention_tp_main_groups + w.attention_tp_peer_groups == + n_out_group && + layer_idx >= 0 && + (size_t) layer_idx < w.attention_tp_layers.size(); + const bool split_attention = + mirror_attention_cache && n_tokens > 1 && + attention_impl == DeepSeek4AttentionImpl::Explicit; + const bool split_attention_projection_only = split_attention && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_PROJECTION_ONLY"); + const bool split_attention_core_only = + split_attention && !split_attention_projection_only && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_CORE_ONLY"); + const bool split_attention_q = + split_attention && !split_attention_projection_only && + !split_attention_core_only; + // Cross-architecture F32 score/value GEMMs differ by a few ULPs. In this + // routed model those tiny changes can flip a later expert choice and + // destroy speculative acceptance. Keep the proven bit-identical Q-head + // projection split, then join before score/value attention. A future + // deterministic attention kernel can safely extend the owner boundary. + const bool split_attention_values = split_attention_core_only || + (split_attention_q && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_VALUES")); + const bool split_attention_direct_kv = split_attention_values && + (split_attention_core_only || + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_DIRECT_KV")); + const bool split_attention_flash = split_attention_values && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_FLASH"); + const bool split_attention_output_projection = + split_attention_values && !split_attention_core_only; + const bool split_attention_output_b = + (split_attention_output_projection || + split_attention_projection_only) && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_OUTPUT_B"); + const int heads_per_group = total_head / n_out_group; + const int main_groups = split_attention + ? w.attention_tp_main_groups : n_out_group; + const int peer_groups = split_attention + ? w.attention_tp_peer_groups : 0; + const int n_head = main_groups * heads_per_group; + const int peer_head = peer_groups * heads_per_group; + const DeepSeek4AttentionTpLayer * peer_weights = split_attention + ? &w.attention_tp_layers[(size_t) layer_idx] : nullptr; + + auto track_main = [attention_tp](ggml_tensor * tensor) { + if (attention_tp && tensor) attention_tp->main_nodes.push_back(tensor); + return tensor; + }; + auto track_peer = [attention_tp](ggml_tensor * tensor) { + if (attention_tp && tensor) attention_tp->peer_nodes.push_back(tensor); + return tensor; + }; + const bool debug_attention_tp = split_attention_q && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_DEBUG_AB"); + auto add_debug_pair = [&](const char * name, + ggml_tensor * main, + ggml_tensor * peer) { + if (!debug_attention_tp || !main || !peer) return; + GGML_ASSERT(main->type == GGML_TYPE_F32); + GGML_ASSERT(peer->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(main) == ggml_nelements(peer)); + // Snapshot both sides into owned contiguous outputs. Several useful + // checkpoints are views into graph inputs or softmax results; retaining + // those views directly lets gallocr reuse their storage and also makes + // a flat host read invalid when a stride gap is present. + main = track_main(ggml_cont(ctx, main)); + peer = track_peer(ggml_cont(ctx, peer)); + ggml_set_output(main); + ggml_set_output(peer); + ggml_build_forward_expand(gf, main); + ggml_build_forward_expand(gf, peer); + attention_tp->debug_pairs.push_back({name, main, peer}); + }; + auto add_debug_main_pair = [&](const char * name, + ggml_tensor * expected, + ggml_tensor * actual) { + if (!debug_attention_tp || !expected || !actual) return; + GGML_ASSERT(expected->type == GGML_TYPE_F32); + GGML_ASSERT(actual->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(expected) == ggml_nelements(actual)); + expected = track_main(ggml_cont(ctx, expected)); + actual = track_main(ggml_cont(ctx, actual)); + ggml_set_output(expected); + ggml_set_output(actual); + ggml_build_forward_expand(gf, expected); + ggml_build_forward_expand(gf, actual); + attention_tp->debug_pairs.push_back({name, expected, actual}); + }; + ggml_tensor * debug_peer_q_ref = nullptr; + ggml_tensor * debug_peer_context_ref = nullptr; + ggml_tensor * debug_full_q_ref = nullptr; + ggml_tensor * debug_full_context_ref = nullptr; + // ── Q path: cur → q_a → norm → q_b → per-head norm ───────────── // q_a: [n_embd, n_tokens] → [n_lora_q, n_tokens] ggml_tensor * qr = ggml_mul_mat(ctx, L.attn_q_a, cur); // qr_norm is reused by the ratio-4 indexer before the main q_b projection. qr = build_rms_norm(ctx, qr, L.attn_q_a_norm, w.rms_eps); - // q_b: [n_lora_q, n_tokens] → [n_head * head_dim, n_tokens] - ggml_tensor * q = ggml_mul_mat(ctx, L.attn_q_b, qr); + // Each TP owner projects only its persistent, contiguous head partition. + ggml_tensor * main_q_b = L.attn_q_b; + if (split_attention_q) { + main_q_b = ggml_view_2d( + ctx, L.attn_q_b, w.n_lora_q, (int64_t) n_head * head_dim, + L.attn_q_b->nb[1], 0); + } + ggml_tensor * q = ggml_mul_mat(ctx, main_q_b, qr); // Reshape to [head_dim, n_head, n_tokens] for per-head ops - q = ggml_reshape_3d(ctx, q, head_dim, n_head, n_tokens); + const int q_head = split_attention_q ? n_head : total_head; + q = ggml_reshape_3d(ctx, q, head_dim, q_head, n_tokens); // Reference DS4 applies unweighted RMSNorm independently to every Q head. q = ggml_rms_norm(ctx, q, w.rms_eps); @@ -1676,7 +1870,7 @@ static ggml_tensor * build_mla_attention( const bool fuse_q_rope = attention_impl != DeepSeek4AttentionImpl::Explicit && n_tokens > 1 && head_dim == 512 && n_rot == 64; if (!fuse_q_rope) { - q = build_tail_rope_3d(ctx, q, rope_pos, n_rot, head_dim, n_head, n_tokens, + q = build_tail_rope_3d(ctx, q, rope_pos, n_rot, head_dim, q_head, n_tokens, rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); } @@ -1684,6 +1878,24 @@ static ggml_tensor * build_mla_attention( rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + if (debug_attention_tp) { + debug_full_q_ref = track_main(ggml_mul_mat( + ctx, L.attn_q_b, qr)); + debug_full_q_ref = ggml_reshape_3d( + ctx, debug_full_q_ref, head_dim, total_head, n_tokens); + debug_full_q_ref = track_main(ggml_rms_norm( + ctx, debug_full_q_ref, w.rms_eps)); + debug_full_q_ref = track_main(build_tail_rope_3d( + ctx, debug_full_q_ref, rope_pos, n_rot, head_dim, + total_head, n_tokens, rope_freq, rope_scale, rope_ext, + rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + rope_n_ctx_orig)); + ggml_tensor * debug_main_q_ref = ggml_view_3d( + ctx, debug_full_q_ref, head_dim, n_head, n_tokens, + debug_full_q_ref->nb[1], debug_full_q_ref->nb[2], 0); + add_debug_main_pair("q_main_vs_full", debug_main_q_ref, q); + } + // ── Causal batched step (exact multi-token target semantics) ─── // The target model is causal: token i must not attend to batch tokens // j > i, must see the compressed-row count as of its own position, and — @@ -1789,6 +2001,7 @@ static ggml_tensor * build_mla_attention( ggml_tensor * cur_last = ggml_view_2d( ctx, cur, n_embd, 1, cur->nb[1], (size_t)(n_tokens - 1) * cur->nb[1]); ggml_tensor * comp_kv_source = lc.comp_kv; + ggml_tensor * comp_values = nullptr; if (ratio > 0 && L.attn_compressor_kv) { build_compressor_step(ctx, gf, cur_last, L.attn_compressor_ape, @@ -1817,7 +2030,9 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->flush_rows : nullptr, (causal_batch || fused_causal) ? cur : nullptr, n_tokens, - kv_start); + kv_start, + false, + &comp_values); } ggml_tensor * index_comp_kv_source = lc.index_comp_kv; @@ -1929,6 +2144,378 @@ static ggml_tensor * build_mla_attention( } // kv_attn: [head_dim, n_attn] + // Build one compact main-to-peer fork payload. Keeping Q-low-rank, current + // KV, and compressor emissions together prevents the peer branch from + // being split at each new dependency. The scheduler copies this single + // tensor before launching the peer split. Do not use the deferred peer-read + // op in this direction: direct gfx1151 reads from R9700 VRAM returned zeros + // in the model-backed A/B harness, while the normal main-to-peer copy is a + // supported HIP transfer. The reverse peer-to-main result path retains the + // exact deferred peer-read operation qualified for this pair. + ggml_tensor * peer_q = nullptr; + ggml_tensor * peer_kv_attn = nullptr; + const bool split_attention_worker = + split_attention_q || split_attention_core_only; + if (split_attention_worker) { + if ((!split_attention_direct_kv && !attention_tp_cache->raw_kv) || + (!split_attention_core_only && !peer_weights->attn_q_b) || + (split_attention_output_projection && + !peer_weights->attn_output_a) || + (split_attention_output_b && !peer_weights->attn_output_b) || + (split_attention_q && !attention_tp_peer_inputs->rope_pos) || + !attention_tp_peer_inputs->neg_pos || + !attention_tp_peer_inputs->attn_row_mask) { + return nullptr; + } + + auto flat_f32_main = [&](ggml_tensor * tensor) { + tensor = ds4_cast_if_needed(ctx, tensor, GGML_TYPE_F32); + if (!ggml_is_contiguous(tensor)) { + tensor = track_main(ggml_cont(ctx, tensor)); + } + return ggml_reshape_1d(ctx, tensor, ggml_nelements(tensor)); + }; + const bool dst_stream_staging = attention_tp_fork_state && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_DST_STREAM_STAGE"); + const bool packed_dst_stream_staging = dst_stream_staging && + ds4_env_flag("DFLASH_DS4_ATTENTION_TP_PACKED_STAGE"); + auto stage_for_peer = [&](ggml_tensor * source, + ggml_tensor *& storage, + ggml_tensor *& state, + int64_t & elements, + const char * name) { + source = flat_f32_main(source); + const int64_t source_elements = ggml_nelements(source); + if (!storage) { + storage = track_main(ggml_new_tensor_1d( + ctx, GGML_TYPE_F32, source_elements)); + ggml_set_name(storage, name); + // The base allocation must outlive every asynchronous + // destination-stream read made from its state views. + ggml_set_output(storage); + state = storage; + elements = source_elements; + } else { + // A staging state is scoped to one fused verifier shape and + // one selected compression family. Reusing it at another size + // would alias an in-flight payload and is a graph-build bug. + GGML_ASSERT(elements == source_elements); + } + + state = track_main(ggml_cpy(ctx, source, state)); + ggml_set_name(state, name); + ggml_set_output(state); + state->flags |= GGML_TENSOR_FLAG_DST_STREAM_COPY; + return state; + }; + + ggml_tensor * peer_qr = nullptr; + ggml_tensor * peer_kv = nullptr; + ggml_tensor * peer_comp_values = nullptr; + if (split_attention_core_only) { + // Project and normalize Q once at its native full width on the + // R9700. Only the independent tail heads and the already assembled + // latent KV matrix cross to Strix. Both become inputs of the same + // peer graph segment, avoiding the extra early-Q scheduler wave. + ggml_tensor * main_q_view = ggml_view_3d( + ctx, q, head_dim, n_head, n_tokens, + q->nb[1], q->nb[2], 0); + ggml_tensor * peer_q_view = ggml_view_3d( + ctx, q, head_dim, peer_head, n_tokens, + q->nb[1], q->nb[2], + (size_t) n_head * q->nb[1]); + q = track_main(ggml_cont(ctx, main_q_view)); + ggml_tensor * main_peer_q = track_main( + ggml_cont(ctx, peer_q_view)); + ggml_tensor * main_kv = ggml_reshape_2d( + ctx, flat_f32_main(kv_attn), head_dim, n_attn); + + // One packed producer gives the scheduler a single cross-device + // dependency. Separate Q and KV copies caused an extra main/peer + // wave per selected layer even though neither peer calculation can + // start before the assembled KV is ready in this core-only mode. + const int64_t peer_q_elements = ggml_nelements(main_peer_q); + ggml_tensor * main_fork = track_main(ggml_concat( + ctx, flat_f32_main(main_peer_q), + flat_f32_main(main_kv), 0)); + if (!ggml_is_contiguous(main_fork)) { + main_fork = track_main(ggml_cont(ctx, main_fork)); + } + ggml_set_output(main_fork); + ggml_build_forward_expand(gf, main_fork); + ggml_tensor * peer_fork = track_peer( + ggml_cont(ctx, main_fork)); + peer_q = ggml_view_3d( + ctx, peer_fork, head_dim, peer_head, n_tokens, + (size_t) head_dim * sizeof(float), + (size_t) head_dim * peer_head * sizeof(float), 0); + peer_kv = ggml_view_2d( + ctx, peer_fork, head_dim, n_attn, + (size_t) head_dim * sizeof(float), + (size_t) peer_q_elements * sizeof(float)); + } else if (split_attention_direct_kv) { + // Keep QR and the assembled KV block as two scheduler inputs so + // the peer Q projection can start before KV assembly completes. + ggml_tensor * main_qr = ggml_reshape_2d( + ctx, flat_f32_main(qr), w.n_lora_q, n_tokens); + ggml_tensor * main_kv = ggml_reshape_2d( + ctx, flat_f32_main(kv_attn), head_dim, n_attn); + + if (packed_dst_stream_staging) { + // Wait for assembled KV, then send QR+KV as one persistent + // packet. This trades a little early peer-Q overlap for one + // cross-device input and two fewer scheduler transitions per + // selected layer. + ggml_tensor * flat_qr = flat_f32_main(main_qr); + ggml_tensor * flat_kv = flat_f32_main(main_kv); + const int64_t qr_elements = ggml_nelements(flat_qr); + ggml_tensor * main_packet = track_main(ggml_concat( + ctx, flat_qr, flat_kv, 0)); + if (!ggml_is_contiguous(main_packet)) { + main_packet = track_main(ggml_cont(ctx, main_packet)); + } + + DeepSeek4AttentionTpForkState & fork = + *attention_tp_fork_state; + ggml_tensor * staged_packet = stage_for_peer( + main_packet, fork.packed_storage, fork.packed_state, + fork.packed_elements, "ds4_attn_packed_stage"); + staged_packet->flags |= GGML_TENSOR_FLAG_DST_STREAM_COPY; + ggml_tensor * peer_packet = track_peer( + ggml_cont(ctx, staged_packet)); + peer_qr = ggml_view_2d( + ctx, peer_packet, w.n_lora_q, n_tokens, + (size_t) w.n_lora_q * sizeof(float), 0); + peer_kv = ggml_view_2d( + ctx, peer_packet, head_dim, n_attn, + (size_t) head_dim * sizeof(float), + (size_t) qr_elements * sizeof(float)); + add_debug_pair("fork_qr", main_qr, peer_qr); + add_debug_pair("fork_kv", main_kv, peer_kv); + } else if (dst_stream_staging) { + DeepSeek4AttentionTpForkState & fork = + *attention_tp_fork_state; + ggml_tensor * staged_qr = stage_for_peer( + main_qr, fork.qr_storage, fork.qr_state, + fork.qr_elements, "ds4_attn_qr_stage"); + ggml_tensor * staged_kv = stage_for_peer( + main_kv, fork.kv_storage, fork.kv_state, + fork.kv_elements, "ds4_attn_kv_stage"); + main_qr = ggml_reshape_2d( + ctx, staged_qr, w.n_lora_q, n_tokens); + main_kv = ggml_reshape_2d( + ctx, staged_kv, head_dim, n_attn); + // Scheduler cross-input copies see the reshaped views, so mark + // those exact tensors in addition to their persistent states. + main_qr->flags |= GGML_TENSOR_FLAG_DST_STREAM_COPY; + main_kv->flags |= GGML_TENSOR_FLAG_DST_STREAM_COPY; + track_main(main_qr); + track_main(main_kv); + peer_qr = track_peer(ggml_cont(ctx, main_qr)); + peer_kv = track_peer(ggml_cont(ctx, main_kv)); + add_debug_pair("fork_qr", main_qr, peer_qr); + add_debug_pair("fork_kv", main_kv, peer_kv); + } else { + peer_qr = track_peer(ggml_cont(ctx, main_qr)); + peer_kv = track_peer(ggml_cont(ctx, main_kv)); + add_debug_pair("fork_qr", main_qr, peer_qr); + add_debug_pair("fork_kv", main_kv, peer_kv); + } + } else { + ggml_tensor * main_bundle = nullptr; + int64_t bundle_elements = 0; + int64_t qr_offset = -1; + int64_t kv_offset = -1; + int64_t comp_offset = -1; + auto append_bundle = [&](ggml_tensor * tensor, + int64_t * offset) { + if (!tensor) return; + tensor = flat_f32_main(tensor); + *offset = bundle_elements; + bundle_elements += ggml_nelements(tensor); + main_bundle = main_bundle + ? track_main(ggml_concat(ctx, main_bundle, tensor, 0)) + : tensor; + }; + append_bundle(qr, &qr_offset); + if (split_attention_values) { + append_bundle(kv, &kv_offset); + if (comp_values) append_bundle(comp_values, &comp_offset); + } + if (!main_bundle || qr_offset < 0 || + (split_attention_values && kv_offset < 0)) { + return nullptr; + } + if (!ggml_is_contiguous(main_bundle)) { + main_bundle = track_main(ggml_cont(ctx, main_bundle)); + } + if (dst_stream_staging) { + DeepSeek4AttentionTpForkState & fork = + *attention_tp_fork_state; + main_bundle = stage_for_peer( + main_bundle, + fork.incremental_storage, fork.incremental_state, + fork.incremental_elements, + "ds4_attn_incremental_stage"); + main_bundle->flags |= GGML_TENSOR_FLAG_DST_STREAM_COPY; + } else { + ggml_set_output(main_bundle); + ggml_build_forward_expand(gf, main_bundle); + } + + ggml_tensor * peer_bundle = track_peer( + ggml_cont(ctx, main_bundle)); + add_debug_pair("fork_bundle", main_bundle, peer_bundle); + + auto peer_view_2d = [&](int64_t offset, int64_t width, + int64_t rows) { + GGML_ASSERT(offset >= 0 && width > 0 && rows > 0); + return ggml_view_2d( + ctx, peer_bundle, width, rows, + (size_t) width * sizeof(float), + (size_t) offset * sizeof(float)); + }; + peer_qr = peer_view_2d( + qr_offset, w.n_lora_q, n_tokens); + peer_kv = split_attention_values + ? peer_view_2d(kv_offset, head_dim, n_tokens) + : nullptr; + peer_comp_values = + split_attention_values && comp_offset >= 0 + ? peer_view_2d( + comp_offset, head_dim, + ggml_nelements(comp_values) / head_dim) + : nullptr; + } + + if (split_attention_q) { + peer_q = track_peer(ggml_mul_mat( + ctx, peer_weights->attn_q_b, peer_qr)); + peer_q = ggml_reshape_3d( + ctx, peer_q, head_dim, peer_head, n_tokens); + peer_q = track_peer(ggml_rms_norm(ctx, peer_q, w.rms_eps)); + peer_q = track_peer(build_tail_rope_3d( + ctx, peer_q, attention_tp_peer_inputs->rope_pos, + n_rot, head_dim, peer_head, n_tokens, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + rope_n_ctx_orig, &attention_tp->peer_nodes)); + + if (debug_attention_tp) { + const int64_t main_head_dim = + (int64_t) n_head * head_dim; + ggml_tensor * main_peer_q_b = ggml_view_2d( + ctx, L.attn_q_b, w.n_lora_q, + (int64_t) peer_head * head_dim, + L.attn_q_b->nb[1], + (size_t) main_head_dim * L.attn_q_b->nb[1]); + debug_peer_q_ref = track_main(ggml_mul_mat( + ctx, main_peer_q_b, qr)); + debug_peer_q_ref = ggml_reshape_3d( + ctx, debug_peer_q_ref, head_dim, peer_head, n_tokens); + debug_peer_q_ref = track_main(ggml_rms_norm( + ctx, debug_peer_q_ref, w.rms_eps)); + debug_peer_q_ref = track_main(build_tail_rope_3d( + ctx, debug_peer_q_ref, rope_pos, + n_rot, head_dim, peer_head, n_tokens, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + rope_n_ctx_orig)); + add_debug_pair("q", debug_peer_q_ref, peer_q); + } + + if (!split_attention_values) { + // Q-only mode joins here because score/value attention still + // runs as one full-head main-device operation. Full attention + // TP must not create this dependency: its main 75% and peer + // 25% shards consume their own Q tensors independently and + // join only after value aggregation. + peer_q = track_peer(ggml_cont(ctx, peer_q)); + ggml_set_output(peer_q); + ggml_build_forward_expand(gf, peer_q); + ggml_tensor * peer_q_ready = + ggml_ds4_deferred_peer_copy(ctx, peer_q); + ggml_set_input(peer_q_ready); + ggml_set_output(peer_q_ready); + attention_tp->deferred_peer_copy_nodes.push_back(peer_q_ready); + track_main(peer_q_ready); + q = track_main(ggml_concat(ctx, q, peer_q_ready, 1)); + if (debug_attention_tp) { + add_debug_main_pair("q_after_join_vs_full", + debug_full_q_ref, q); + } + } + } + + // Preserve overwritten peer rows before updating the persistent ring. + // They are needed only by q>1 causal verification; q=1 merely mirrors + // its new row so the next wide verify never requires a bulk cache copy. + ggml_tensor * peer_old_rows = nullptr; + if (split_attention_values && !split_attention_direct_kv) { + for (int ti = 0; ti < n_tokens; ++ti) { + ggml_tensor * slot = ggml_view_2d( + ctx, attention_tp_cache->raw_kv, head_dim, 1, + attention_tp_cache->raw_kv->nb[1], + (size_t) ((kv_start + ti) % w.n_swa) * + attention_tp_cache->raw_kv->nb[1]); + ggml_tensor * saved = track_peer(ggml_cont(ctx, slot)); + ggml_build_forward_expand(gf, saved); + peer_old_rows = peer_old_rows + ? track_peer(ggml_concat( + ctx, peer_old_rows, saved, 1)) + : saved; + } + peer_old_rows = track_peer(ds4_cast_if_needed( + ctx, peer_old_rows, GGML_TYPE_F32)); + } + + ggml_tensor * peer_raw_source = nullptr; + if (split_attention_values && !split_attention_direct_kv) { + peer_raw_source = track_peer(ggml_set_rows( + ctx, attention_tp_cache->raw_kv, peer_kv, + attention_tp_peer_inputs->raw_kv_rows)); + ggml_build_forward_expand(gf, peer_raw_source); + } + + ggml_tensor * peer_comp_source = attention_tp_cache->comp_kv; + if (split_attention_values && !split_attention_direct_kv && + peer_comp_values) { + if (!peer_comp_source || + !attention_tp_peer_inputs->attn_comp_rows) { + return nullptr; + } + peer_comp_source = track_peer(ggml_set_rows( + ctx, peer_comp_source, peer_comp_values, + attention_tp_peer_inputs->attn_comp_rows)); + ggml_build_forward_expand(gf, peer_comp_source); + } + + if (split_attention_values) { + if (split_attention_direct_kv) { + peer_kv_attn = peer_kv; + } else { + ggml_tensor * peer_ring = ggml_view_2d( + ctx, peer_raw_source, head_dim, w.n_swa, + peer_raw_source->nb[1], 0); + peer_kv_attn = track_peer(ds4_cast_if_needed( + ctx, peer_ring, GGML_TYPE_F32)); + if (n_comp_attn > 0) { + if (!peer_comp_source || !comp_kv_source) return nullptr; + ggml_tensor * peer_comp = ggml_view_2d( + ctx, peer_comp_source, head_dim, n_comp_attn, + peer_comp_source->nb[1], 0); + peer_comp = track_peer(ds4_cast_if_needed( + ctx, peer_comp, GGML_TYPE_F32)); + peer_kv_attn = track_peer(ggml_concat( + ctx, peer_kv_attn, peer_comp, 1)); + } + peer_kv_attn = track_peer(ggml_concat( + ctx, peer_kv_attn, peer_old_rows, 1)); + } + add_debug_pair("kv", kv_attn, peer_kv_attn); + } + } + // Build one additive mask tensor and share it between the explicit and // flash-attention implementations. ggml flash attention expects // [n_kv,n_query] F16; the explicit path broadcasts the same values over @@ -2027,6 +2614,7 @@ static ggml_tensor * build_mla_attention( } } ggml_tensor * context = nullptr; + ggml_tensor * peer_context = nullptr; bool inverse_rope_fused = false; const bool use_flash = attention_impl != DeepSeek4AttentionImpl::Explicit && n_tokens > 1; @@ -2198,49 +2786,146 @@ static ggml_tensor * build_mla_attention( w.rope_yarn_beta_slow, rope_n_ctx_orig, fuse_q_rope); inverse_rope_fused = true; } + } } else { - // Flatten q to [head_dim, n_head*n_tokens] for batched matmul. - ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, - n_head * n_tokens); - ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); - scores = ggml_scale(ctx, scores, kq_scale); - if (score_mask) { - if (n_tokens > 1) { - ggml_tensor * m3 = ggml_reshape_3d(ctx, score_mask, - n_attn, 1, n_tokens); - ggml_tensor * s3 = ggml_reshape_3d(ctx, scores, - n_attn, n_head, n_tokens); - scores = ggml_reshape_2d(ctx, ggml_add(ctx, s3, m3), - n_attn, n_head * n_tokens); + struct ExplicitAttentionStages { + ggml_tensor * scores = nullptr; + ggml_tensor * probs = nullptr; + ggml_tensor * value_result = nullptr; + }; + auto build_explicit_attention = [&](ggml_tensor * q_shard, + ggml_tensor * kv_shard, + ggml_tensor * shard_score_mask, + ggml_tensor * sinks, + int shard_heads, + bool peer, + ExplicitAttentionStages * stages) { + auto track = [&](ggml_tensor * tensor) { + return peer ? track_peer(tensor) : track_main(tensor); + }; + ggml_tensor * q_flat = ggml_reshape_2d( + ctx, q_shard, head_dim, shard_heads * n_tokens); + ggml_tensor * scores = track(ggml_mul_mat( + ctx, kv_shard, q_flat)); + scores = track(ggml_scale(ctx, scores, kq_scale)); + if (shard_score_mask) { + if (n_tokens > 1) { + ggml_tensor * m3 = ggml_reshape_3d( + ctx, shard_score_mask, n_attn, 1, n_tokens); + ggml_tensor * s3 = ggml_reshape_3d( + ctx, scores, n_attn, shard_heads, n_tokens); + scores = track(ggml_reshape_2d( + ctx, ggml_add(ctx, s3, m3), n_attn, + shard_heads * n_tokens)); + } else { + scores = track(ggml_add( + ctx, scores, shard_score_mask)); + } + } + if (stages) stages->scores = scores; + + ggml_tensor * probs = nullptr; + if (sinks) { + ggml_tensor * sink_scores = ggml_reshape_2d( + ctx, sinks, 1, shard_heads); + if (n_tokens > 1) { + ggml_tensor * sink_shape = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, + shard_heads * n_tokens); + sink_scores = track(ggml_repeat( + ctx, sink_scores, sink_shape)); + } + ggml_tensor * scores_with_sink = track(ggml_concat( + ctx, scores, sink_scores, 0)); + ggml_tensor * probs_with_sink = track(ggml_soft_max( + ctx, scores_with_sink)); + probs = ggml_view_2d( + ctx, probs_with_sink, n_attn, + shard_heads * n_tokens, probs_with_sink->nb[1], 0); } else { - scores = ggml_add(ctx, scores, score_mask); - } - } - - // DS4 adds one learned per-head sink logit to the denominator, but the - // sink contributes no value vector. - ggml_tensor * probs = nullptr; - if (L.attn_sinks) { - ggml_tensor * sink_scores = ggml_reshape_2d(ctx, L.attn_sinks, - 1, n_head); - if (n_tokens > 1) { - ggml_tensor * sink_shape = ggml_new_tensor_2d( - ctx, GGML_TYPE_F32, 1, n_head * n_tokens); - sink_scores = ggml_repeat(ctx, sink_scores, sink_shape); - } - ggml_tensor * scores_with_sink = ggml_concat(ctx, scores, - sink_scores, 0); - ggml_tensor * probs_with_sink = ggml_soft_max(ctx, - scores_with_sink); - probs = ggml_view_2d(ctx, probs_with_sink, n_attn, - n_head * n_tokens, probs_with_sink->nb[1], 0); - } else { - probs = ggml_soft_max(ctx, scores); + probs = track(ggml_soft_max(ctx, scores)); + } + if (stages) stages->probs = probs; + ggml_tensor * kv_t = track(ggml_cont( + ctx, ggml_transpose(ctx, kv_shard))); + ggml_tensor * result = track(ggml_mul_mat(ctx, kv_t, probs)); + if (stages) stages->value_result = result; + return track(ggml_reshape_3d( + ctx, result, head_dim, shard_heads, n_tokens)); + }; + + ggml_tensor * main_sinks = L.attn_sinks; + if (split_attention_values && main_sinks) { + main_sinks = ggml_view_1d(ctx, main_sinks, n_head, 0); + } + const int main_attention_heads = split_attention_values + ? n_head : (split_attention ? total_head : n_head); + context = build_explicit_attention( + q, kv_attn, score_mask, main_sinks, main_attention_heads, + false, nullptr); + if (split_attention_values) { + ggml_tensor * peer_score_mask = ggml_reshape_2d( + ctx, attention_tp_peer_inputs->attn_row_mask, + n_attn, n_tokens); + ExplicitAttentionStages peer_stages{}; + if (split_attention_flash) { + // One peer flash-attention launch replaces the score, scale, + // mask, softmax, transpose, and value kernels. The main shard + // retains the exact explicit path; this is independently + // qualified because only the peer-head result can differ. + ggml_tensor * peer_q_fa = track_peer(ggml_permute( + ctx, peer_q, 0, 2, 1, 3)); + ggml_tensor * peer_k_fa = track_peer(ggml_reshape_3d( + ctx, peer_kv_attn, head_dim, n_attn, 1)); + ggml_tensor * peer_mask_fa = track_peer(ds4_cast_if_needed( + ctx, peer_score_mask, GGML_TYPE_F16)); + if (!ggml_is_contiguous(peer_mask_fa)) { + peer_mask_fa = track_peer(ggml_cont(ctx, peer_mask_fa)); + } + peer_context = track_peer(ggml_flash_attn_ext( + ctx, peer_q_fa, peer_k_fa, peer_k_fa, peer_mask_fa, + kq_scale, 0.0f, 0.0f)); + if (peer_weights->attn_sinks) { + ggml_flash_attn_ext_add_sinks( + peer_context, peer_weights->attn_sinks); + } + ggml_flash_attn_ext_set_prec(peer_context, GGML_PREC_F32); + ggml_flash_attn_ext_set_ds4_sparse( + peer_context, n_raw, w.n_swa, 0, 32); + } else { + peer_context = build_explicit_attention( + peer_q, peer_kv_attn, peer_score_mask, + peer_weights->attn_sinks, + peer_head, true, &peer_stages); + } + if (debug_attention_tp) { + ExplicitAttentionStages ref_stages{}; + ggml_tensor * main_peer_sinks = L.attn_sinks + ? ggml_view_1d( + ctx, L.attn_sinks, peer_head, + (size_t) n_head * L.attn_sinks->nb[0]) + : nullptr; + debug_peer_context_ref = build_explicit_attention( + debug_peer_q_ref, kv_attn, score_mask, + main_peer_sinks, peer_head, false, &ref_stages); + add_debug_pair("score_mask", score_mask, peer_score_mask); + if (!split_attention_flash) { + add_debug_pair("scores", ref_stages.scores, + peer_stages.scores); + add_debug_pair("probs", ref_stages.probs, + peer_stages.probs); + add_debug_pair("value_result", ref_stages.value_result, + peer_stages.value_result); + } + add_debug_pair("context_pre_rope", + debug_peer_context_ref, peer_context); + + debug_full_context_ref = build_explicit_attention( + debug_full_q_ref, kv_attn, score_mask, + L.attn_sinks, total_head, false, nullptr); + } } - ggml_tensor * kv_t = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); - context = ggml_mul_mat(ctx, kv_t, probs); - context = ggml_reshape_3d(ctx, context, head_dim, n_head, n_tokens); } // ── Inverse tail RoPE on attention output ─────────────────────── @@ -2254,15 +2939,66 @@ static ggml_tensor * build_mla_attention( i32_array_inputs.push_back({neg_pos, std::move(neg_vals)}); } context = build_tail_rope_3d( - ctx, context, neg_pos, n_rot, head_dim, n_head, n_tokens, - rope_freq, rope_scale, rope_ext, rope_attn, + ctx, context, neg_pos, n_rot, head_dim, + split_attention_values + ? n_head + : (split_attention ? total_head : n_head), + n_tokens, rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, - rope_n_ctx_orig); + rope_n_ctx_orig, + split_attention ? &attention_tp->main_nodes : nullptr); + if (split_attention_values) { + // Keep the peer shard through inverse RoPE as well. Joining the + // head contexts here transferred four times more data and left + // the complete grouped output-A projection on the main device. + peer_context = track_peer(build_tail_rope_3d( + ctx, peer_context, attention_tp_peer_inputs->neg_pos, + n_rot, head_dim, peer_head, n_tokens, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + rope_n_ctx_orig, &attention_tp->peer_nodes)); + } + if (debug_attention_tp && split_attention_values) { + debug_full_context_ref = track_main(build_tail_rope_3d( + ctx, debug_full_context_ref, neg_pos, n_rot, head_dim, + total_head, n_tokens, rope_freq, rope_scale, rope_ext, + rope_attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope_n_ctx_orig)); + ggml_tensor * debug_main_context_ref = ggml_view_3d( + ctx, debug_full_context_ref, head_dim, n_head, n_tokens, + debug_full_context_ref->nb[1], + debug_full_context_ref->nb[2], 0); + ggml_tensor * debug_peer_context_post_ref = ggml_view_3d( + ctx, debug_full_context_ref, head_dim, peer_head, n_tokens, + debug_full_context_ref->nb[1], + debug_full_context_ref->nb[2], + (size_t) n_head * debug_full_context_ref->nb[1]); + add_debug_main_pair("main_context_post_rope", + debug_main_context_ref, context); + add_debug_pair("peer_context_post_rope", + debug_peer_context_post_ref, peer_context); + } + } + + if (split_attention_core_only) { + // Core-only TP rejoins immediately after score/value attention. The + // large Q and output projections therefore keep their efficient full + // width on the R9700, while only the context-dependent attention core + // is divided between devices. + peer_context = track_peer(ggml_cont(ctx, peer_context)); + ggml_set_output(peer_context); + ggml_build_forward_expand(gf, peer_context); + ggml_tensor * peer_context_ready = + ggml_ds4_deferred_peer_copy(ctx, peer_context); + ggml_set_input(peer_context_ready); + ggml_set_output(peer_context_ready); + attention_tp->deferred_peer_copy_nodes.push_back( + peer_context_ready); + track_main(peer_context_ready); + context = track_main(ggml_concat( + ctx, context, peer_context_ready, 1)); } - // Flatten to [head_dim*n_head, n_tokens] for output projection - ggml_tensor * attn_out = ggml_reshape_2d(ctx, context, head_dim * n_head, n_tokens); - // ── Grouped output projection ────────────────────────────────── // DS4 output uses grouped low-rank projection: // attn_out: [head_dim*n_head, n_tokens] → reshape [group_dim, n_tokens, n_groups] @@ -2270,32 +3006,145 @@ static ggml_tensor * build_mla_attention( // batched matmul over n_groups: → [n_lora_o, n_tokens, n_groups] // → reshape [n_lora_o*n_groups, n_tokens] // out_b: [n_lora_o*n_groups, n_embd] → final: [n_embd, n_tokens] - const int group_dim = head_dim * (n_head / n_out_group); // 512 * 8 = 4096 - // Reshape attn_out: [32768, n_tokens] → [4096, 8, n_tokens] → permute to [4096, n_tokens, 8] - attn_out = ggml_reshape_3d(ctx, attn_out, group_dim, n_out_group, n_tokens); - attn_out = ggml_permute(ctx, attn_out, 0, 2, 1, 3); - if (n_tokens == 1) { - attn_out = ggml_cont(ctx, attn_out); + const int group_dim = head_dim * heads_per_group; // 512 * 8 = 4096 + auto build_grouped_output_a = [&](ggml_tensor * shard_context, + ggml_tensor * shard_weights, + int shard_heads, + int shard_groups, + bool peer) { + auto track = [&](ggml_tensor * tensor) { + return peer ? track_peer(tensor) : track_main(tensor); + }; + ggml_tensor * shard_out = ggml_reshape_2d( + ctx, shard_context, head_dim * shard_heads, n_tokens); + shard_out = ggml_reshape_3d( + ctx, shard_out, group_dim, shard_groups, n_tokens); + shard_out = ggml_permute(ctx, shard_out, 0, 2, 1, 3); + if (n_tokens == 1) { + shard_out = track(ggml_cont(ctx, shard_out)); + } + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, shard_weights, group_dim, n_lora_o, shard_groups); + return track(ggml_mul_mat(ctx, weights_3d, shard_out)); + }; + + ggml_tensor * attn_low = nullptr; + ggml_tensor * split_main_output = nullptr; + if (split_attention_output_projection || + split_attention_projection_only) { + ggml_tensor * main_projection_context = context; + ggml_tensor * peer_projection_context = peer_context; + if (split_attention_projection_only) { + // Keep the latency-sensitive score/value attention on the main + // GPU, then fork the independent output-head groups. This costs + // one compact context copy per layer and lets both devices stream + // their own persistent output-A/B weights in parallel without + // duplicating the KV cache or its many small kernels on the peer. + main_projection_context = ggml_view_3d( + ctx, context, head_dim, n_head, n_tokens, + context->nb[1], context->nb[2], 0); + // Slicing the head axis leaves the full-head token stride, so the + // main shard is not reshapeable until it is packed. The peer + // shard is packed by its cross-device cont below for the same + // reason. + main_projection_context = track_main( + ggml_cont(ctx, main_projection_context)); + ggml_tensor * peer_context_view = ggml_view_3d( + ctx, context, head_dim, peer_head, n_tokens, + context->nb[1], context->nb[2], + (size_t) n_head * context->nb[1]); + peer_projection_context = track_peer( + ggml_cont(ctx, peer_context_view)); + } + + ggml_tensor * main_output_a = ggml_view_2d( + ctx, L.attn_output_a, group_dim, + (int64_t) n_lora_o * main_groups, + L.attn_output_a->nb[1], 0); + ggml_tensor * main_low = build_grouped_output_a( + main_projection_context, main_output_a, + n_head, main_groups, false); + ggml_tensor * peer_low = build_grouped_output_a( + peer_projection_context, peer_weights->attn_output_a, + peer_head, peer_groups, true); + + if (split_attention_output_b) { + const int64_t main_lora = (int64_t) n_lora_o * main_groups; + ggml_tensor * main_output_b = ggml_view_2d( + ctx, L.attn_output_b, main_lora, n_embd, + L.attn_output_b->nb[1], 0); + auto project_output_b = [&](ggml_tensor * weights, + ggml_tensor * low, + int groups, + bool peer) { + ggml_tensor * result = groups > 1 + ? ggml_mul_mat_grouped_src(ctx, weights, low) + : ggml_mul_mat( + ctx, weights, + ggml_reshape_2d( + ctx, low, n_lora_o, n_tokens)); + return peer ? track_peer(result) : track_main(result); + }; + split_main_output = project_output_b( + main_output_b, main_low, main_groups, false); + ggml_tensor * split_peer_output = project_output_b( + peer_weights->attn_output_b, peer_low, + peer_groups, true); + + // Carry both partial output-B projections directly into the + // existing fused HC-post join. This is the longest owner boundary: + // Q, attention, inverse RoPE, output-A, and output-B all overlap, + // while HC post performs the one required peer+main addition. + split_main_output = track_main(ggml_cont( + ctx, split_main_output)); + split_peer_output = track_peer(ggml_cont( + ctx, split_peer_output)); + ggml_set_output(split_peer_output); + ggml_build_forward_expand(gf, split_peer_output); + ggml_tensor * split_peer_ready = + ggml_ds4_deferred_peer_copy(ctx, split_peer_output); + ggml_set_input(split_peer_ready); + ggml_set_output(split_peer_ready); + attention_tp->deferred_peer_copy_nodes.push_back( + split_peer_ready); + track_main(split_peer_ready); + attention_tp->main_output = split_main_output; + attention_tp->peer_output = split_peer_ready; + } else { + // Join the low-rank group result. For q=5 and a 2/8 peer + // split this is 40 KiB/layer instead of 160 KiB of contexts. + // Output-B remains one unchanged main-device reduction. + peer_low = track_peer(ggml_cont(ctx, peer_low)); + ggml_set_output(peer_low); + ggml_build_forward_expand(gf, peer_low); + ggml_tensor * peer_low_ready = + ggml_ds4_deferred_peer_copy(ctx, peer_low); + ggml_set_input(peer_low_ready); + ggml_set_output(peer_low_ready); + attention_tp->deferred_peer_copy_nodes.push_back(peer_low_ready); + track_main(peer_low_ready); + attn_low = track_main(ggml_concat( + ctx, main_low, peer_low_ready, 2)); + } + } else { + const int output_n_head = split_attention ? total_head : n_head; + attn_low = build_grouped_output_a( + context, L.attn_output_a, output_n_head, + n_out_group, false); } - // attn_out is now [group_dim, n_tokens, n_out_group] - ggml_tensor * out_a_3d = ggml_reshape_3d(ctx, L.attn_output_a, group_dim, n_lora_o, n_out_group); - // out_a_3d: [group_dim, n_lora_o, n_out_group] — ne[2] matches - ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); // attn_low: [n_lora_o, n_tokens, n_out_group] + if (split_attention_output_b) { + return split_main_output; + } + ggml_tensor * out = nullptr; - const bool grouped_output_projection = - n_tokens > 1 && - !ds4_env_flag("DFLASH_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION"); - if (grouped_output_projection) { + if (n_tokens > 1) { // Batched ROCmFPX MMQ consumes src1's channel stride directly. This // avoids materializing both permutations (~256 MiB/layer at 2K). out = ggml_mul_mat_grouped_src(ctx, L.attn_output_b, attn_low); } else { - // Preserve the established single-token graph and provide an exact - // fallback for heterogeneous runtimes that cannot retain grouped-view - // metadata across a scheduler copy. At verifier widths (q <= 4), this - // materializes at most 128 KiB per layer rather than the long-prefill - // volume avoided by the grouped path. + // Preserve the established single-token graph and its numerical + // behavior. Decode is intentionally outside the prefill fast path. attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); attn_low = ggml_reshape_2d( ctx, attn_low, n_lora_o * n_out_group, n_tokens); @@ -3021,8 +3870,7 @@ static bool eval_ds4_hybrid( ggml_tensor * ffn_normed_backend = nullptr, const MoeHybridDeviceOutputs * device_outputs = nullptr) { const auto ffn_t0 = Ds4TimingClock::now(); - if (!storage.cold_expert_ids.empty() && - !storage.down_cold && !storage.gate_up_cold && + if (!storage.down_cold && !storage.gate_up_cold && !(expert_compute && expert_layer)) { if (!hybrid_owner || !stream_engine || !stream_engine->is_ready() || !hybrid_owner->has_mmap() || @@ -4581,13 +5429,10 @@ bool deepseek4_step( // while a variant recurs, which is what the ggml-cuda/HIP graph cache keys // on, enabling graph replay for the bulk of decode steps. -static bool ds4_fused_decode_enabled(const DeepSeek4Weights & w) { - // The supported control is --ds4-fused-decode, propagated through the - // loaded weights. Keep the old environment spelling as a compatibility - // fallback for existing launch scripts. - static const bool legacy_env_enabled = +static bool ds4_fused_decode_enabled() { + static const bool enabled = ds4_env_flag("DFLASH_DS4_FUSED_DECODE"); - return w.fused_decode || legacy_env_enabled; + return enabled; } struct DeepSeek4FusedDecodeGraph { @@ -4606,8 +5451,22 @@ struct DeepSeek4FusedDecodeGraph { ggml_tensor * i32_bundle = nullptr; ggml_tensor * i64_bundle = nullptr; ggml_tensor * mask_bundle = nullptr; // additive score mask (0 / -1e30), may be null + // Stable inputs allocated directly on the attention peer. They mirror the + // tiny host-authored values and therefore never traverse a captured HIP + // peer-copy edge. + // Keep positive and negative positions as independent peer inputs. A + // second-half view into one bundled input produced the correct positive + // Q rotation but stale values for the inverse output rotation when the + // fused graph was scheduled across gfx1201 + gfx1151. Dedicated tensors + // avoid that scheduler/view-offset hazard without adding a device copy. + ggml_tensor * attention_tp_pos_q = nullptr; + ggml_tensor * attention_tp_neg_q = nullptr; + ggml_tensor * attention_tp_rawrows = nullptr; + ggml_tensor * attention_tp_i64_bundle = nullptr; // compressor row ids + ggml_tensor * attention_tp_mask_bundle = nullptr; std::vector hash_ids; std::vector hybrid_inputs; + std::vector attention_tp_inputs; std::vector authoritative_routes; ggml_tensor * logits = nullptr; ggml_backend_sched_t sched = nullptr; @@ -4617,9 +5476,15 @@ struct DeepSeek4FusedDecodeGraph { i32_bundle = nullptr; i64_bundle = nullptr; mask_bundle = nullptr; + attention_tp_pos_q = nullptr; + attention_tp_neg_q = nullptr; + attention_tp_rawrows = nullptr; + attention_tp_i64_bundle = nullptr; + attention_tp_mask_bundle = nullptr; logits = nullptr; hash_ids.clear(); hybrid_inputs.clear(); + attention_tp_inputs.clear(); authoritative_routes.clear(); shape_key.clear(); last_use = 0; @@ -4651,6 +5516,11 @@ struct DeepSeek4FusedDecodeGraph { // Native graph executables outlive ggml graph metadata in the backend // cache. Retire them before either the scheduler or metadata arena is // released, otherwise a rebuilt slot can inherit the same pointer key. + // A heterogeneous graph may have queued work on both device streams; + // finish it before invalidating native graph executables or events. + if (sched) { + ggml_backend_sched_synchronize(sched); + } invalidate_native_graphs(main_backend, peer_backend); if (sched) { ggml_backend_sched_free(sched); @@ -6075,7 +6945,6 @@ static int ds4_try_layer_major_prefill( int kv_start, std::vector & out_logits, const int32_t * token_ids, - Ds4VerifyHooks * verify_hooks, DeepSeek4StepTelemetry * telemetry) { if (!backend || !embed || n_tokens <= 4 || n_tokens > DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS || @@ -6083,14 +6952,6 @@ static int ds4_try_layer_major_prefill( return 0; } if (cache.prefill_mode == PrefillAttentionMode::Exact) return 0; - // Layer-major prefill returns only the final-position logits. DSpark's - // per-layer feature capture is supported below, but verifier requests for - // every position's logits/argmax must retain the generic path. - if (verify_hooks && - (verify_hooks->all_logits_out || verify_hooks->argmax_out || - verify_hooks->prefer_argmax_only)) { - return 0; - } if (!ds4_backend_is_gpu(backend) || !hc_out_weights.loaded || hc_out_weights.scale_data.empty() || !w.output_hc_fn || !w.output_hc_base) { @@ -6124,64 +6985,6 @@ static int ds4_try_layer_major_prefill( const int64_t mix_dim = 2 * (int64_t) n_hc + (int64_t) n_hc * n_hc; const int next_pos = kv_start + n_tokens; - const std::vector * capture_layer_ids = - verify_hooks ? verify_hooks->capture_layer_ids : nullptr; - std::vector * capture_out = - verify_hooks ? verify_hooks->capture_out : nullptr; - const bool capture_enabled = capture_layer_ids && capture_out && - !capture_layer_ids->empty(); - const int capture_begin = verify_hooks - ? std::clamp(verify_hooks->capture_token_begin, 0, n_tokens) : 0; - const int requested_capture_end = - verify_hooks && verify_hooks->capture_token_end >= 0 - ? verify_hooks->capture_token_end : n_tokens; - const int capture_end = std::clamp( - requested_capture_end, capture_begin, n_tokens); - const int capture_tokens = capture_end - capture_begin; - std::vector capture_hc_state; - if (capture_out) { - capture_out->clear(); - } - if (capture_enabled) { - capture_out->assign( - (size_t) n_tokens * capture_layer_ids->size() * n_embd, 0.0f); - capture_hc_state.resize((size_t) hc_dim * capture_tokens); - } - const auto capture_layer = [&](int layer, ggml_tensor * state) { - if (!capture_enabled || capture_tokens <= 0 || !state) return; - const auto first = std::find(capture_layer_ids->begin(), - capture_layer_ids->end(), layer); - if (first == capture_layer_ids->end()) return; - - const auto read_t0 = Ds4TimingClock::now(); - const size_t capture_offset = - (size_t) capture_begin * hc_dim * sizeof(float); - ggml_backend_tensor_get(state, capture_hc_state.data(), capture_offset, - sizeof(float) * capture_hc_state.size()); - if (telemetry) { - telemetry->full_graph_read_us += ds4_elapsed_us( - read_t0, Ds4TimingClock::now()); - } - - const size_t n_capture = capture_layer_ids->size(); - for (size_t ci = 0; ci < n_capture; ++ci) { - if ((*capture_layer_ids)[ci] != layer) continue; - for (int t = capture_begin; t < capture_end; ++t) { - float * dst = capture_out->data() + - ((size_t) t * n_capture + ci) * n_embd; - const float * src = capture_hc_state.data() + - (size_t) (t - capture_begin) * hc_dim; - for (int d = 0; d < n_embd; ++d) { - float sum = 0.0f; - for (int h = 0; h < n_hc; ++h) { - sum += src[(size_t) h * n_embd + d]; - } - dst[d] = sum / (float) n_hc; - } - } - } - }; - Ds4LayerMajorGraphCache * graph_cache = nullptr; bool cache_hit = false; bool cache_build = false; @@ -6351,7 +7154,6 @@ static int ds4_try_layer_major_prefill( telemetry->full_graph_compute_us += ds4_elapsed_us( compute_t0, Ds4TimingClock::now()); } - capture_layer(il, (il & 1) == 0 ? state_b : state_a); if (layer.logits) { out_logits.resize((size_t) w.n_vocab); ggml_backend_tensor_get( @@ -6578,8 +7380,6 @@ static int ds4_try_layer_major_prefill( compute_t0, Ds4TimingClock::now()); } - capture_layer(il, state_out); - if (logits) { out_logits.resize((size_t) w.n_vocab); ggml_backend_tensor_get(logits, out_logits.data(), 0, @@ -6764,19 +7564,6 @@ bool deepseek4_step_layer_range( n_tokens > 4 && n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && layer_begin == 0 && is_last_shard && out_logits && ds4_backend_is_gpu(backend); - const bool layer_major_hooks_supported = - !verify_hooks || - (!verify_hooks->all_logits_out && !verify_hooks->argmax_out && - !verify_hooks->prefer_argmax_only); - // The standard layer-major pipeline owns an exact batched compressor. - // Let it see the wide prompt before the generic boundary splitter turns - // the request into ratio-sized (typically four-token) forwards. It also - // owns DSpark feature capture, so the final capture window stays batched. - const bool standard_layer_major_prefill = - !w.moe_hybrid && cache.prefill_mode != PrefillAttentionMode::Exact && - n_tokens > 4 && n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && - layer_begin == 0 && is_last_shard && out_logits && - ds4_backend_is_gpu(backend) && layer_major_hooks_supported; // These graphs are rebuilt around an owner join on every layer, so tensor // metadata addresses can be recycled for different topologies. Until // the full heterogeneous layer is captured as one stable scheduler graph, @@ -6794,8 +7581,7 @@ bool deepseek4_step_layer_range( // as sequential execution while retaining safe batched prefixes. const int first_chunk = deepseek4_safe_compressor_batch_tokens(w, kv_start, n_tokens); if (first_chunk > 0 && first_chunk < n_tokens && - !fused_verify_candidate && !heterogeneous_sparse_prefill && - !standard_layer_major_prefill) { + !fused_verify_candidate && !heterogeneous_sparse_prefill) { const int input_width = layer_begin == 0 ? n_embd : hc_dim; std::vector hc_all; std::vector shard_out_all; @@ -6953,13 +7739,14 @@ bool deepseek4_step_layer_range( // Large full-model prefill batches use the device-resident layer-major // pipeline. DSpark verification remains on its exact q=2..4 path below. - if (standard_layer_major_prefill) { + if (!fused_verify_candidate && n_tokens > 4 && + n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && layer_begin == 0 && + is_last_shard && out_logits && ds4_backend_is_gpu(backend)) { const int prc = ds4_try_layer_major_prefill( fused_decode_graph_cache, backend, w, cache, hc_layer_weights_range, hc_output_weights_range, hash_routing_tables_range, scratch.hash_expert_ids, embed, - n_tokens, kv_start, *out_logits, token_ids, verify_hooks, - telemetry); + n_tokens, kv_start, *out_logits, token_ids, telemetry); if (prc < 0) return false; if (prc > 0) { if (telemetry) { @@ -7003,6 +7790,21 @@ bool deepseek4_step_layer_range( Ds4FusedVerifyCache & graph_cache = q1_feature_capture ? layer_range_cache.fused_capture_graph_cache : layer_range_cache.fused_verify_graph_cache; + // q>1 executes the head split; q=1 still mirrors its newly written KV + // rows through the same graph. Treat both as coherent-cache steps so a + // fallback token cannot force a 103 MiB bulk synchronization later. + const bool attention_tp_step = + w.attention_tp_backend && w.attention_tp_peer_groups > 0; + if (attention_tp_step) { + if (!deepseek4_sync_attention_tp_cache(cache)) { + std::fprintf(stderr, + "[deepseek4-attention-tp] peer KV sync failed\n"); + return false; + } + // A failed or unavailable graph must force a resync before any + // later attempt; publish the new position only after success. + cache.attention_tp_cur_pos = -1; + } const int vrc = ds4_try_fused_verify_step( graph_cache, q1_feature_capture, fused_decode_graph_cache, backend, w, cache, @@ -7020,6 +7822,9 @@ bool deepseek4_step_layer_range( if (vratio == 4) cache.layers[il].n_index_comp = std::max(cache.layers[il].n_index_comp, np / (int) vratio); } cache.cur_pos = np; + if (attention_tp_step) { + cache.attention_tp_cur_pos = np; + } if (telemetry) telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); return true; } @@ -7038,8 +7843,7 @@ bool deepseek4_step_layer_range( if (!moe_hybrid && n_tokens == 1 && allow_decode_graph_reuse && layer_begin == 0 && is_last_shard && !(verify_hooks && verify_hooks->capture_layer_ids && verify_hooks->capture_out) && - out_logits && ds4_backend_is_gpu(backend) && - ds4_fused_decode_enabled(w)) { + out_logits && ds4_backend_is_gpu(backend) && ds4_fused_decode_enabled()) { const int rc = ds4_try_fused_decode_step( fused_decode_graph_cache, backend, w, cache, hc_layer_weights_range, hc_output_weights_range, hash_routing_tables_range, scratch.hash_expert_ids, @@ -7127,38 +7931,6 @@ bool deepseek4_step_layer_range( hc_state.data(), 0, sizeof(float) * hc_state.size()); hc_state_backend = cached_decode_hc_post_graph.residual_hc; } - const auto capture_requested = [&](int layer) { - if (!verify_hooks || !verify_hooks->capture_layer_ids || - !verify_hooks->capture_out) { - return false; - } - const std::vector & ids = *verify_hooks->capture_layer_ids; - return std::find(ids.begin(), ids.end(), layer) != ids.end(); - }; - const auto capture_hc_layer = [&](int layer, const float * state) { - if (!state || !capture_requested(layer)) return; - const std::vector & ids = *verify_hooks->capture_layer_ids; - std::vector & capture = *verify_hooks->capture_out; - if ((int) capture.size() != (int) ids.size() * n_embd * n_tokens) { - capture.assign( - (size_t) ids.size() * n_embd * n_tokens, 0.0f); - } - for (size_t ci = 0; ci < ids.size(); ++ci) { - if (ids[ci] != layer) continue; - for (int t = 0; t < n_tokens; ++t) { - float * dst = capture.data() + - (size_t) t * ids.size() * n_embd + ci * n_embd; - const float * hs = state + (size_t) t * hc_dim; - for (int d = 0; d < n_embd; ++d) { - float sum = 0.0f; - for (int h = 0; h < n_hc; ++h) { - sum += hs[(size_t) h * n_embd + d]; - } - dst[d] = sum / (float) n_hc; - } - } - } - }; for (int il = layer_begin; il < layer_end; ++il) { const DeepSeek4Layer & L = w.layers[(size_t)il]; DeepSeek4LayerCache & lc = cache.layers[(size_t)il]; @@ -7909,15 +8681,25 @@ bool deepseek4_step_layer_range( n_hc); std::memcpy(hc_state.data(), next_hc.data(), next_hc.size() * sizeof(float)); if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); - capture_hc_layer(il, hc_state.data()); - } - if ((use_backend_prefill_hc || use_backend_decode_hc_graph || - use_backend_decode_hc_direct) && - hc_state_backend && capture_requested(il)) { - ggml_backend_tensor_get( - hc_state_backend, hc_state.data(), 0, - sizeof(float) * hc_state.size()); - capture_hc_layer(il, hc_state.data()); + if (verify_hooks && verify_hooks->capture_layer_ids && verify_hooks->capture_out) { + const std::vector & _ids = *verify_hooks->capture_layer_ids; + for (size_t _ci = 0; _ci < _ids.size(); ++_ci) { + if (_ids[_ci] != il) continue; + const int _ncap = (int) _ids.size(); + std::vector & _cap = *verify_hooks->capture_out; + if ((int) _cap.size() != _ncap * n_embd * n_tokens) + _cap.assign((size_t) _ncap * n_embd * n_tokens, 0.0f); + for (int _t = 0; _t < n_tokens; ++_t) { + float * _dst = _cap.data() + (size_t) _t * _ncap * n_embd + (size_t) _ci * n_embd; + const float * _hs = hc_state.data() + (size_t) _t * hc_dim; + for (int _d = 0; _d < n_embd; ++_d) { + float _acc = 0.0f; + for (int _h = 0; _h < n_hc; ++_h) _acc += _hs[(size_t) _h * n_embd + _d]; + _dst[_d] = _acc / (float) n_hc; + } + } + } + } } } } @@ -7966,9 +8748,7 @@ bool deepseek4_step_layer_range( ggml_context * ctx = ggml_init(params); if (!ctx) return false; - const bool need_all_logits = - verify_hooks && verify_hooks->all_logits_out; - const bool last_only = n_tokens > 1 && !need_all_logits; + const bool last_only = n_tokens > 1; const int output_tokens = last_only ? 1 : n_tokens; ggml_tensor * inp = ggml_new_tensor_2d( ctx, GGML_TYPE_F32, n_embd, output_tokens); @@ -8121,6 +8901,57 @@ bool create_deepseek4_cache(ggml_backend_t backend, } ggml_backend_buffer_clear(out.buf, 0); + if (w.attention_tp_backend && w.attention_tp_peer_groups > 0) { + ggml_init_params peer_params{}; + peer_params.mem_size = ggml_tensor_overhead() * + (size_t) (w.n_layer * 2 + 8) + 4096; + peer_params.no_alloc = true; + out.attention_tp_ctx = ggml_init(peer_params); + if (!out.attention_tp_ctx) { + ggml_backend_buffer_free(out.buf); + out.buf = nullptr; + ggml_free(out.ctx); + out.ctx = nullptr; + return false; + } + out.attention_tp_layers.resize((size_t) w.n_layer); + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4LayerCache & source = out.layers[(size_t) il]; + DeepSeek4AttentionTpCacheLayer & destination = + out.attention_tp_layers[(size_t) il]; + destination.raw_kv = ggml_dup_tensor( + out.attention_tp_ctx, source.raw_kv); + char name[64]; + std::snprintf(name, sizeof(name), "ds4_raw_kv_peer_%d", il); + ggml_set_name(destination.raw_kv, name); + if (source.comp_kv) { + destination.comp_kv = ggml_dup_tensor( + out.attention_tp_ctx, source.comp_kv); + std::snprintf(name, sizeof(name), + "ds4_comp_kv_peer_%d", il); + ggml_set_name(destination.comp_kv, name); + } + } + out.attention_tp_buf = ggml_backend_alloc_ctx_tensors( + out.attention_tp_ctx, w.attention_tp_backend); + if (!out.attention_tp_buf) { + out.attention_tp_layers.clear(); + ggml_free(out.attention_tp_ctx); + out.attention_tp_ctx = nullptr; + ggml_backend_buffer_free(out.buf); + out.buf = nullptr; + ggml_free(out.ctx); + out.ctx = nullptr; + return false; + } + out.attention_tp_backend = w.attention_tp_backend; + out.attention_tp_cur_pos = 0; + ggml_backend_buffer_clear(out.attention_tp_buf, 0); + std::fprintf(stderr, + "[deepseek4-attention-tp] peer KV cache: %.1f MB\n", + ggml_backend_buffer_get_size(out.attention_tp_buf) / + 1024.0 / 1024.0); + } const size_t total_bytes = ggml_backend_buffer_get_size(out.buf); std::fprintf(stderr, "[deepseek4] KV cache: %.1f MB for ctx=%d\n", (double)total_bytes / (1024.0 * 1024.0), max_ctx); @@ -8130,8 +8961,19 @@ bool create_deepseek4_cache(ggml_backend_t backend, void free_deepseek4_cache(DeepSeek4Cache & c) { delete c.layer_range_cache; c.layer_range_cache = nullptr; - if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } + if (c.attention_tp_buf) { + ggml_backend_buffer_free(c.attention_tp_buf); + c.attention_tp_buf = nullptr; + } + if (c.attention_tp_ctx) { + ggml_free(c.attention_tp_ctx); + c.attention_tp_ctx = nullptr; + } + c.attention_tp_backend = nullptr; + c.attention_tp_layers.clear(); + c.attention_tp_cur_pos = -1; if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } + if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } c.layers.clear(); c.hc_state = nullptr; } @@ -8145,6 +8987,56 @@ void reset_deepseek4_cache(DeepSeek4Cache & c) { if (c.buf) { ggml_backend_buffer_clear(c.buf, 0); } + if (c.attention_tp_buf) { + ggml_backend_buffer_clear(c.attention_tp_buf, 0); + c.attention_tp_cur_pos = 0; + } +} + +bool deepseek4_sync_attention_tp_cache(DeepSeek4Cache & c) { + if (!c.attention_tp_buf || !c.attention_tp_backend) return true; + if (c.attention_tp_cur_pos == c.cur_pos) return true; + // The qualified path splits only Q-head projection and never reads the + // peer KV mirror. Avoid a full cache copy after every speculative commit. + // Keep the old mirror available solely for the explicit value-split + // diagnostic while its deterministic-kernel work is unfinished. + if (!ds4_env_flag("DFLASH_DS4_ATTENTION_TP_MIRROR_KV")) { + c.attention_tp_cur_pos = c.cur_pos; + return true; + } + if (c.attention_tp_layers.size() != c.layers.size()) return false; + + const auto same_layout = [](const ggml_tensor * a, + const ggml_tensor * b) { + if (!a || !b || a->type != b->type) return false; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (a->ne[d] != b->ne[d] || a->nb[d] != b->nb[d]) { + return false; + } + } + return true; + }; + for (size_t il = 0; il < c.layers.size(); ++il) { + const DeepSeek4LayerCache & source = c.layers[il]; + DeepSeek4AttentionTpCacheLayer & destination = + c.attention_tp_layers[il]; + if (!same_layout(source.raw_kv, destination.raw_kv)) { + return false; + } + ggml_backend_tensor_copy(source.raw_kv, destination.raw_kv); + if (!!source.comp_kv != !!destination.comp_kv) return false; + if (source.comp_kv) { + if (!same_layout(source.comp_kv, destination.comp_kv)) { + return false; + } + ggml_backend_tensor_copy(source.comp_kv, destination.comp_kv); + } + } + c.attention_tp_cur_pos = c.cur_pos; + std::fprintf(stderr, + "[deepseek4-attention-tp] synchronized peer KV at pos=%d\n", + c.cur_pos); + return true; } namespace { @@ -8159,49 +9051,6 @@ ggml_tensor * clone_snapshot_tensor(ggml_context * ctx, return dst; } -ggml_tensor * clone_snapshot_rows(ggml_context * ctx, - const ggml_tensor * src, - int live_rows, - const char * name) { - if (!ctx || !src || ggml_n_dims(src) != 2 || live_rows < 0 || - live_rows > src->ne[1]) { - return nullptr; - } - // GGML tensors cannot have an empty physical dimension. Keep one - // allocated row for an empty logical prefix, but copy zero bytes below. - const int64_t allocated_rows = std::max(1, live_rows); - ggml_tensor * dst = ggml_new_tensor_2d( - ctx, src->type, src->ne[0], allocated_rows); - if (!dst) return nullptr; - if (name && *name) ggml_set_name(dst, name); - return dst; -} - -size_t tensor_prefix_bytes(const ggml_tensor * tensor, int rows) { - if (!tensor || rows <= 0) return 0; - return ggml_row_size(tensor->type, tensor->ne[0]) * (size_t) rows; -} - -bool copy_tensor_prefix_from_backend(const ggml_tensor * src, - ggml_tensor * dst, - int rows) { - if (!src || !dst || rows < 0) return false; - const size_t bytes = tensor_prefix_bytes(src, rows); - if (bytes > ggml_nbytes(src) || bytes > ggml_nbytes(dst)) return false; - if (bytes > 0) ggml_backend_tensor_get(src, dst->data, 0, bytes); - return true; -} - -bool copy_tensor_prefix_to_backend(const ggml_tensor * src, - ggml_tensor * dst, - int rows) { - if (!src || !dst || rows < 0) return false; - const size_t bytes = tensor_prefix_bytes(src, rows); - if (bytes > ggml_nbytes(src) || bytes > ggml_nbytes(dst)) return false; - if (bytes > 0) ggml_backend_tensor_set(dst, src->data, 0, bytes); - return true; -} - bool copy_tensor_from_backend(const ggml_tensor * src, ggml_tensor * dst) { if (!src || !dst) return false; const size_t bytes = ggml_nbytes(src); @@ -8228,45 +9077,15 @@ bool tensors_compatible(const ggml_tensor * a, const ggml_tensor * b) { return true; } -bool prefix_tensors_compatible(const ggml_tensor * snap, - const ggml_tensor * cache, - int live_rows) { - if (!!snap != !!cache) return false; - if (!snap) return live_rows == 0; - // A right-sized zero/one-row tensor reports one logical GGML dimension, - // while its full-capacity cache tensor reports two. Compare the physical - // row layout instead of ggml_n_dims() so those valid snapshots restore. - if (live_rows < 0 || cache->ne[1] <= 0 || - snap->type != cache->type || - snap->ne[0] != cache->ne[0] || live_rows > cache->ne[1]) { - return false; - } - for (int i = 2; i < GGML_MAX_DIMS; ++i) { - if (snap->ne[i] != cache->ne[i]) return false; - } - return snap->ne[1] == std::max(1, live_rows); -} - } // 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 || cache.cur_pos < 0 || - cache.cur_pos > cache.max_ctx) { + cache.layers.size() != (size_t)cache.n_layer) { return false; } - for (const auto & layer : cache.layers) { - if (layer.n_comp < 0 || layer.n_index_comp < 0 || - (layer.comp_kv && layer.n_comp > layer.comp_kv->ne[1]) || - (!layer.comp_kv && layer.n_comp != 0) || - (layer.index_comp_kv && - layer.n_index_comp > layer.index_comp_kv->ne[1]) || - (!layer.index_comp_kv && layer.n_index_comp != 0)) { - return false; - } - } free_deepseek4_snapshot(out); @@ -8289,13 +9108,8 @@ bool deepseek4_snapshot_save(const DeepSeek4Cache & cache, 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 = src.comp_kv - ? clone_snapshot_rows(out.ctx, src.comp_kv, src.n_comp, nullptr) - : nullptr; - dst.index_comp_kv = src.index_comp_kv - ? clone_snapshot_rows(out.ctx, src.index_comp_kv, - src.n_index_comp, nullptr) - : 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 = @@ -8332,13 +9146,9 @@ bool deepseek4_snapshot_save(const DeepSeek4Cache & cache, 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_prefix_from_backend(src.comp_kv, dst.comp_kv, - src.n_comp)) || + (src.comp_kv && !copy_tensor_from_backend(src.comp_kv, dst.comp_kv)) || (src.index_comp_kv && - !copy_tensor_prefix_from_backend(src.index_comp_kv, - dst.index_comp_kv, - src.n_index_comp)) || + !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)) || @@ -8363,79 +9173,30 @@ bool deepseek4_snapshot_save(const DeepSeek4Cache & cache, 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() || snap.cur_pos < 0 || - snap.cur_pos > cache.max_ctx) { - std::fprintf(stderr, - "[deepseek4] snapshot restore: invalid header " - "(snap_ctx=%d cache_ctx=%d snap_layers=%zu " - "cache_layers=%zu pos=%d max_ctx=%d)\n", - snap.ctx != nullptr, cache.ctx != nullptr, - snap.layers.size(), cache.layers.size(), - snap.cur_pos, cache.max_ctx); + snap.layers.size() != cache.layers.size()) { return false; } - if (!tensors_compatible(snap.hc_state_snap, cache.hc_state)) { - std::fprintf(stderr, - "[deepseek4] snapshot restore: incompatible HC state\n"); + if (!tensors_compatible(snap.hc_state_snap, cache.hc_state) || + !copy_tensor_to_backend(snap.hc_state_snap, cache.hc_state)) { return false; } - // Validate the complete layout before changing the live cache. Compressed - // tensors are deliberately right-sized to their logical row counts; - // inactive capacity rows are not part of the snapshot contract. for (size_t il = 0; il < cache.layers.size(); ++il) { const auto & src = snap.layers[il]; - const auto & dst = cache.layers[il]; - const bool raw_ok = tensors_compatible(src.raw_kv, dst.raw_kv); - const bool comp_ok = prefix_tensors_compatible( - src.comp_kv, dst.comp_kv, src.n_comp); - const bool index_ok = prefix_tensors_compatible( - src.index_comp_kv, dst.index_comp_kv, src.n_index_comp); - const bool attn_kv_ok = tensors_compatible( - src.attn_compressor.state_kv, dst.attn_compressor.state_kv); - const bool attn_score_ok = tensors_compatible( - src.attn_compressor.state_score, dst.attn_compressor.state_score); - const bool index_kv_ok = tensors_compatible( - src.indexer_compressor.state_kv, dst.indexer_compressor.state_kv); - const bool index_score_ok = tensors_compatible( - src.indexer_compressor.state_score, - dst.indexer_compressor.state_score); - if (!raw_ok || !comp_ok || !index_ok || !attn_kv_ok || - !attn_score_ok || !index_kv_ok || !index_score_ok) { - std::fprintf(stderr, - "[deepseek4] snapshot restore: incompatible layer %zu " - "(raw=%d comp=%d[%d/%lld/%lld] " - "index=%d[%d/%lld/%lld] states=%d/%d/%d/%d)\n", - il, raw_ok, comp_ok, src.n_comp, - (long long) (src.comp_kv ? src.comp_kv->ne[1] : 0), - (long long) (dst.comp_kv ? dst.comp_kv->ne[1] : 0), - index_ok, src.n_index_comp, - (long long) (src.index_comp_kv - ? src.index_comp_kv->ne[1] : 0), - (long long) (dst.index_comp_kv - ? dst.index_comp_kv->ne[1] : 0), - attn_kv_ok, attn_score_ok, - index_kv_ok, index_score_ok); + 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(snap.hc_state_snap, cache.hc_state)) { - std::fprintf(stderr, - "[deepseek4] snapshot restore: HC copy failed\n"); - return false; - } - for (size_t il = 0; il < cache.layers.size(); ++il) { - const auto & src = snap.layers[il]; - auto & dst = cache.layers[il]; if (!copy_tensor_to_backend(src.raw_kv, dst.raw_kv) || - (src.comp_kv && - !copy_tensor_prefix_to_backend(src.comp_kv, dst.comp_kv, - src.n_comp)) || + (src.comp_kv && !copy_tensor_to_backend(src.comp_kv, dst.comp_kv)) || (src.index_comp_kv && - !copy_tensor_prefix_to_backend(src.index_comp_kv, - dst.index_comp_kv, - src.n_index_comp)) || + !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)) || @@ -8447,10 +9208,7 @@ bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, dst.indexer_compressor.state_kv)) || (src.indexer_compressor.state_score && !copy_tensor_to_backend(src.indexer_compressor.state_score, - dst.indexer_compressor.state_score))) { - std::fprintf(stderr, - "[deepseek4] snapshot restore: layer %zu copy failed\n", - il); + dst.indexer_compressor.state_score))) { return false; } dst.n_comp = src.n_comp; @@ -8458,6 +9216,9 @@ bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, } cache.cur_pos = snap.cur_pos; + if (cache.attention_tp_buf) { + cache.attention_tp_cur_pos = -1; + } return true; } @@ -9310,6 +10071,23 @@ bool deepseek4_dspark_draft_read_async_output( return true; } +bool deepseek4_dspark_draft_device_outputs( + ggml_backend_t backend, + ggml_tensor ** out_hidden, + ggml_tensor ** confidence_hidden) { + DsparkDraftCache & C = g_dspark_draft_cache; + if (!backend || backend != C.backend || !C.ctx || !C.out || + C.block <= 0 || !C.drafter || !out_hidden || + (confidence_hidden && !C.confidence_out)) { + return false; + } + *out_hidden = C.out; + if (confidence_hidden) { + *confidence_hidden = C.confidence_out; + } + return true; +} + void deepseek4_dspark_draft_wait(ggml_backend_t backend) { ggml_backend_synchronize(backend); } diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index b0e80ec00..2c989d97d 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -152,6 +152,16 @@ struct DeepSeek4Layer { // ─── Global weights ───────────────────────────────────────────────────── +// Compact copies of one contiguous attention-head partition. Each output +// group owns n_head / n_out_group complete heads, so keeping these tensors +// together lets the peer execute Q -> attention -> output locally. +struct DeepSeek4AttentionTpLayer { + ggml_tensor * attn_q_b = nullptr; + ggml_tensor * attn_sinks = nullptr; + ggml_tensor * attn_output_a = nullptr; + ggml_tensor * attn_output_b = nullptr; +}; + struct DeepSeek4Weights { ggml_context * ctx = nullptr; ggml_backend_t backend = nullptr; @@ -160,6 +170,15 @@ struct DeepSeek4Weights { // owns per-device allocations while the tensor metadata stays in ctx. ggml_backend_buffer_t dense_split_buf = nullptr; + // Optional persistent attention-head partition on the in-process peer. + // The backend is borrowed; this struct owns the context and weight buffer. + ggml_context * attention_tp_ctx = nullptr; + ggml_backend_buffer_t attention_tp_buf = nullptr; + ggml_backend_t attention_tp_backend = nullptr; + std::vector attention_tp_layers; + int attention_tp_main_groups = 0; + int attention_tp_peer_groups = 0; + // Global tensors ggml_tensor * tok_embd = nullptr; // [n_embd, n_vocab] ggml_tensor * out_norm = nullptr; // [n_embd] @@ -268,6 +287,13 @@ struct DeepSeek4LayerCache { DeepSeek4CompressorState indexer_compressor; }; +// The peer needs only attention-visible KV, not compressor or indexer state: +// the primary computes each new row once and transfers just those new rows. +struct DeepSeek4AttentionTpCacheLayer { + ggml_tensor * raw_kv = nullptr; + ggml_tensor * comp_kv = nullptr; +}; + // Per-shard runtime state for deepseek4_step_layer_range (host-side HC weight // cache + cached decode graphs). Defined in deepseek4_graph.cpp; owned by the // DeepSeek4Cache below and released by free_deepseek4_cache(). @@ -289,6 +315,14 @@ struct DeepSeek4Cache { ggml_context * ctx = nullptr; ggml_backend_buffer_t buf = nullptr; + + ggml_context * attention_tp_ctx = nullptr; + ggml_backend_buffer_t attention_tp_buf = nullptr; + ggml_backend_t attention_tp_backend = nullptr; + std::vector attention_tp_layers; + // Position represented by the peer cache. A mismatch triggers one bulk + // synchronization after prefill, prefix restore, or recovery. + int attention_tp_cur_pos = -1; }; struct DeepSeek4Snapshot; @@ -324,6 +358,11 @@ bool load_deepseek4_gguf_partial(const std::string & path, void free_deepseek4_weights(DeepSeek4Weights & w); +bool init_deepseek4_attention_tp(DeepSeek4Weights & w, + ggml_backend_t peer_backend, + int peer_groups, + std::string * err); + // Release graph allocators and host mirrors that retain model tensor pointers. // This must run before the owning ggml context is destroyed. void deepseek4_release_runtime_graphs(const DeepSeek4Weights & w); @@ -335,6 +374,7 @@ bool create_deepseek4_cache(ggml_backend_t backend, void free_deepseek4_cache(DeepSeek4Cache & c); void reset_deepseek4_cache(DeepSeek4Cache & c); +bool deepseek4_sync_attention_tp_cache(DeepSeek4Cache & c); int deepseek4_previous_raw_ring_spans( int kv_start, int n_swa, @@ -382,10 +422,6 @@ bool deepseek4_step( struct Ds4VerifyHooks { const std::vector * capture_layer_ids = nullptr; // e.g. {40,41,42} std::vector * capture_out = nullptr; // [n_cap*n_embd * n_tokens] - // Optional relative token range for layer-major feature readback. Generic - // verifier paths may ignore this and return the complete batch. - int capture_token_begin = 0; - int capture_token_end = -1; // exclusive; -1 = n_tokens std::vector * all_logits_out = nullptr; // [n_vocab * n_tokens] std::vector * argmax_out = nullptr; // [n_tokens], optional GPU result bool prefer_argmax_only = false; // skip logits D2H when available diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index aa4256430..ee6f6073c 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -1063,14 +1063,203 @@ bool build_deepseek4_moe_hybrid_storage_from_file( path, backend, w, placement, nullptr, out, err); } +bool init_deepseek4_attention_tp( + DeepSeek4Weights & w, + ggml_backend_t peer_backend, + int peer_groups, + std::string * err) { + auto fail = [err](const std::string & message) { + if (err) *err = message; + return false; + }; + if (!peer_backend) return fail("attention TP peer backend is null"); + if (w.attention_tp_ctx || w.attention_tp_buf || + !w.attention_tp_layers.empty()) { + return fail("attention TP is already initialized"); + } + if (w.n_out_group <= 1 || peer_groups <= 0 || + peer_groups >= w.n_out_group) { + return fail("attention TP peer groups must leave at least one group " + "on each owner"); + } + if (w.n_head <= 0 || w.n_head % w.n_out_group != 0 || + w.head_dim <= 0 || w.n_lora_q <= 0 || w.n_lora_o <= 0 || + w.n_embd <= 0 || (int) w.layers.size() != w.n_layer) { + return fail("attention TP model dimensions are inconsistent"); + } + + const int main_groups = w.n_out_group - peer_groups; + const int heads_per_group = w.n_head / w.n_out_group; + const int64_t main_heads = (int64_t) main_groups * heads_per_group; + const int64_t peer_heads = (int64_t) peer_groups * heads_per_group; + const int64_t group_dim = (int64_t) w.head_dim * heads_per_group; + const int64_t main_head_dim = main_heads * w.head_dim; + const int64_t peer_head_dim = peer_heads * w.head_dim; + const int64_t main_lora = (int64_t) main_groups * w.n_lora_o; + const int64_t peer_lora = (int64_t) peer_groups * w.n_lora_o; + + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & layer = w.layers[(size_t) il]; + const ggml_tensor * q = layer.attn_q_b; + const ggml_tensor * a = layer.attn_output_a; + const ggml_tensor * b = layer.attn_output_b; + if (!q || !a || !b || + q->ne[0] != w.n_lora_q || + q->ne[1] != (int64_t) w.n_head * w.head_dim || + a->ne[0] != group_dim || + a->ne[1] != (int64_t) w.n_lora_o * w.n_out_group || + b->ne[0] != (int64_t) w.n_lora_o * w.n_out_group || + b->ne[1] != w.n_embd) { + return fail("unexpected attention TP tensor shape at layer " + + std::to_string(il)); + } + if (layer.attn_sinks && + ggml_nelements(layer.attn_sinks) != w.n_head) { + return fail("unexpected attention sink shape at layer " + + std::to_string(il)); + } + if (!ggml_is_contiguous(q) || !ggml_is_contiguous(a) || + !ggml_is_contiguous(b) || + w.n_lora_q % ggml_blck_size(q->type) != 0 || + group_dim % ggml_blck_size(a->type) != 0 || + peer_lora % ggml_blck_size(b->type) != 0 || + main_lora % ggml_blck_size(b->type) != 0) { + return fail("unsupported attention TP tensor layout at layer " + + std::to_string(il)); + } + } + + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * + (size_t) (4 * w.n_layer + 8) + 4096; + params.no_alloc = true; + ggml_context * peer_ctx = ggml_init(params); + if (!peer_ctx) { + return fail("failed to allocate attention TP metadata"); + } + + std::vector peer_layers((size_t) w.n_layer); + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & source = w.layers[(size_t) il]; + DeepSeek4AttentionTpLayer & destination = peer_layers[(size_t) il]; + destination.attn_q_b = ggml_new_tensor_2d( + peer_ctx, source.attn_q_b->type, w.n_lora_q, peer_head_dim); + if (source.attn_sinks) { + destination.attn_sinks = ggml_new_tensor_1d( + peer_ctx, source.attn_sinks->type, peer_heads); + } + destination.attn_output_a = ggml_new_tensor_2d( + peer_ctx, source.attn_output_a->type, group_dim, peer_lora); + destination.attn_output_b = ggml_new_tensor_2d( + peer_ctx, source.attn_output_b->type, peer_lora, w.n_embd); + ggml_format_name(destination.attn_q_b, + "blk.%d.attn_q_b.peer", il); + if (destination.attn_sinks) { + ggml_format_name(destination.attn_sinks, + "blk.%d.attn_sinks.peer", il); + } + ggml_format_name(destination.attn_output_a, + "blk.%d.attn_output_a.peer", il); + ggml_format_name(destination.attn_output_b, + "blk.%d.attn_output_b.peer", il); + } + + ggml_backend_buffer_t peer_buf = + ggml_backend_alloc_ctx_tensors(peer_ctx, peer_backend); + if (!peer_buf) { + ggml_free(peer_ctx); + return fail("failed to allocate attention TP peer weights"); + } + ggml_backend_buffer_set_usage(peer_buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + std::vector staging; + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & source = w.layers[(size_t) il]; + const DeepSeek4AttentionTpLayer & destination = + peer_layers[(size_t) il]; + + size_t bytes = ggml_nbytes(destination.attn_q_b); + staging.resize(bytes); + ggml_backend_tensor_get( + source.attn_q_b, staging.data(), + (size_t) main_head_dim * source.attn_q_b->nb[1], bytes); + ggml_backend_tensor_set(destination.attn_q_b, staging.data(), 0, bytes); + + if (source.attn_sinks) { + bytes = ggml_nbytes(destination.attn_sinks); + staging.resize(bytes); + ggml_backend_tensor_get( + source.attn_sinks, staging.data(), + (size_t) main_heads * source.attn_sinks->nb[0], bytes); + ggml_backend_tensor_set( + destination.attn_sinks, staging.data(), 0, bytes); + } + + bytes = ggml_nbytes(destination.attn_output_a); + staging.resize(bytes); + ggml_backend_tensor_get( + source.attn_output_a, staging.data(), + (size_t) main_lora * source.attn_output_a->nb[1], bytes); + ggml_backend_tensor_set( + destination.attn_output_a, staging.data(), 0, bytes); + + const size_t main_row_bytes = + ggml_row_size(source.attn_output_b->type, main_lora); + const size_t peer_row_bytes = + ggml_row_size(source.attn_output_b->type, peer_lora); + bytes = ggml_nbytes(destination.attn_output_b); + if (main_row_bytes + peer_row_bytes > source.attn_output_b->nb[1] || + peer_row_bytes != destination.attn_output_b->nb[1]) { + ggml_backend_buffer_free(peer_buf); + ggml_free(peer_ctx); + return fail("attention TP row packing mismatch at layer " + + std::to_string(il)); + } + staging.resize(bytes); + ggml_backend_tensor_get_2d( + source.attn_output_b, staging.data(), main_row_bytes, + peer_row_bytes, (size_t) w.n_embd, + source.attn_output_b->nb[1], + destination.attn_output_b->nb[1]); + ggml_backend_tensor_set( + destination.attn_output_b, staging.data(), 0, bytes); + } + + w.attention_tp_ctx = peer_ctx; + w.attention_tp_buf = peer_buf; + w.attention_tp_backend = peer_backend; + w.attention_tp_layers = std::move(peer_layers); + w.attention_tp_main_groups = main_groups; + w.attention_tp_peer_groups = peer_groups; + std::fprintf(stderr, + "[deepseek4-attention-tp] groups=%d+%d heads=%lld+%lld " + "peer_weights=%.1f MiB\n", + main_groups, peer_groups, + (long long) main_heads, (long long) peer_heads, + ggml_backend_buffer_get_size(peer_buf) / 1024.0 / 1024.0); + return true; +} + void free_deepseek4_weights(DeepSeek4Weights & w) { deepseek4_release_runtime_graphs(w); - if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } + if (w.attention_tp_buf) { + ggml_backend_buffer_free(w.attention_tp_buf); + w.attention_tp_buf = nullptr; + } + if (w.attention_tp_ctx) { + ggml_free(w.attention_tp_ctx); + w.attention_tp_ctx = nullptr; + } + w.attention_tp_backend = nullptr; + w.attention_tp_layers.clear(); + w.attention_tp_main_groups = 0; + w.attention_tp_peer_groups = 0; if (w.dense_split_buf) { ggml_backend_buffer_free(w.dense_split_buf); w.dense_split_buf = nullptr; } if (w.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; } + if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } w.layers.clear(); w.embedder.tok_embd_owned.clear(); w.embedder.tok_embd_bytes = nullptr; diff --git a/server/test/test_deepseek4_mmid_grouped_cuda.cpp b/server/test/test_deepseek4_mmid_grouped_cuda.cpp index 975b772fd..019c342c1 100644 --- a/server/test/test_deepseek4_mmid_grouped_cuda.cpp +++ b/server/test/test_deepseek4_mmid_grouped_cuda.cpp @@ -2,6 +2,7 @@ #include "ggml-backend.h" #include "ggml-cuda.h" #include "ggml.h" +#include "rocmfpx.h" #include @@ -81,8 +82,20 @@ static bool run_case( } std::vector weights_q(ggml_nbytes(weights)); - const size_t quantized = ggml_quantize_chunk( - type, weights_f.data(), weights_q.data(), 0, n_rows * n_experts, k_dim, nullptr); + size_t quantized = 0; + if (type == GGML_TYPE_Q2_0_ROCMFP2) { + quantized = rocmfpx_quantize_fp2( + weights_f.data(), weights_q.data(), n_rows*n_experts, k_dim, + nullptr); + } else if (type == GGML_TYPE_Q3_0_ROCMFPX) { + quantized = rocmfpx_quantize_fp3( + weights_f.data(), weights_q.data(), n_rows*n_experts, k_dim, + nullptr); + } else { + quantized = ggml_quantize_chunk( + type, weights_f.data(), weights_q.data(), 0, + n_rows*n_experts, k_dim, nullptr); + } if (quantized != weights_q.size()) { std::fprintf(stderr, "quantize size mismatch type=%s got=%zu expected=%zu\n", ggml_type_name(type), quantized, weights_q.size()); @@ -153,8 +166,9 @@ static int run_child(const char * mode, const char * output_path) { std::ofstream output(output_path, std::ios::binary | std::ios::trunc); const ggml_type types[] = { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q5_K, + GGML_TYPE_Q2_0_ROCMFP2, GGML_TYPE_Q3_0_ROCMFPX, }; - const int widths[] = {2, 4, 8, 9, 16}; + const int widths[] = {2, 4, 5, 8, 9, 16}; bool ok = output.good(); for (ggml_type type : types) { for (int width : widths) { @@ -294,13 +308,20 @@ static std::string child_command( const std::string & output_path, const std::string & log_path) { #if defined(_WIN32) - return "set \"DFLASH_MMID_TELEMETRY=1\" && set \"DFLASH_MMID_GROUPED_TYPES=7\" && " + return "set \"DFLASH_MMID_TELEMETRY=1\" && set \"DFLASH_MMID_GROUPED_TYPES=15\" && " + "set \"DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24=1\" && " + "set \"DFLASH_CUDA_MMVQ_MOE_GROUP_REUSE=" + + std::string(std::strcmp(mode, "grouped") == 0 ? "1" : "0") + "\" && " "set \"DFLASH_MMID_GROUPED=" + std::string(std::strcmp(mode, "grouped") == 0 ? "1" : "0") + "\" && " + shell_quote(executable) + " --child " + mode + " " + shell_quote(output_path) + " 2>" + shell_quote(log_path); #else - return "DFLASH_MMID_TELEMETRY=1 DFLASH_MMID_GROUPED_TYPES=7 DFLASH_MMID_GROUPED=" + + return "DFLASH_MMID_TELEMETRY=1 DFLASH_MMID_GROUPED_TYPES=15 " + "DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24=1 " + "DFLASH_CUDA_MMVQ_MOE_GROUP_REUSE=" + + std::string(std::strcmp(mode, "grouped") == 0 ? "1" : "0") + + " DFLASH_MMID_GROUPED=" + std::string(std::strcmp(mode, "grouped") == 0 ? "1" : "0") + " " + shell_quote(executable) + " --child " + mode + " " + shell_quote(output_path) + " 2>" + shell_quote(log_path); @@ -345,26 +366,36 @@ int main(int argc, char ** argv) { const std::vector grouped_log = read_file(grouped_log_path.c_str()); const size_t legacy_grouped = count_records(legacy_log, "variant=grouped"); const size_t grouped_grouped = count_records(grouped_log, "variant=grouped"); + const size_t grouped_reuse = count_records(grouped_log, "variant=group-reuse"); const ggml_type types[] = { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q5_K, + GGML_TYPE_Q2_0_ROCMFP2, GGML_TYPE_Q3_0_ROCMFPX, }; - const int widths[] = {2, 4, 8, 9, 16}; + const int widths[] = {2, 4, 5, 8, 9, 16}; size_t offset = 0; size_t compared_bytes = 0; int compared_cases = 0; int exact_cases = 0; int tolerant_cases = 0; + size_t expected_grouped_records = 0; + size_t expected_reuse_records = 0; bool output_parity = legacy.size() == grouped.size() && !legacy.empty(); bool grouped_dispatch = true; for (ggml_type type : types) { for (int width : widths) { const bool legacy_mmvq = has_mmvq_record(legacy_log, type, width); + const bool expect_reuse = + (type == GGML_TYPE_Q2_0_ROCMFP2 || + type == GGML_TYPE_Q3_0_ROCMFPX) && width <= 5; for (bool fused_ds4 : {false, true}) { const size_t case_bytes = (size_t) 128 * 8 * width * sizeof(float); const bool require_exact = legacy_mmvq && !fused_ds4; grouped_dispatch = - has_mmvq_record(grouped_log, type, width, "grouped") && grouped_dispatch; + has_mmvq_record( + grouped_log, type, width, + expect_reuse ? "group-reuse" : "grouped") && + grouped_dispatch; if (output_parity) { output_parity = compare_case_outputs( legacy, grouped, offset, case_bytes, require_exact); @@ -381,11 +412,16 @@ int main(int argc, char ** argv) { tolerant_cases += require_exact ? 0 : 1; offset += case_bytes; } + // One non-fused MUL_MAT_ID plus two inputs to the fused DS4 graph. + (expect_reuse ? expected_reuse_records : + expected_grouped_records) += 3; } } - output_parity = output_parity && offset == legacy.size() && compared_cases == 50; + output_parity = output_parity && offset == legacy.size() && compared_cases == 84; const bool pass = legacy_status == 0 && grouped_status == 0 && - output_parity && grouped_dispatch && legacy_grouped == 0 && grouped_grouped == 75; + output_parity && grouped_dispatch && legacy_grouped == 0 && + grouped_grouped == expected_grouped_records && + grouped_reuse == expected_reuse_records; if (pass) { std::remove(legacy_path.c_str()); std::remove(grouped_path.c_str()); @@ -394,10 +430,11 @@ int main(int argc, char ** argv) { } std::printf("[mmid-grouped-test] legacy_status=%d grouped_status=%d bytes=%zu " "compared_cases=%d exact_cases=%d tolerant_cases=%d compared_bytes=%zu " - "legacy_grouped=%zu grouped_grouped=%zu " + "legacy_grouped=%zu grouped_grouped=%zu grouped_reuse=%zu " "parity=%s\n", legacy_status, grouped_status, legacy.size(), compared_cases, exact_cases, tolerant_cases, compared_bytes, legacy_grouped, grouped_grouped, + grouped_reuse, pass ? "PASS" : "FAIL"); return pass ? 0 : 1; } diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index bea27f0a8..684cec545 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -843,6 +843,106 @@ static void test_grouped_output_projection_cpu(ggml_backend_t backend) { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_attention_tp_weight_partition_cpu(ggml_backend_t backend) { + std::fprintf(stderr, " test_attention_tp_weight_partition_cpu ..."); + + constexpr int n_groups = 4; + constexpr int peer_groups = 2; + constexpr int heads_per_group = 2; + constexpr int head_dim = 2; + constexpr int n_lora_q = 3; + constexpr int n_lora_o = 3; + constexpr int n_embd = 5; + constexpr int n_head = n_groups * heads_per_group; + constexpr int group_dim = heads_per_group * head_dim; + constexpr int main_heads = (n_groups - peer_groups) * heads_per_group; + constexpr int peer_heads = peer_groups * heads_per_group; + constexpr int main_lora = (n_groups - peer_groups) * n_lora_o; + constexpr int peer_lora = peer_groups * n_lora_o; + + DeepSeek4Weights weights; + weights.backend = backend; + weights.n_layer = 1; + weights.n_head = n_head; + weights.head_dim = head_dim; + weights.n_out_group = n_groups; + weights.n_lora_q = n_lora_q; + weights.n_lora_o = n_lora_o; + weights.n_embd = n_embd; + weights.layers.resize(1); + weights.ctx = make_test_context(1u << 16); + TEST_ASSERT_MSG(weights.ctx != nullptr, "weight metadata allocation failed"); + if (!weights.ctx) return; + + DeepSeek4Layer & layer = weights.layers[0]; + layer.attn_q_b = ggml_new_tensor_2d( + weights.ctx, GGML_TYPE_F32, n_lora_q, n_head * head_dim); + layer.attn_sinks = ggml_new_tensor_1d( + weights.ctx, GGML_TYPE_F32, n_head); + layer.attn_output_a = ggml_new_tensor_2d( + weights.ctx, GGML_TYPE_F32, group_dim, n_groups * n_lora_o); + layer.attn_output_b = ggml_new_tensor_2d( + weights.ctx, GGML_TYPE_F32, n_groups * n_lora_o, n_embd); + weights.buf = ggml_backend_alloc_ctx_tensors(weights.ctx, backend); + TEST_ASSERT_MSG(weights.buf != nullptr, "weight buffer allocation failed"); + if (!weights.buf) { + free_deepseek4_weights(weights); + return; + } + + std::vector q((size_t) n_lora_q * n_head * head_dim); + std::vector sinks((size_t) n_head); + std::vector a((size_t) group_dim * n_groups * n_lora_o); + std::vector b((size_t) n_groups * n_lora_o * n_embd); + for (size_t i = 0; i < q.size(); ++i) q[i] = (float) i; + for (size_t i = 0; i < sinks.size(); ++i) sinks[i] = 1000.0f + (float) i; + for (size_t i = 0; i < a.size(); ++i) a[i] = 2000.0f + (float) i; + for (size_t i = 0; i < b.size(); ++i) b[i] = 3000.0f + (float) i; + ggml_backend_tensor_set(layer.attn_q_b, q.data(), 0, q.size() * sizeof(float)); + ggml_backend_tensor_set(layer.attn_sinks, sinks.data(), 0, + sinks.size() * sizeof(float)); + ggml_backend_tensor_set(layer.attn_output_a, a.data(), 0, + a.size() * sizeof(float)); + ggml_backend_tensor_set(layer.attn_output_b, b.data(), 0, + b.size() * sizeof(float)); + + std::string err; + TEST_ASSERT_MSG(init_deepseek4_attention_tp( + weights, backend, peer_groups, &err), + err.c_str()); + if (!weights.attention_tp_layers.empty()) { + const DeepSeek4AttentionTpLayer & peer = + weights.attention_tp_layers[0]; + std::vector peer_q((size_t) n_lora_q * peer_heads * head_dim); + std::vector peer_sinks((size_t) peer_heads); + std::vector peer_a((size_t) group_dim * peer_lora); + std::vector peer_b((size_t) peer_lora * n_embd); + ggml_backend_tensor_get(peer.attn_q_b, peer_q.data(), 0, + peer_q.size() * sizeof(float)); + ggml_backend_tensor_get(peer.attn_sinks, peer_sinks.data(), 0, + peer_sinks.size() * sizeof(float)); + ggml_backend_tensor_get(peer.attn_output_a, peer_a.data(), 0, + peer_a.size() * sizeof(float)); + ggml_backend_tensor_get(peer.attn_output_b, peer_b.data(), 0, + peer_b.size() * sizeof(float)); + const size_t q_offset = (size_t) main_heads * head_dim * n_lora_q; + const size_t a_offset = (size_t) main_lora * group_dim; + TEST_ASSERT(std::equal(peer_q.begin(), peer_q.end(), q.begin() + q_offset)); + TEST_ASSERT(std::equal(peer_sinks.begin(), peer_sinks.end(), + sinks.begin() + main_heads)); + TEST_ASSERT(std::equal(peer_a.begin(), peer_a.end(), a.begin() + a_offset)); + for (int row = 0; row < n_embd; ++row) { + for (int col = 0; col < peer_lora; ++col) { + TEST_ASSERT(peer_b[(size_t) row * peer_lora + col] == + b[(size_t) row * n_groups * n_lora_o + + main_lora + col]); + } + } + } + free_deepseek4_weights(weights); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_ds4_flash_attention_cpu_rejected(ggml_backend_t backend) { std::fprintf(stderr, " test_ds4_flash_attention_cpu_rejected ..."); ggml_context * ctx = make_test_context(); @@ -1206,15 +1306,6 @@ static std::vector read_tensor_bytes(const ggml_tensor * tensor) { return data; } -static std::vector read_tensor_rows(const ggml_tensor * tensor, - int rows) { - const size_t bytes = rows > 0 - ? ggml_row_size(tensor->type, tensor->ne[0]) * (size_t) rows : 0; - std::vector data(bytes); - if (bytes > 0) ggml_backend_tensor_get(tensor, data.data(), 0, bytes); - return data; -} - static bool init_snapshot_test_shard(DeepSeek4LayerSplitAdapter & adapter) { adapter.shards_.resize(1); auto & shard = adapter.shards_[0]; @@ -1790,8 +1881,8 @@ static void test_snapshot_save_restore() { write_tensor_pattern(layer.indexer_compressor.state_score, 79); const std::vector raw_before = read_tensor_bytes(layer.raw_kv); - const std::vector comp_before = read_tensor_rows(layer.comp_kv, 5); - const std::vector index_before = read_tensor_rows(layer.index_comp_kv, 3); + const std::vector comp_before = read_tensor_bytes(layer.comp_kv); + const std::vector index_before = read_tensor_bytes(layer.index_comp_kv); const std::vector attn_kv_before = read_tensor_bytes(layer.attn_compressor.state_kv); const std::vector attn_score_before = @@ -1812,8 +1903,6 @@ static void test_snapshot_save_restore() { TEST_ASSERT(adapter.snapshot_save(0)); TEST_ASSERT(adapter.snapshot_used(0)); TEST_ASSERT(adapter.snapshot_cur_pos(0) == 7); - TEST_ASSERT(adapter.snapshots_[0].shards[0].layers[0].comp_kv->ne[1] == 5); - TEST_ASSERT(adapter.snapshots_[0].shards[0].layers[0].index_comp_kv->ne[1] == 3); adapter.cur_pos_ = 0; adapter.last_tok_ = -1; @@ -1841,31 +1930,13 @@ static void test_snapshot_save_restore() { TEST_ASSERT(layer.n_comp == 5); TEST_ASSERT(layer.n_index_comp == 3); TEST_ASSERT(read_tensor_bytes(layer.raw_kv) == raw_before); - TEST_ASSERT(read_tensor_rows(layer.comp_kv, 5) == comp_before); - TEST_ASSERT(read_tensor_rows(layer.index_comp_kv, 3) == index_before); + TEST_ASSERT(read_tensor_bytes(layer.comp_kv) == comp_before); + TEST_ASSERT(read_tensor_bytes(layer.index_comp_kv) == index_before); TEST_ASSERT(read_tensor_bytes(layer.attn_compressor.state_kv) == attn_kv_before); TEST_ASSERT(read_tensor_bytes(layer.attn_compressor.state_score) == attn_score_before); TEST_ASSERT(read_tensor_bytes(layer.indexer_compressor.state_kv) == index_kv_before); TEST_ASSERT(read_tensor_bytes(layer.indexer_compressor.state_score) == index_score_before); - // A zero-row compressed prefix is represented by one physical GGML row. - // ggml_n_dims() reports that tensor as 1D, so restore must validate its - // row layout rather than reject it against the full 2D cache capacity. - adapter.snapshot_free(0); - layer.n_comp = 0; - layer.n_index_comp = 0; - cache.cur_pos = 1; - adapter.cur_pos_ = 1; - adapter.last_tok_ = 7; - TEST_ASSERT(adapter.snapshot_save(0)); - TEST_ASSERT(adapter.snapshots_[0].shards[0].layers[0].comp_kv->ne[1] == 1); - TEST_ASSERT(adapter.snapshots_[0].shards[0].layers[0].index_comp_kv->ne[1] == 1); - cache.cur_pos = 0; - TEST_ASSERT(adapter.snapshot_restore(0)); - TEST_ASSERT(cache.cur_pos == 1); - TEST_ASSERT(layer.n_comp == 0); - TEST_ASSERT(layer.n_index_comp == 0); - adapter.snapshot_free(0); TEST_ASSERT(!adapter.snapshot_used(0)); TEST_ASSERT(adapter.snapshots_[0].hc_state.empty()); @@ -1879,172 +1950,6 @@ static void test_snapshot_save_restore() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } -static bool init_monolithic_snapshot_test_backend(DeepSeek4Backend & backend) { - backend.backend_ = ggml_backend_cpu_init(); - backend.snap_backend_ = ggml_backend_cpu_init(); - if (!backend.backend_ || !backend.snap_backend_) return false; - backend.w_.n_layer = 1; - backend.w_.n_embd = 4; - backend.w_.n_hc = 1; - backend.w_.n_vocab = 3; - backend.w_.head_dim = 4; - backend.w_.n_swa = 8; - backend.w_.n_indexer_head_dim = 2; - backend.w_.compress_ratios = {4}; - return create_deepseek4_cache(backend.backend_, backend.w_, 16, - backend.cache_); -} - -static void test_monolithic_snapshot_preserves_decode_state() { - std::fprintf(stderr, - " test_monolithic_snapshot_preserves_decode_state ..."); - - DeepSeek4BackendConfig cfg; - DeepSeek4Backend backend(cfg); - TEST_ASSERT(init_monolithic_snapshot_test_backend(backend)); - if (!backend.cache_.buf || backend.cache_.layers.empty()) { - std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); - return; - } - - auto & layer = backend.cache_.layers[0]; - write_tensor_pattern(layer.raw_kv, 13); - write_tensor_pattern(layer.comp_kv, 31); - write_tensor_pattern(layer.index_comp_kv, 47); - write_tensor_pattern(layer.attn_compressor.state_kv, 59); - write_tensor_pattern(layer.attn_compressor.state_score, 71); - write_tensor_pattern(layer.indexer_compressor.state_kv, 83); - write_tensor_pattern(layer.indexer_compressor.state_score, 97); - write_tensor_pattern(backend.cache_.hc_state, 109); - - layer.n_comp = 5; - layer.n_index_comp = 3; - backend.cache_.cur_pos = 7; - backend.last_logits_ = {1.0f, 4.0f, 2.0f}; - backend.last_logits_pos_ = 7; - backend.spec_feat_window_ = {9.0f, 8.0f, 7.0f, 6.0f}; - - const auto raw_before = read_tensor_bytes(layer.raw_kv); - const auto comp_before = read_tensor_rows(layer.comp_kv, layer.n_comp); - const auto index_before = - read_tensor_rows(layer.index_comp_kv, layer.n_index_comp); - const auto hc_before = read_tensor_bytes(backend.cache_.hc_state); - - TEST_ASSERT(backend.snapshot_save(0)); - TEST_ASSERT(backend.snapshot_used(0)); - TEST_ASSERT(backend.snapshot_cur_pos(0) == 7); - TEST_ASSERT(backend.snapshots_[0].layers[0].comp_kv->ne[1] == 5); - TEST_ASSERT(backend.snapshots_[0].layers[0].index_comp_kv->ne[1] == 3); - - ggml_backend_buffer_clear(backend.cache_.buf, 0); - backend.cache_.cur_pos = 0; - layer.n_comp = 0; - layer.n_index_comp = 0; - backend.last_logits_ = {-1.0f}; - backend.last_logits_pos_ = -1; - backend.spec_feat_window_.clear(); - - TEST_ASSERT(backend.snapshot_restore(0)); - TEST_ASSERT(backend.cache_.cur_pos == 7); - TEST_ASSERT(layer.n_comp == 5); - TEST_ASSERT(layer.n_index_comp == 3); - TEST_ASSERT(read_tensor_bytes(layer.raw_kv) == raw_before); - TEST_ASSERT(read_tensor_rows(layer.comp_kv, 5) == comp_before); - TEST_ASSERT(read_tensor_rows(layer.index_comp_kv, 3) == index_before); - TEST_ASSERT(read_tensor_bytes(backend.cache_.hc_state) == hc_before); - TEST_ASSERT(backend.last_logits_ == std::vector({1.0f, 4.0f, 2.0f})); - TEST_ASSERT(backend.spec_feat_window_ == - std::vector({9.0f, 8.0f, 7.0f, 6.0f})); - TEST_ASSERT(backend.last_logits_pos_ == 7); - - GenerateRequest exact; - exact.prompt.assign(7, 0); - exact.n_gen = 1; - const GenerateResult exact_result = - backend.restore_and_generate_impl(0, exact, DaemonIO{}); - TEST_ASSERT(exact_result.ok()); - TEST_ASSERT(exact_result.tokens == std::vector({1})); - TEST_ASSERT(backend.cache_.cur_pos == 7); - TEST_ASSERT(backend.last_logits_ == std::vector({1.0f, 4.0f, 2.0f})); - - // A cache state that advanced without corresponding logits (for example, - // after DSpark) must not be persisted with stale decode state. - backend.cache_.cur_pos = 8; - TEST_ASSERT(!backend.snapshot_save(1)); - backend.cache_.cur_pos = 7; - TEST_ASSERT(!backend.snapshot_save(-1)); - TEST_ASSERT(!backend.snapshot_save(ModelBackend::kMaxSlots)); - - // Parking the target releases both the core tensors and the potentially - // large host-side logits/feature vectors for every populated slot. - TEST_ASSERT(backend.park(ParkTarget::TargetModel)); - TEST_ASSERT(!backend.snapshot_used(0)); - TEST_ASSERT(backend.snapshot_aux_[0].last_logits.empty()); - TEST_ASSERT(backend.snapshot_aux_[0].spec_feat_window.empty()); - TEST_ASSERT(!backend.snapshot_restore(0)); - - std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); -} - -static void test_spec_feature_tail_is_bounded() { - std::fprintf(stderr, " test_spec_feature_tail_is_bounded ..."); - - DeepSeek4BackendConfig cfg; - DeepSeek4Backend backend(cfg); - backend.w_.n_embd = 2; - backend.w_.n_swa = 3; - backend.spec_drafter_ = std::make_unique(); - backend.spec_drafter_->n_target_layers = 1; - - std::vector features = { - 0.0f, 1.0f, - 2.0f, 3.0f, - 4.0f, 5.0f, - 6.0f, 7.0f, - 8.0f, 9.0f, - }; - backend.keep_spec_feature_tail(features, 3); - TEST_ASSERT(features == - std::vector({4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f})); - - features.push_back(10.0f); // malformed partial feature row - backend.keep_spec_feature_tail(features, 3); - TEST_ASSERT(features.empty()); - - std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); -} - -static void test_dspark_prefill_capture_boundaries() { - std::fprintf(stderr, " test_dspark_prefill_capture_boundaries ..."); - - using Backend = DeepSeek4Backend; - // Layer-major prefill captures only the requested tail from a wide graph, - // while generic paths still stop exactly at the final feature window. - TEST_ASSERT(Backend::capture_safe_prefill_tokens( - 0, 2048, 1920, true, false, 0, 0) == 2048); - TEST_ASSERT(Backend::capture_safe_prefill_tokens( - 0, 2048, 1920, false, false, 0, 0) == 1920); - TEST_ASSERT(Backend::capture_safe_prefill_tokens( - 1920, 128, 1920, false, false, 0, 0) == 128); - - // A pending checkpoint contributes both edges of its capture window. The - // resulting batches are either wholly hooked or wholly unhooked. - TEST_ASSERT(Backend::capture_safe_prefill_tokens( - 0, 2048, 1920, true, true, 384, 512) == 384); - TEST_ASSERT(Backend::capture_safe_prefill_tokens( - 384, 1664, 1920, true, true, 384, 512) == 128); - TEST_ASSERT(Backend::capture_safe_prefill_tokens( - 512, 1536, 1920, false, false, 384, 512) == 1408); - - // Boundaries at a batch edge and empty requests need no extra split. - TEST_ASSERT(Backend::capture_safe_prefill_tokens( - 0, 128, 128, false, true, 128, 256) == 128); - TEST_ASSERT(Backend::capture_safe_prefill_tokens( - 10, 0, 20, false, true, 12, 18) == 0); - - std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); -} - static void test_reset_request_state() { std::fprintf(stderr, " test_reset_request_state ..."); @@ -3150,6 +3055,121 @@ static void test_ds4_flash_attention_inverse_rope_fallback_gpu() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_tail_inverse_rope_matches_across_gpus() { + std::fprintf(stderr, + " test_tail_inverse_rope_matches_across_gpus ..."); + if (ggml_backend_cuda_get_device_count() < 2) { + std::fprintf(stderr, " skipped (requires two GPUs)\n"); + return; + } + + constexpr int head_dim = 512; + constexpr int n_rot = 64; + constexpr int n_head = 16; + constexpr int n_tokens = 5; + constexpr int kv_start = 2048; + const size_t count = + (size_t) head_dim * n_head * n_tokens; + std::vector input(count); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = ((int) (i % 257) - 128) * 0.003125f; + } + std::vector positions((size_t) 2 * n_tokens); + for (int i = 0; i < n_tokens; ++i) { + positions[(size_t) i] = kv_start + i; + positions[(size_t) n_tokens + i] = -(kv_start + i); + } + + auto run = [&](int device, bool use_positive_rope_back) { + std::vector result(count, 0.0f); + ggml_backend_t backend = ggml_backend_cuda_init(device); + TEST_ASSERT_MSG(backend != nullptr, + "failed to initialize inverse-RoPE GPU"); + if (!backend) return result; + ggml_context * ctx = make_test_context(8u << 20); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + ggml_backend_free(backend); + return result; + } + + ggml_tensor * x = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, head_dim, n_head, n_tokens); + ggml_tensor * pos_bundle = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, 2 * n_tokens); + ggml_set_input(x); + ggml_set_input(pos_bundle); + ggml_tensor * rope_pos = ggml_view_1d( + ctx, pos_bundle, n_tokens, + use_positive_rope_back + ? 0 + : (size_t) n_tokens * sizeof(int32_t)); + + const int nope_dim = head_dim - n_rot; + ggml_tensor * nope = ggml_view_3d( + ctx, x, nope_dim, n_head, n_tokens, + x->nb[1], x->nb[2], 0); + ggml_tensor * tail = ggml_view_3d( + ctx, x, n_rot, n_head, n_tokens, + x->nb[1], x->nb[2], + (size_t) nope_dim * x->nb[0]); + tail = ggml_cont(ctx, tail); + tail = use_positive_rope_back + ? ggml_rope_ext_back( + ctx, tail, rope_pos, nullptr, + n_rot, GGML_ROPE_TYPE_NORMAL, 65536, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f) + : ggml_rope_ext( + ctx, tail, rope_pos, nullptr, + n_rot, GGML_ROPE_TYPE_NORMAL, 65536, + 10000.0f, 1.0f, 0.0f, 1.0f, 32.0f, 1.0f); + ggml_tensor * output = ggml_cont( + ctx, ggml_concat(ctx, ggml_cont(ctx, nope), tail, 0)); + ggml_set_output(output); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + const bool allocated = + alloc && ggml_gallocr_alloc_graph(alloc, graph); + TEST_ASSERT_MSG(allocated, + "inverse-RoPE graph allocation failed"); + if (allocated) { + ggml_backend_tensor_set( + x, input.data(), 0, input.size() * sizeof(float)); + ggml_backend_tensor_set( + pos_bundle, positions.data(), 0, + positions.size() * sizeof(int32_t)); + const bool computed = + ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS; + TEST_ASSERT_MSG(computed, + "inverse-RoPE graph compute failed"); + if (computed) { + ggml_backend_tensor_get( + output, result.data(), 0, + result.size() * sizeof(float)); + } + } + if (alloc) ggml_gallocr_free(alloc); + ggml_free(ctx); + ggml_backend_free(backend); + return result; + }; + + const std::vector main = run(0, false); + const std::vector peer = run(1, true); + float max_abs = 0.0f; + for (size_t i = 0; i < count; ++i) { + max_abs = std::max(max_abs, std::fabs(main[i] - peer[i])); + TEST_ASSERT_MSG( + nearly_equal(main[i], peer[i], 2.0e-5f, 2.0e-5f), + "inverse tail RoPE differs across GPUs"); + } + std::fprintf(stderr, " max_abs=%.3g", max_abs); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_hc_post_strided_split_gpu() { std::fprintf(stderr, " test_hc_post_strided_split_gpu ..."); ggml_backend_t backend = ggml_backend_cuda_init(0); @@ -4108,6 +4128,7 @@ int main() { test_rmsnorm_correctness(backend); test_grouped_output_projection_shape(); test_grouped_output_projection_cpu(backend); + test_attention_tp_weight_partition_cpu(backend); test_ds4_flash_attention_cpu_rejected(backend); test_indexer_qat_cpu(backend); test_indexer_score_cpu(backend); @@ -4134,9 +4155,6 @@ int main() { test_dspark_park_all_releases_drafter(); test_dspark_raw_ring_rollback_after_wrap(backend); test_snapshot_save_restore(); - test_monolithic_snapshot_preserves_decode_state(); - test_spec_feature_tail_is_bounded(); - test_dspark_prefill_capture_boundaries(); test_reset_request_state(); test_reset_deepseek4_cache(backend); test_adapter_guard_paths(); @@ -4150,6 +4168,7 @@ int main() { test_ds4_indexer_score_packed_q4_gpu(); test_ds4_topk_block_radix_gpu(); test_ds4_flash_attention_inverse_rope_fallback_gpu(); + test_tail_inverse_rope_matches_across_gpus(); test_hc_post_strided_split_gpu(); test_moe_id_alignment_q5_gpu(); test_hc_pre_kernel_gpu();