diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 09c88513c..486288877 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -177,7 +177,9 @@ jobs: -DCMAKE_BUILD_TYPE=Release cmake --build build \ --target test_flash_attn_sparse test_deepseek4_mmid_grouped_cuda \ - test_deepseek4_unit \ + test_deepseek4_unit test_rocmfp3_mix_registry \ + test_rocmfp_mix_slice_matvec test_rocmfp_mix_gateup_glu \ + test_ds4_mix_registry_teardown \ -j"$(nproc)" - name: Run flash-attn sparse kernel test on the 3090 @@ -191,6 +193,11 @@ jobs: - name: Run DeepSeek4 graph-generation regression on the 3090 run: ./server/build/test_deepseek4_unit + - name: Run mixed-quantization registry safety tests on the 3090 + run: | + ctest --test-dir server/build --output-on-failure \ + -R 'rocmfp3_mix_registry|rocmfp_mix_slice_matvec|rocmfp_mix_gateup_glu|ds4_mix_registry_teardown' + # Optional model-backed end-to-end smoke (real spec-decode on the 3090), # disabled by default because it builds dflash_server and lazy-loads the # ~16 GB Qwen3.6-27B target + draft (~1-2 min). The weights are already @@ -299,11 +306,13 @@ jobs: cmake --build "$RUNNER_TEMP/rocmfp-build" \ --target test_rocmfp4 test_rocmfpx test_rocmfp4_hip_tail test_rocmfpx_mmq \ test_deepseek4_mmid_grouped_cuda test_deepseek4_unit \ - test_recurrent_snapshot test_server_unit \ + test_recurrent_snapshot test_server_unit test_rocmfp3_mix_registry \ + test_rocmfp_mix_slice_matvec test_rocmfp_mix_gateup_glu \ + test_ds4_mix_registry_teardown \ --parallel 8 ctest --test-dir "$RUNNER_TEMP/rocmfp-build" \ --output-on-failure \ - -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|deepseek4_unit|recurrent_snapshot|ChainRollbackPolicy' + -R 'rocmfp4_reference|rocmfpx_reference|rocmfp4_hip_tail|rocmfpx_mmq|deepseek4_mmid_grouped_cuda|deepseek4_unit|recurrent_snapshot|ChainRollbackPolicy|rocmfp3_mix_registry|rocmfp_mix_slice_matvec|rocmfp_mix_gateup_glu|ds4_mix_registry_teardown' build-windows: name: Build Windows (MSVC + CUDA, library + server targets) diff --git a/harness/qualification/deepseek4/qualify_ds4_q5_amd.sh b/harness/qualification/deepseek4/qualify_ds4_q5_amd.sh index 756f99709..deaec3f6a 100755 --- a/harness/qualification/deepseek4/qualify_ds4_q5_amd.sh +++ b/harness/qualification/deepseek4/qualify_ds4_q5_amd.sh @@ -32,6 +32,10 @@ case "$Q5_VERIFY" in esac FP4_Q5_X4_PLUS1="${FP4_Q5_X4_PLUS1:-auto}" CRITICAL_PATH_PLACEMENT="${CRITICAL_PATH_PLACEMENT:-0}" +MAIN_TO_PEER_RATE_EXPLICIT=0 +if [[ -n "${MAIN_TO_PEER_RATE:-}" ]]; then + MAIN_TO_PEER_RATE_EXPLICIT=1 +fi MAIN_TO_PEER_RATE="${MAIN_TO_PEER_RATE:-3.4}" BALANCE_MIN_HOT="${BALANCE_MIN_HOT:-0}" EXPERT_BUDGET_MB="${EXPERT_BUDGET_MB:-13200}" @@ -44,14 +48,12 @@ HASH_MODELS="${HASH_MODELS:-0}" CUDA_GRAPH_STATS_EVERY="${CUDA_GRAPH_STATS_EVERY:-200}" CUDA_DISABLE_GRAPHS_DEVICES="${CUDA_DISABLE_GRAPHS_DEVICES:-}" DYNAMIC_ROUTE_BALANCE="${DYNAMIC_ROUTE_BALANCE:-0}" -DYNAMIC_MAIN_SLOTS="${DYNAMIC_MAIN_SLOTS:-3}" +DYNAMIC_MAIN_SLOTS="${DYNAMIC_MAIN_SLOTS:-auto}" DYNAMIC_MAIN_SLOTS_X2="${DYNAMIC_MAIN_SLOTS_X2:-}" DYNAMIC_MAIN_SLOTS_X4="${DYNAMIC_MAIN_SLOTS_X4:-}" -VERIFY_WIDTH=$((4 + Q5_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}-$(date -u +%Y%m%dT%H%M%SZ)}" +EXPERT_TOP_K="${EXPERT_TOP_K:-4}" +RUN_ID="${RUN_ID:-}" OUT_ROOT="${OUT_ROOT:-$CHECKOUT/results/ds4_q5_context_qualification}" -OUT_DIR="$OUT_ROOT/$RUN_ID" -SERVER_LOG="$OUT_DIR/server.log" for executable in "$SERVER_BIN" "$TOKENIZER_HARNESS"; do if [[ ! -f "$executable" || ! -x "$executable" ]]; then @@ -95,18 +97,42 @@ case "$DYNAMIC_ROUTE_BALANCE" in 0|1) ;; *) echo "DYNAMIC_ROUTE_BALANCE must be 0 or 1" >&2; exit 2 ;; esac -case "$DYNAMIC_MAIN_SLOTS" in - 1|2|3|4|5|6) ;; - *) echo "DYNAMIC_MAIN_SLOTS must be an integer from 1 through 6" >&2; exit 2 ;; -esac -case "$DYNAMIC_MAIN_SLOTS_X2" in - ""|2|3|4|5|6|7|8|9|10|11|12) ;; - *) echo "DYNAMIC_MAIN_SLOTS_X2 must be empty or an integer from 2 through 12" >&2; exit 2 ;; -esac -case "$DYNAMIC_MAIN_SLOTS_X4" in - ""|4|5|6|7|8|9|10|11|12|13|14|15|16|17|18|19|20|21|22|23|24) ;; - *) echo "DYNAMIC_MAIN_SLOTS_X4 must be empty or an integer from 4 through 24" >&2; exit 2 ;; -esac +if [[ ! "$EXPERT_TOP_K" =~ ^[1-9][0-9]*$ ]] || ((EXPERT_TOP_K > 6)); then + echo "EXPERT_TOP_K must be an integer from 1 through 6" >&2 + exit 2 +fi +if [[ "$DYNAMIC_MAIN_SLOTS" != auto ]] && + { [[ ! "$DYNAMIC_MAIN_SLOTS" =~ ^[1-9][0-9]*$ ]] || + ((DYNAMIC_MAIN_SLOTS > EXPERT_TOP_K)); }; then + echo "DYNAMIC_MAIN_SLOTS must be auto or an integer from 1 through EXPERT_TOP_K ($EXPERT_TOP_K)" >&2 + exit 2 +fi +if [[ -n "$DYNAMIC_MAIN_SLOTS_X2" ]] && + { [[ ! "$DYNAMIC_MAIN_SLOTS_X2" =~ ^[1-9][0-9]*$ ]] || + ((DYNAMIC_MAIN_SLOTS_X2 < 2 || DYNAMIC_MAIN_SLOTS_X2 > 2 * EXPERT_TOP_K)); }; then + echo "DYNAMIC_MAIN_SLOTS_X2 must be empty or an integer from 2 through $((2 * EXPERT_TOP_K))" >&2 + exit 2 +fi +if [[ -n "$DYNAMIC_MAIN_SLOTS_X4" ]] && + { [[ ! "$DYNAMIC_MAIN_SLOTS_X4" =~ ^[1-9][0-9]*$ ]] || + ((DYNAMIC_MAIN_SLOTS_X4 < 4 || DYNAMIC_MAIN_SLOTS_X4 > 4 * EXPERT_TOP_K)); }; then + echo "DYNAMIC_MAIN_SLOTS_X4 must be empty or an integer from 4 through $((4 * EXPERT_TOP_K))" >&2 + exit 2 +fi +explicit_route_quotas=0 +if [[ "$DYNAMIC_MAIN_SLOTS" != auto ]]; then + ((explicit_route_quotas += 1)) +fi +if [[ -n "$DYNAMIC_MAIN_SLOTS_X2" ]]; then + ((explicit_route_quotas += 1)) +fi +if [[ -n "$DYNAMIC_MAIN_SLOTS_X4" ]]; then + ((explicit_route_quotas += 1)) +fi +if ((explicit_route_quotas > 1)); then + echo "set at most one dynamic main-slot quota" >&2 + exit 2 +fi case "$FP4_Q5_X4_PLUS1" in auto|0|1) ;; *) echo "FP4_Q5_X4_PLUS1 must be auto, 0, or 1" >&2; exit 2 ;; @@ -139,12 +165,6 @@ case "$HASH_MODELS" in 0|1) ;; *) echo "HASH_MODELS must be 0 or 1" >&2; exit 2 ;; esac -case "$RUN_ID" in - ""|.|..|*[!A-Za-z0-9._-]*) - echo "RUN_ID may contain only letters, numbers, dot, underscore, and hyphen" >&2 - exit 2 - ;; -esac for numeric_setting in PORT MAX_CTX EXPERT_BUDGET_MB WARMUP RUNS MAX_TOKENS \ VRAM_MONITOR_SECONDS CUDA_GRAPH_STATS_EVERY; do numeric_value="${!numeric_setting}" @@ -178,6 +198,48 @@ for target in "${target_args[@]}"; do fi done +DYNAMIC_BALANCE_ENV_NAME="" +DYNAMIC_BALANCE_ENV_VALUE="" +DYNAMIC_BALANCE_LABEL="off" +if [[ "$DYNAMIC_ROUTE_BALANCE" == 1 ]]; then + if [[ "$DYNAMIC_MAIN_SLOTS" != auto ]]; then + DYNAMIC_BALANCE_ENV_NAME="DFLASH_MOE_TP_DYNAMIC_MAIN_SLOTS" + DYNAMIC_BALANCE_ENV_VALUE="$DYNAMIC_MAIN_SLOTS" + DYNAMIC_BALANCE_LABEL="s${DYNAMIC_MAIN_SLOTS}" + elif [[ -n "$DYNAMIC_MAIN_SLOTS_X2" ]]; then + DYNAMIC_BALANCE_ENV_NAME="DFLASH_MOE_TP_DYNAMIC_MAIN_SLOTS_X2" + DYNAMIC_BALANCE_ENV_VALUE="$DYNAMIC_MAIN_SLOTS_X2" + DYNAMIC_BALANCE_LABEL="s2x${DYNAMIC_MAIN_SLOTS_X2}" + elif [[ -n "$DYNAMIC_MAIN_SLOTS_X4" ]]; then + DYNAMIC_BALANCE_ENV_NAME="DFLASH_MOE_TP_DYNAMIC_MAIN_SLOTS_X4" + DYNAMIC_BALANCE_ENV_VALUE="$DYNAMIC_MAIN_SLOTS_X4" + DYNAMIC_BALANCE_LABEL="s4x${DYNAMIC_MAIN_SLOTS_X4}" + elif [[ "$EXPERT_TOP_K" == 4 && "$MAIN_TO_PEER_RATE_EXPLICIT" == 0 ]]; then + # Preserve the qualified top-4 default. Automatic rate-based scaling is + # for widened top-k or an explicit operator override. + DYNAMIC_BALANCE_ENV_NAME="DFLASH_MOE_TP_DYNAMIC_MAIN_SLOTS" + DYNAMIC_BALANCE_ENV_VALUE=3 + DYNAMIC_BALANCE_LABEL="s3" + else + DYNAMIC_BALANCE_ENV_NAME="DFLASH_MOE_TP_MAIN_TO_PEER_RATE" + DYNAMIC_BALANCE_ENV_VALUE="$MAIN_TO_PEER_RATE" + DYNAMIC_BALANCE_LABEL="r${MAIN_TO_PEER_RATE}" + fi +fi + +VERIFY_WIDTH=$((4 + Q5_VERIFY)) +if [[ -z "$RUN_ID" ]]; then + RUN_ID="ds4-q${VERIFY_WIDTH}-k${EXPERT_TOP_K}-fr${FORCE_GRAPH_REPLAY}-direct${DIRECT_INDEXER_TOPK}-radix${BLOCK_RADIX_TOPK}-x4p1${FP4_Q5_X4_PLUS1}-cp${CRITICAL_PATH_PLACEMENT}-bal${DYNAMIC_BALANCE_LABEL}-$(date -u +%Y%m%dT%H%M%SZ)" +fi +case "$RUN_ID" in + .|..|*[!A-Za-z0-9._-]*) + echo "RUN_ID may contain only letters, numbers, dot, underscore, and hyphen" >&2 + exit 2 + ;; +esac +OUT_DIR="$OUT_ROOT/$RUN_ID" +SERVER_LOG="$OUT_DIR/server.log" + # The script changes physical cards 0 and 1 with rocm-smi. A visibility mask # that reorders those cards would apply the performance levels to the wrong # logical devices, so this qualification only accepts the canonical order. @@ -279,7 +341,7 @@ server_env=( "GGML_CUDA_GRAPH_STATS=1" "GGML_CUDA_GRAPH_STATS_EVERY=$CUDA_GRAPH_STATS_EVERY" "LUCE_CUDA_I32_REPEAT=1" - "DFLASH_DS4_TOPK=4" + "DFLASH_DS4_TOPK=$EXPERT_TOP_K" "DFLASH_DS4_FUSED_VERIFY=1" "DFLASH_DS4_FUSED_HYBRID_DECODE=1" "DFLASH_DS4_TIMING=1" @@ -343,19 +405,9 @@ if [[ -n "$DECODE_HOTNESS_CSV" ]]; then fi if [[ "$DYNAMIC_ROUTE_BALANCE" == 1 ]]; then server_env+=( - "DFLASH_DS4_TP_DYNAMIC_ROUTE_BALANCE=1" - "DFLASH_DS4_TP_DYNAMIC_MAIN_SLOTS=$DYNAMIC_MAIN_SLOTS" + "DFLASH_MOE_TP_DYNAMIC_ROUTE_BALANCE=1" + "$DYNAMIC_BALANCE_ENV_NAME=$DYNAMIC_BALANCE_ENV_VALUE" ) - if [[ -n "$DYNAMIC_MAIN_SLOTS_X2" ]]; then - server_env+=( - "DFLASH_DS4_TP_DYNAMIC_MAIN_SLOTS_X2=$DYNAMIC_MAIN_SLOTS_X2" - ) - fi - if [[ -n "$DYNAMIC_MAIN_SLOTS_X4" ]]; then - server_env+=( - "DFLASH_DS4_TP_DYNAMIC_MAIN_SLOTS_X4=$DYNAMIC_MAIN_SLOTS_X4" - ) - fi fi if [[ -n "$CUDA_DISABLE_GRAPHS_DEVICES" ]]; then server_env+=( @@ -417,7 +469,7 @@ server_args=( --hard-limit-reply-budget 0 --chunk 2048 --ds4-fused-decode - --ds4-expert-top-k 4 + --ds4-expert-top-k "$EXPERT_TOP_K" --ds4-prefill sparse --peer-access ) @@ -442,6 +494,7 @@ server_args=( echo "dynamic_main_slots=$DYNAMIC_MAIN_SLOTS" echo "dynamic_main_slots_x2=$DYNAMIC_MAIN_SLOTS_X2" echo "dynamic_main_slots_x4=$DYNAMIC_MAIN_SLOTS_X4" + echo "expert_top_k=$EXPERT_TOP_K" echo "cache_slots=$CACHE_SLOTS" echo "mmvq_max_ncols=$MMVQ_MAX_NCOLS" echo "targets=$TARGETS" diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index a969fa9d7..0ed5743fb 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1106,6 +1106,7 @@ if(DFLASH27B_TESTS) endif() target_include_directories(test_rocmfp3_mix_registry PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src/ggml-cuda) target_link_libraries(test_rocmfp3_mix_registry PRIVATE ggml ggml-base ${DFLASH27B_GGML_BACKEND_TARGET}) @@ -1132,6 +1133,7 @@ if(DFLASH27B_TESTS) endif() target_include_directories(test_rocmfp_mix_slice_matvec PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src/ggml-cuda) target_link_libraries(test_rocmfp_mix_slice_matvec PRIVATE ggml ggml-base ${DFLASH27B_GGML_BACKEND_TARGET}) @@ -1158,6 +1160,7 @@ if(DFLASH27B_TESTS) endif() target_include_directories(test_rocmfp_mix_gateup_glu PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/src/ggml-cuda) target_link_libraries(test_rocmfp_mix_gateup_glu PRIVATE ggml ggml-base ${DFLASH27B_GGML_BACKEND_TARGET}) diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index ac2fc92d4..b759f94e6 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -89,6 +89,19 @@ GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void); GGML_BACKEND_API bool ggml_backend_cuda_topk_rows(const struct ggml_tensor * logits, int k, float * probs_out, int32_t * ids_out); +// Attach learned per-expert decode tables to a mixed-precision tensor. The +// host variants copy the tables to the device that owns `base`. Call the +// matching unregister function before releasing the tensor's backing buffer. +// Returns false without registering when validation or device setup fails. +GGML_BACKEND_API bool ggml_cuda_rocmfp3_mix_register_host( + const void * base, size_t expert_stride, int n_experts, int out, int in, + const void * codebooks_bf16_host, const uint8_t * modes_host); +GGML_BACKEND_API bool ggml_cuda_rocmfp2_mix_register_host( + const void * base, size_t expert_stride, int n_experts, int out, int in, + const void * codebooks_bf16_host, const uint8_t * modes_host); +GGML_BACKEND_API void ggml_cuda_rocmfp2_mix_unregister(const void * base); +GGML_BACKEND_API void ggml_cuda_rocmfp3_mix_unregister(const void * base); + #ifdef __cplusplus } #endif diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu index d782aaba1..08a47c79d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu @@ -5,6 +5,36 @@ #include "rocmfp2_mix.cuh" #include "rocmfp3_mix.cuh" +namespace { + +class mix_registry_dispatch_guard { +public: + explicit mix_registry_dispatch_guard(ggml_type type) : type_(type) { + if (type_ == GGML_TYPE_Q2_1_ROCMFP2_MIX) { + ggml_cuda_rocmfp2_mix_registry_lock(); + } else if (type_ == GGML_TYPE_Q3_1_ROCMFP3_MIX) { + ggml_cuda_rocmfp3_mix_registry_lock(); + } + } + + ~mix_registry_dispatch_guard() { + if (type_ == GGML_TYPE_Q2_1_ROCMFP2_MIX) { + ggml_cuda_rocmfp2_mix_registry_unlock(); + } else if (type_ == GGML_TYPE_Q3_1_ROCMFP3_MIX) { + ggml_cuda_rocmfp3_mix_registry_unlock(); + } + } + + mix_registry_dispatch_guard(const mix_registry_dispatch_guard &) = delete; + mix_registry_dispatch_guard & operator=( + const mix_registry_dispatch_guard &) = delete; + +private: + ggml_type type_; +}; + +} // namespace + static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { const bool is_mix_type = args.type_x == GGML_TYPE_Q2_1_ROCMFP2_MIX || @@ -137,6 +167,10 @@ static void ggml_cuda_mul_mat_q_impl( const float * src1_d = (const float *) src1->data; float * dst_d = (float *) dst->data; + // Keep mix side data alive until every MMQ launch using it is enqueued. + // Registry teardown takes the same lock and drains the owning device before + // freeing those buffers. + mix_registry_dispatch_guard mix_guard(src0->type); const void * mix_codebooks_raw = nullptr; const uint8_t * mix_modes = nullptr; if (src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX) { @@ -448,6 +482,7 @@ void ggml_cuda_op_mul_mat_q( const bool use_stream_k = ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) || GGML_CUDA_CC_IS_CDNA(cc)) && src1_ncols == ne11; + mix_registry_dispatch_guard mix_guard(src0->type); const void * mix_codebooks_raw = nullptr; const uint8_t * mix_modes = nullptr; if (src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu index 6cfb4b163..0541c58f2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu @@ -3,6 +3,7 @@ #include "ggml-cuda/dequantize.cuh" #include "ggml-cuda/mmvq.cuh" +#include #include static __device__ __forceinline__ float silu_f32(float x) { @@ -444,21 +445,20 @@ static __global__ void ds4_balanced_owner_ids_kernel( int n_routes, int n_tokens, int n_expert, - int main_slots_x4, + int main_quota, bool main_owner) { if (blockIdx.x != 0 || threadIdx.x != 0) { return; } - // Assign one shared quota over the complete verification batch. Quarter- - // route units expose the otherwise unreachable 16/30 split at q=5 while - // keeping the decision device-local and identical on both owners. - const int main_quota = (main_slots_x4 * n_tokens) / 4; - int assigned_main = 0; + // Assign one shared, host-rounded quota over the complete verification + // batch while keeping the decision device-local and identical on both + // owners. + int64_t assigned_main = 0; for (int token = 0; token < n_tokens; ++token) { - const int row = token * n_routes; + const int64_t row = (int64_t) token * n_routes; for (int route = 0; route < n_routes; ++route) { - const int index = row + route; + const int64_t index = row + route; const int32_t global_id = global_ids[index]; const bool active = router_weights[index] != 0.0f; const bool valid_global = global_id >= 0 && global_id < n_expert; @@ -660,22 +660,32 @@ void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) GGML_ASSERT(local_lut && local_lut->type == GGML_TYPE_I32); GGML_ASSERT(candidate_lut && candidate_lut->type == GGML_TYPE_F32); GGML_ASSERT(dst->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_are_same_shape(global_ids, weights)); + GGML_ASSERT(ggml_are_same_shape(global_ids, dst)); GGML_ASSERT(ggml_is_contiguous(global_ids)); GGML_ASSERT(ggml_is_contiguous(weights)); GGML_ASSERT(ggml_is_contiguous(local_lut)); GGML_ASSERT(ggml_is_contiguous(candidate_lut)); GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(global_ids->ne[0] > 0 && global_ids->ne[0] <= INT_MAX / 4 && + global_ids->ne[1] > 0 && global_ids->ne[1] <= INT_MAX && + global_ids->ne[2] == 1 && global_ids->ne[3] == 1); const int n_routes = (int) global_ids->ne[0]; const int n_tokens = (int) global_ids->ne[1]; - // LUT rows are replicated across verifier tokens. Only dimension 1 is - // the base expert domain; using the total element count would permit a - // positive out-of-range ID to index a later replica. - GGML_ASSERT(local_lut->ne[0] == 1 && local_lut->ne[1] > 0); + GGML_ASSERT(local_lut->ne[0] == 1 && local_lut->ne[2] == n_tokens && + local_lut->ne[3] == 1); GGML_ASSERT(candidate_lut->ne[0] == 1 && - candidate_lut->ne[1] == local_lut->ne[1]); + candidate_lut->ne[1] == local_lut->ne[1] && + candidate_lut->ne[2] == n_tokens && + candidate_lut->ne[3] == 1); + GGML_ASSERT(local_lut->ne[1] > 0 && local_lut->ne[1] <= INT_MAX); + // Dimension 1 is the base expert domain; later dimensions are only + // verifier-token replicas and must not expand the valid ID range. const int n_expert = (int) local_lut->ne[1]; - const int main_slots_x4 = ggml_get_op_params_i32(dst, 1); + const int main_quota = ggml_get_op_params_i32(dst, 1); + GGML_ASSERT(main_quota > 0 && + (int64_t) main_quota <= (int64_t) n_routes * n_tokens); const bool main_owner = ggml_get_op_params_i32(dst, 2) != 0; ds4_balanced_owner_ids_kernel<<<1, 1, 0, ctx.stream()>>>( (const int32_t *) global_ids->data, @@ -683,7 +693,7 @@ void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) (const int32_t *) local_lut->data, (const float *) candidate_lut->data, (int32_t *) dst->data, - n_routes, n_tokens, n_expert, main_slots_x4, main_owner); + n_routes, n_tokens, n_expert, main_quota, main_owner); return; } if (mode == GGML_MOE_FUSED_ALIGN_IDS) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu index 342e396fc..c61f695bc 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu @@ -1,12 +1,14 @@ // Runtime decode for GGML_TYPE_Q2_1_ROCMFP2_MIX (106). // A per-tensor registry supplies the per-expert codebook/mode that the ggml // to_fp16 converter signature cannot carry; the deepseek4 loader registers each -// fused down-expert tensor after staging its sidecar side-data to device memory. +// mixed tensor after staging its decode tables to device memory. #include "rocmfp2_mix.cuh" #include "convert.cuh" // For ggml_cuda_op_swiglu_ds4_single: the fused path must apply the EXACT function the // standalone swiglu_ds4 kernel applies, not a re-derivation of the formula. #include "unary.cuh" +#include +#include #include #include @@ -49,15 +51,16 @@ namespace { struct MixEntry { const void * base; size_t nb02; // byte stride between experts + size_t expert_bytes; // payload bytes inside each expert stride int n_experts, out, in; const nv_bfloat16 * codebooks; // n_experts * 2 * 4 const uint8_t * modes; // n_experts - const uint8_t * rotations; // n_experts (unused until p3 rotation lands) - bool owns_device; // true => this entry cudaMalloc'd the 3 buffers above - // (register_host) and must free them on erase/update int device; // device the side-data lives on; frees must happen in that context }; -std::mutex g_mix_mtx; +// Dispatch wrappers keep this lock from lookup through kernel enqueue. It is +// recursive because MMQ takes the public dispatch lock and then calls the +// ordinary lookup helper, which also protects itself. +std::recursive_mutex g_mix_mtx; std::vector g_mix_registry; // Resolve the device that owns `p`. The mix side-data must be allocated on the SAME device as @@ -71,7 +74,8 @@ static int mix_device_of(const void * p) { cudaPointerAttributes attr{}; if (p && cudaPointerGetAttributes(&attr, p) == cudaSuccess) { // cudaMemoryTypeUnregistered/Host leave `device` meaningless; only trust device memory. - if (attr.type == cudaMemoryTypeDevice && attr.device >= 0) { + if ((attr.type == cudaMemoryTypeDevice || + attr.type == cudaMemoryTypeManaged) && attr.device >= 0) { return attr.device; } } @@ -83,23 +87,35 @@ static int mix_device_of(const void * p) { struct MixDeviceGuard { int prev = -1; bool active = false; + bool valid = false; explicit MixDeviceGuard(int dev) { if (cudaGetDevice(&prev) != cudaSuccess) return; - if (dev == prev) return; - if (cudaSetDevice(dev) == cudaSuccess) active = true; + if (dev == prev) { + valid = true; + return; + } + if (cudaSetDevice(dev) == cudaSuccess) active = valid = true; } ~MixDeviceGuard() { if (active) cudaSetDevice(prev); } }; -// Free an entry's device side-data if it owns it. Caller holds g_mix_mtx. +// Free an entry's device side-data. Caller holds g_mix_mtx, so no new launch +// can acquire these pointers. Synchronizing here drains launches that released +// the lock after enqueueing but are still running on the device. void mix_free_entry_device(MixEntry & e) { - if (!e.owns_device) return; MixDeviceGuard guard(e.device); // free where it was allocated - if (e.codebooks) cudaFree((void *) e.codebooks); - if (e.modes) cudaFree((void *) e.modes); - if (e.rotations) cudaFree((void *) e.rotations); - e.codebooks = nullptr; e.modes = nullptr; e.rotations = nullptr; - e.owns_device = false; + if (!guard.valid) { + GGML_ABORT("rocmfp2_mix: failed to select the side-data device"); + } + CUDA_CHECK(cudaDeviceSynchronize()); + const cudaError_t codebook_err = e.codebooks + ? cudaFree((void *) e.codebooks) : cudaSuccess; + const cudaError_t mode_err = e.modes + ? cudaFree((void *) e.modes) : cudaSuccess; + e.codebooks = nullptr; + e.modes = nullptr; + if (codebook_err != cudaSuccess) CUDA_CHECK(codebook_err); + if (mode_err != cudaSuccess) CUDA_CHECK(mode_err); } // Enforce the wide-load invariant for EVERY registration path. mix_block_accum reads the @@ -108,27 +124,70 @@ void mix_free_entry_device(MixEntry & e) { // is 0 only when nb % 4 == 0, i.e. in % 128 == 0. Any other `in` reads up to 6 B past the // tensor allocation, and for the last tensor in a buffer that is a fault. // -// This lives in mix_register_impl (the single chokepoint both public entrypoints funnel -// through) rather than in one of them: the non-owning ggml_cuda_rocmfp2_mix_register() -// accepts caller-managed device side-data and would otherwise still be able to register an -// unsafe shape. Checking here makes it structurally impossible for a registration path to -// skip the guard, including any added later. -void mix_validate_shape(int in) { +bool mix_validate_shape(int in) { if (in % 128 != 0) { - GGML_ABORT("rocmfp2_mix: in=%d must be a multiple of 128 (block count nb=%d must be " - "a multiple of 4) so the 16 B wide-load window stays in bounds on the " - "final block; got in %% 128 = %d", - in, in / 32, in % 128); + GGML_LOG_ERROR("rocmfp2_mix: in=%d must be a multiple of 128\n", in); + return false; + } + return true; +} + +bool mix_validate_registration( + const void * base, size_t nb02, int n_experts, int out, int in, + const void * codebooks, const void * modes) { + if (!base || nb02 == 0 || n_experts <= 0 || out <= 0 || in <= 0 || + !codebooks || !modes) { + GGML_LOG_ERROR("rocmfp2_mix: invalid registration metadata\n"); + return false; + } + if (!mix_validate_shape(in)) return false; + + const size_t row_bytes = (size_t) (in / MIX_QK) * MIX_BLOCK_BYTES; + if ((size_t) out > std::numeric_limits::max() / row_bytes) { + GGML_LOG_ERROR("rocmfp2_mix: expert tensor size overflows\n"); + return false; + } + const size_t expert_bytes = (size_t) out * row_bytes; + if (nb02 < expert_bytes) { + GGML_LOG_ERROR("rocmfp2_mix: expert stride is smaller than the tensor shape\n"); + return false; + } + const std::uintptr_t address = reinterpret_cast(base); + const std::uintptr_t stride = (std::uintptr_t) nb02; + const std::uintptr_t expert_span = (std::uintptr_t) expert_bytes; + if ((size_t) stride != nb02 || (size_t) expert_span != expert_bytes) { + GGML_LOG_ERROR("rocmfp2_mix: expert span does not fit an address\n"); + return false; + } + const std::uintptr_t available = + std::numeric_limits::max() - address; + if ((std::uintptr_t) (n_experts - 1) > available / stride) { + GGML_LOG_ERROR("rocmfp2_mix: registered expert address range overflows\n"); + return false; + } + const std::uintptr_t last_offset = + (std::uintptr_t) (n_experts - 1) * stride; + if (expert_span - 1 > available - last_offset) { + GGML_LOG_ERROR("rocmfp2_mix: final expert address range overflows\n"); + return false; } + constexpr size_t table_bytes = 2 * MIX_K * sizeof(nv_bfloat16); + if ((size_t) n_experts > + std::numeric_limits::max() / table_bytes) { + GGML_LOG_ERROR("rocmfp2_mix: codebook allocation size overflows\n"); + return false; + } + return true; } void mix_register_impl(const void * base, size_t nb02, int n_experts, int out, int in, const nv_bfloat16 * codebooks, const uint8_t * modes, - const uint8_t * rotations, bool owns_device) { - mix_validate_shape(in); - std::lock_guard lk(g_mix_mtx); - MixEntry ne{base, nb02, n_experts, out, in, codebooks, modes, rotations, - owns_device, mix_device_of(base)}; + int device) { + std::lock_guard lk(g_mix_mtx); + const size_t expert_bytes = + (size_t) out * (size_t) (in / MIX_QK) * MIX_BLOCK_BYTES; + MixEntry ne{base, nb02, expert_bytes, n_experts, out, in, codebooks, modes, + device}; for (auto & e : g_mix_registry) { if (e.base == base) { // update in place — free the old owned buffers first mix_free_entry_device(e); @@ -140,58 +199,54 @@ void mix_register_impl(const void * base, size_t nb02, int n_experts, int out, i } } // namespace -// Non-owning registration: codebooks/modes/rotations are device buffers whose -// lifetime the CALLER manages (unregister will not free them). -extern "C" void ggml_cuda_rocmfp2_mix_register( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks, const void * modes, const void * rotations) { - mix_register_impl(base, nb02, n_experts, out, in, - (const nv_bfloat16 *) codebooks, (const uint8_t *) modes, - (const uint8_t *) rotations, /*owns_device=*/false); -} - // Host-side convenience for the deepseek4 loader: stage per-expert codebooks // (bf16) and modes from host memory into device buffers, then register. The -// registry OWNS these buffers and frees them on unregister/update. rotations -// host array optional (nullptr => none rotated). On a cudaMalloc failure the -// already-allocated buffers are freed before propagating the error, so a failed -// registration leaks nothing. -extern "C" void ggml_cuda_rocmfp2_mix_register_host( +// registry owns these buffers and frees them on unregister/update. On a +// cudaMalloc failure, allocated buffers are freed before the error propagates. +extern "C" bool ggml_cuda_rocmfp2_mix_register_host( const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host) { - // Validate up front, before any cudaMalloc: mix_register_impl checks this too (it is - // the chokepoint), but reaching it would mean having already allocated three device - // buffers for a shape we are about to reject. - mix_validate_shape(in); - const size_t cb_bytes = (size_t) n_experts * 2 * 4 * sizeof(nv_bfloat16); - void * cb_dev = nullptr; void * modes_dev = nullptr; void * rots_dev = nullptr; + const void * codebooks_bf16_host, const uint8_t * modes_host) { + if (!mix_validate_registration( + base, nb02, n_experts, out, in, + codebooks_bf16_host, modes_host)) { + return false; + } + for (int i = 0; i < n_experts; ++i) { + if (modes_host[i] > 1) { + GGML_LOG_ERROR("rocmfp2_mix: unsupported mode %u\n", + (unsigned) modes_host[i]); + return false; + } + } + const size_t cb_bytes = (size_t) n_experts * 2 * MIX_K * sizeof(nv_bfloat16); + void * cb_dev = nullptr; void * modes_dev = nullptr; // Allocate where the WEIGHTS are. Without this the side-data lands on the current // device while the kernel runs on the model's device, and the first request // segfaults on a multi-GPU host. - MixDeviceGuard guard(mix_device_of(base)); + const int device = mix_device_of(base); + MixDeviceGuard guard(device); + if (!guard.valid) { + GGML_LOG_ERROR("rocmfp2_mix: failed to select the tensor's device\n"); + return false; + } cudaError_t err = cudaMalloc(&cb_dev, cb_bytes); if (err == cudaSuccess) err = cudaMemcpy(cb_dev, codebooks_bf16_host, cb_bytes, cudaMemcpyHostToDevice); if (err == cudaSuccess) err = cudaMalloc(&modes_dev, (size_t) n_experts); if (err == cudaSuccess) err = cudaMemcpy(modes_dev, modes_host, (size_t) n_experts, cudaMemcpyHostToDevice); - if (err == cudaSuccess && rotations_host) { - err = cudaMalloc(&rots_dev, (size_t) n_experts); - if (err == cudaSuccess) err = cudaMemcpy(rots_dev, rotations_host, (size_t) n_experts, cudaMemcpyHostToDevice); - } if (err != cudaSuccess) { if (cb_dev) cudaFree(cb_dev); if (modes_dev) cudaFree(modes_dev); - if (rots_dev) cudaFree(rots_dev); - CUDA_CHECK(err); // report/abort exactly as before, but only after cleanup - return; + GGML_LOG_ERROR("rocmfp2_mix: decode-table upload failed: %s\n", + cudaGetErrorString(err)); + return false; } mix_register_impl(base, nb02, n_experts, out, in, (const nv_bfloat16 *) cb_dev, - (const uint8_t *) modes_dev, (const uint8_t *) rots_dev, - /*owns_device=*/true); + (const uint8_t *) modes_dev, device); + return true; } extern "C" void ggml_cuda_rocmfp2_mix_unregister(const void * base) { - std::lock_guard lk(g_mix_mtx); + std::lock_guard lk(g_mix_mtx); for (size_t i = 0; i < g_mix_registry.size(); ++i) { if (g_mix_registry[i].base == base) { mix_free_entry_device(g_mix_registry[i]); @@ -201,20 +256,40 @@ extern "C" void ggml_cuda_rocmfp2_mix_unregister(const void * base) { } } -static bool mix_lookup(const void * vx, MixEntry & out_e, int & out_expert) { - std::lock_guard lk(g_mix_mtx); - const char * p = (const char *) vx; +static bool mix_lookup( + const void * vx, MixEntry & out_e, int & out_expert, + size_t * out_byte_offset = nullptr) { + std::lock_guard lk(g_mix_mtx); + if (!vx) return false; + const std::uintptr_t address = + reinterpret_cast(vx); for (const auto & e : g_mix_registry) { - const char * b = (const char *) e.base; - if (p >= b && p < b + (size_t) e.n_experts * e.nb02) { + if (!e.base || e.nb02 == 0 || e.n_experts <= 0) continue; + const std::uintptr_t base = + reinterpret_cast(e.base); + if (address < base) continue; + const std::uintptr_t offset = address - base; + const std::uintptr_t expert = offset / e.nb02; + const std::uintptr_t within_expert = offset % e.nb02; + if (expert < static_cast(e.n_experts) && + within_expert < e.expert_bytes) { out_e = e; - out_expert = (int) (((size_t) (p - b)) / e.nb02); + out_expert = (int) expert; + if (out_byte_offset) { + *out_byte_offset = (size_t) within_expert; + } return true; } } return false; } +static bool mix_lookup_expert_base( + const void * vx, MixEntry & out_e, int & out_expert) { + size_t byte_offset = 0; + return mix_lookup(vx, out_e, out_expert, &byte_offset) && byte_offset == 0; +} + // Branchless decode. Different lanes decode different meta bytes, so `e` is // per-lane data-dependent and the two guards (`e > 0x7E`, `exp == 0`) diverge // within a warp -- the branched form pays exec-mask save/restore + v_cmpx per @@ -289,15 +364,33 @@ __global__ void dequantize_rocmfp2_mix_kernel( } void dequantize_rocmfp2_mix_to_fp16_cuda(const void * vx, half * y, int64_t k, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry e; int expert; - if (!mix_lookup(vx, e, expert)) { + size_t byte_offset = 0; + if (!mix_lookup(vx, e, expert, &byte_offset)) { GGML_ABORT("rocmfp2_mix: tensor slice %p not registered", vx); } + const int64_t registered_elements = (int64_t) e.in * e.out; + if (k < 0 || k > registered_elements || (k > 0 && !y)) { + GGML_ABORT("rocmfp2_mix: invalid dequantization range"); + } + if (k == 0) return; + const int64_t requested_blocks = (k + MIX_QK - 1) / MIX_QK; + const size_t available_blocks = + (e.expert_bytes - byte_offset) / MIX_BLOCK_BYTES; + if (byte_offset % MIX_BLOCK_BYTES != 0 || + static_cast(requested_blocks) > available_blocks) { + GGML_ABORT("rocmfp2_mix: invalid dequantization range"); + } const nv_bfloat16 * book = e.codebooks + (size_t) expert * 2 * 4; const uint8_t * mode_ptr = e.modes + expert; const int threads = 256; - const int blocks = (int) ((k + threads - 1) / threads); + const int64_t block_count = (k + threads - 1) / threads; + if (block_count > std::numeric_limits::max()) { + GGML_ABORT("rocmfp2_mix: dequantization grid is too large"); + } + const int blocks = (int) block_count; // Portable launch: triple-chevron compiles under both nvcc and hipcc; the // hipLaunchKernelGGL macro is HIP-only and breaks the default CUDA build, // which still globs this *.cu file. @@ -747,15 +840,17 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_vec_3d( int64_t src1_token_stride, int64_t src1_slice_stride, int64_t dst_token_stride, int64_t dst_slice_stride, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry e; int slice0; - if (!mix_lookup(vx, e, slice0)) { + if (!mix_lookup_expert_base(vx, e, slice0)) { return false; // not registered -> caller keeps the dequant fallback } // The registry must cover every slice this launch will index via blockIdx.y. // Registering a 3-D dense tensor with n_experts < ne02 would silently read a // codebook belonging to another tensor, so refuse rather than corrupt. - if (slice0 != 0 || e.n_experts < nslices) { + if (slice0 != 0 || in != e.in || out != e.out || nslices <= 0 || + e.n_experts < nslices || ntokens <= 0) { return false; } const int warps_per_block = 2; @@ -784,11 +879,16 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id( int64_t ids_s0, int64_t ids_s1, int64_t src1_s1, int64_t src1_s2, int64_t dst_s1, int64_t dst_s2, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry e; int expert0; - if (!mix_lookup(vx, e, expert0)) { + if (!mix_lookup_expert_base(vx, e, expert0)) { return false; // not registered -> caller falls back to sort + dequant } + if (expert0 != 0 || in != e.in || out != e.out || + n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { + return false; + } const int warps_per_block = 2; // 64 threads (mirror the mmvq path) const int threads = warps_per_block * MIX_WARP; // Two output rows per warp (register-blocked activation reuse), so a workgroup @@ -819,12 +919,16 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id_glu( int64_t src1_s1, int64_t src1_s2, int64_t dst_s1, int64_t dst_s2, float glu_limit, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry eu, eg; int expert0_u, expert0_g; - if (!mix_lookup(vx_up, eu, expert0_u) || !mix_lookup(vx_gate, eg, expert0_g)) { + if (!mix_lookup_expert_base(vx_up, eu, expert0_u) || + !mix_lookup_expert_base(vx_gate, eg, expert0_g)) { return false; } - if (eu.in != eg.in || eu.out != eg.out || eu.n_experts != eg.n_experts) { + if (expert0_u != 0 || expert0_g != 0 || in != eu.in || out != eu.out || + eu.in != eg.in || eu.out != eg.out || eu.n_experts != eg.n_experts || + n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { return false; // not a matched pair; the caller's two-launch path is still correct } const int warps_per_block = 2; @@ -842,14 +946,27 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id_glu( bool ggml_cuda_rocmfp2_mix_registered(const void * vx) { MixEntry e; int expert; - return mix_lookup(vx, e, expert); + return mix_lookup_expert_base(vx, e, expert); +} + +void ggml_cuda_rocmfp2_mix_registry_lock() { + g_mix_mtx.lock(); +} + +void ggml_cuda_rocmfp2_mix_registry_unlock() { + g_mix_mtx.unlock(); } bool ggml_cuda_rocmfp2_mix_mmq_info( const void * vx, const void ** codebooks, const uint8_t ** modes) { + if (!codebooks || !modes) { + return false; + } MixEntry e; int expert; - if (!mix_lookup(vx, e, expert)) { + size_t byte_offset = 0; + if (!mix_lookup(vx, e, expert, &byte_offset) || + byte_offset % MIX_BLOCK_BYTES != 0) { return false; } *codebooks = e.codebooks + (size_t) expert * 2 * MIX_K; @@ -861,14 +978,15 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_vec( const void * vx, const float * x, float * y, int in, int out, int ncols, int64_t x_col_stride, int64_t y_col_stride, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry e; int expert; - if (!mix_lookup(vx, e, expert)) { + if (!mix_lookup_expert_base(vx, e, expert)) { return false; // not registered -> caller falls back to dequant->cuBLAS } - // TODO(rotation): the current artifact is rotation-free (e.rotations all 0). - // When block-Hadamard-rotated p3 experts land, fold H_32 here (per-expert - // e.rotations[expert]) exactly as the dequant path will — same hook. + if (in != e.in || out != e.out || ncols <= 0) { + return false; + } const nv_bfloat16 * book = e.codebooks + (size_t) expert * 2 * 4; const uint8_t * mode_ptr = e.modes + expert; // Launch-config-only occupancy lever (bit-exact: one warp still owns one diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cuh index c94ecf348..853f598cc 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cuh @@ -2,36 +2,12 @@ // Runtime decode support for GGML_TYPE_Q2_1_ROCMFP2_MIX (106): per-expert mixed // absmax/adaptive ROCmFP2. The 10-byte block wire is identical to q2_0_rocmfpx, so // a GGUF splice from qtype 107 is offset-preserving; -// the per-expert codebook + mode + rotation flag live out-of-band (loaded from a -// sidecar into device buffers) and are attached here via a base-pointer registry, +// the per-expert codebook and mode come from GGUF metadata or a legacy sidecar, +// then are attached here via a base-pointer registry, // because the ggml to_fp16 converter signature carries no expert/codebook context. #include "common.cuh" -extern "C" { - -// Register per-expert decode side-data for one gate or up expert tensor. -// NOTE: the serving GGUF stores gate and up SEPARATELY (ffn_gate_exps / -// ffn_up_exps), and one codebook pair is fitted per fused expert covering -// BOTH halves -- so the two tensors register the SAME codebook buffer. Do not -// assume a registered base owns its codebook exclusively. -// base : device pointer to the tensor's block data (src0->data) -// nb02 : byte stride between experts (src0->nb[2]) -// n_experts : number of experts (ne02) -// out, in : per-expert weight dims (rows, cols) -// codebooks : device buffer, n_experts * 2 * 4 __hip_bfloat16 (2 books x 4 levels) -// modes : device buffer, n_experts uint8 (0 fixed qtype-107 levels, 1 adaptive 7s1c) -// rotations : device buffer, n_experts uint8 -- always 0 here. The qtype-106 -// encoder never emits rotation and the kernel ignores it; the -// loader passes zeros and rejects anything else. -GGML_API void ggml_cuda_rocmfp2_mix_register( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks, const void * modes, const void * rotations); - -GGML_API void ggml_cuda_rocmfp2_mix_unregister(const void * base); - -} - // Dequantize one expert slice (k elements = out*in) starting at vx to half. // Called from ggml_get_to_fp16_cuda for type 106; resolves the expert + codebook // from the registry using the vx pointer. @@ -96,8 +72,16 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id( // check to confirm the sync-free mul_mat_id path will handle the node). bool ggml_cuda_rocmfp2_mix_registered(const void * vx); +// Hold this lock from side-data lookup through asynchronous kernel launch. +// Unregister/update waits for the lock, synchronizes the owning device, and +// only then frees the registry-owned codebooks and modes. +void ggml_cuda_rocmfp2_mix_registry_lock(); +void ggml_cuda_rocmfp2_mix_registry_unlock(); + // Return the device-resident side data needed by the batched MMQ loader. // Pointers are advanced to the expert containing `vx`, so both whole MoE -// tensors and registered expert slices are safe callers. +// tensors and registered expert slices are safe callers. A caller that uses +// the returned pointers in an asynchronous launch must hold the registry lock +// until that launch has been enqueued. bool ggml_cuda_rocmfp2_mix_mmq_info( const void * vx, const void ** codebooks, const uint8_t ** modes); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu index f24031509..df10f4c38 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu @@ -1,12 +1,14 @@ // Runtime decode for GGML_TYPE_Q3_1_ROCMFP3_MIX (105). // A per-tensor registry supplies the per-expert codebook/mode that the ggml // to_fp16 converter signature cannot carry; the deepseek4 loader registers each -// fused down-expert tensor after staging its sidecar side-data to device memory. +// mixed tensor after staging its decode tables to device memory. #include "rocmfp3_mix.cuh" // For ggml_cuda_op_swiglu_ds4_single -- the fused path applies the EXACT function the // standalone swiglu_ds4 kernel applies, not a re-derivation. #include "unary.cuh" #include "convert.cuh" +#include +#include #include #include @@ -49,17 +51,65 @@ namespace { struct MixEntry { const void * base; size_t nb02; // byte stride between experts + size_t expert_bytes; // payload bytes inside each expert stride int n_experts, out, in; const nv_bfloat16 * codebooks; // n_experts * 2 * 8 const uint8_t * modes; // n_experts - const uint8_t * rotations; // n_experts (unused until p3 rotation lands) - bool owns_device; // true => this entry cudaMalloc'd the 3 buffers above - // (register_host) and must free them on erase/update int device; // device the side-data lives on; frees must happen in that context }; -std::mutex g_mix_mtx; +// Dispatch wrappers keep this lock from lookup through kernel enqueue. It is +// recursive because MMQ takes the public dispatch lock and then calls the +// ordinary lookup helper, which also protects itself. +std::recursive_mutex g_mix_mtx; std::vector g_mix_registry; +bool mix_validate_registration( + const void * base, size_t nb02, int n_experts, int out, int in, + const void * codebooks, const void * modes) { + if (!base || nb02 == 0 || n_experts <= 0 || out <= 0 || in <= 0 || + in % MIX_QK != 0 || !codebooks || !modes) { + GGML_LOG_ERROR("rocmfp3_mix: invalid registration metadata\n"); + return false; + } + + const size_t row_bytes = (size_t) (in / MIX_QK) * MIX_BLOCK_BYTES; + if ((size_t) out > std::numeric_limits::max() / row_bytes) { + GGML_LOG_ERROR("rocmfp3_mix: expert tensor size overflows\n"); + return false; + } + const size_t expert_bytes = (size_t) out * row_bytes; + if (nb02 < expert_bytes) { + GGML_LOG_ERROR("rocmfp3_mix: expert stride is smaller than the tensor shape\n"); + return false; + } + const std::uintptr_t address = reinterpret_cast(base); + const std::uintptr_t stride = (std::uintptr_t) nb02; + const std::uintptr_t expert_span = (std::uintptr_t) expert_bytes; + if ((size_t) stride != nb02 || (size_t) expert_span != expert_bytes) { + GGML_LOG_ERROR("rocmfp3_mix: expert span does not fit an address\n"); + return false; + } + const std::uintptr_t available = + std::numeric_limits::max() - address; + if ((std::uintptr_t) (n_experts - 1) > available / stride) { + GGML_LOG_ERROR("rocmfp3_mix: registered expert address range overflows\n"); + return false; + } + const std::uintptr_t last_offset = + (std::uintptr_t) (n_experts - 1) * stride; + if (expert_span - 1 > available - last_offset) { + GGML_LOG_ERROR("rocmfp3_mix: final expert address range overflows\n"); + return false; + } + constexpr size_t table_bytes = 2 * MIX_K * sizeof(nv_bfloat16); + if ((size_t) n_experts > + std::numeric_limits::max() / table_bytes) { + GGML_LOG_ERROR("rocmfp3_mix: codebook allocation size overflows\n"); + return false; + } + return true; +} + // Resolve the device that owns `p`. The mix side-data must be allocated on the SAME device as // the expert tensor it describes: the kernel dereferences both together, and cudaMalloc uses // the CURRENT device, which is not necessarily the one the model was loaded onto. Returns the @@ -71,7 +121,8 @@ static int mix_device_of(const void * p) { cudaPointerAttributes attr{}; if (p && cudaPointerGetAttributes(&attr, p) == cudaSuccess) { // cudaMemoryTypeUnregistered/Host leave `device` meaningless; only trust device memory. - if (attr.type == cudaMemoryTypeDevice && attr.device >= 0) { + if ((attr.type == cudaMemoryTypeDevice || + attr.type == cudaMemoryTypeManaged) && attr.device >= 0) { return attr.device; } } @@ -83,31 +134,45 @@ static int mix_device_of(const void * p) { struct MixDeviceGuard { int prev = -1; bool active = false; + bool valid = false; explicit MixDeviceGuard(int dev) { if (cudaGetDevice(&prev) != cudaSuccess) return; - if (dev == prev) return; - if (cudaSetDevice(dev) == cudaSuccess) active = true; + if (dev == prev) { + valid = true; + return; + } + if (cudaSetDevice(dev) == cudaSuccess) active = valid = true; } ~MixDeviceGuard() { if (active) cudaSetDevice(prev); } }; -// Free an entry's device side-data if it owns it. Caller holds g_mix_mtx. +// Free an entry's device side-data. Caller holds g_mix_mtx, so no new launch +// can acquire these pointers. Synchronizing here drains launches that released +// the lock after enqueueing but are still running on the device. void mix_free_entry_device(MixEntry & e) { - if (!e.owns_device) return; MixDeviceGuard guard(e.device); // free where it was allocated - if (e.codebooks) cudaFree((void *) e.codebooks); - if (e.modes) cudaFree((void *) e.modes); - if (e.rotations) cudaFree((void *) e.rotations); - e.codebooks = nullptr; e.modes = nullptr; e.rotations = nullptr; - e.owns_device = false; + if (!guard.valid) { + GGML_ABORT("rocmfp3_mix: failed to select the side-data device"); + } + CUDA_CHECK(cudaDeviceSynchronize()); + const cudaError_t codebook_err = e.codebooks + ? cudaFree((void *) e.codebooks) : cudaSuccess; + const cudaError_t mode_err = e.modes + ? cudaFree((void *) e.modes) : cudaSuccess; + e.codebooks = nullptr; + e.modes = nullptr; + if (codebook_err != cudaSuccess) CUDA_CHECK(codebook_err); + if (mode_err != cudaSuccess) CUDA_CHECK(mode_err); } void mix_register_impl(const void * base, size_t nb02, int n_experts, int out, int in, const nv_bfloat16 * codebooks, const uint8_t * modes, - const uint8_t * rotations, bool owns_device) { - std::lock_guard lk(g_mix_mtx); - MixEntry ne{base, nb02, n_experts, out, in, codebooks, modes, rotations, - owns_device, mix_device_of(base)}; + int device) { + std::lock_guard lk(g_mix_mtx); + const size_t expert_bytes = + (size_t) out * (size_t) (in / MIX_QK) * MIX_BLOCK_BYTES; + MixEntry ne{base, nb02, expert_bytes, n_experts, out, in, codebooks, modes, + device}; for (auto & e : g_mix_registry) { if (e.base == base) { // update in place — free the old owned buffers first mix_free_entry_device(e); @@ -119,58 +184,54 @@ void mix_register_impl(const void * base, size_t nb02, int n_experts, int out, i } } // namespace -// Non-owning registration: codebooks/modes/rotations are device buffers whose -// lifetime the CALLER manages (unregister will not free them). -extern "C" void ggml_cuda_rocmfp3_mix_register( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks, const void * modes, const void * rotations) { - mix_register_impl(base, nb02, n_experts, out, in, - (const nv_bfloat16 *) codebooks, (const uint8_t *) modes, - (const uint8_t *) rotations, /*owns_device=*/false); -} - // Host-side convenience for the deepseek4 loader: stage per-expert codebooks // (bf16) and modes from host memory into device buffers, then register. The -// registry OWNS these buffers and frees them on unregister/update. rotations -// host array optional (nullptr => none rotated). On a cudaMalloc failure the -// already-allocated buffers are freed before propagating the error, so a failed -// registration leaks nothing. -extern "C" void ggml_cuda_rocmfp3_mix_register_host( +// registry owns these buffers and frees them on unregister/update. On a +// cudaMalloc failure, allocated buffers are freed before the error propagates. +extern "C" bool ggml_cuda_rocmfp3_mix_register_host( const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host) { - const size_t cb_bytes = (size_t) n_experts * 2 * 8 * sizeof(nv_bfloat16); - void * cb_dev = nullptr; void * modes_dev = nullptr; void * rots_dev = nullptr; + const void * codebooks_bf16_host, const uint8_t * modes_host) { + if (!mix_validate_registration( + base, nb02, n_experts, out, in, + codebooks_bf16_host, modes_host)) { + return false; + } + for (int i = 0; i < n_experts; ++i) { + if (modes_host[i] > 1) { + GGML_LOG_ERROR("rocmfp3_mix: unsupported mode %u\n", + (unsigned) modes_host[i]); + return false; + } + } + const size_t cb_bytes = (size_t) n_experts * 2 * MIX_K * sizeof(nv_bfloat16); + void * cb_dev = nullptr; void * modes_dev = nullptr; // Allocate where the WEIGHTS are. Without this the side-data lands on the current // device while the kernel runs on the model's device, and the first request // segfaults on a multi-GPU host. - MixDeviceGuard guard(mix_device_of(base)); + const int device = mix_device_of(base); + MixDeviceGuard guard(device); + if (!guard.valid) { + GGML_LOG_ERROR("rocmfp3_mix: failed to select the tensor's device\n"); + return false; + } cudaError_t err = cudaMalloc(&cb_dev, cb_bytes); if (err == cudaSuccess) err = cudaMemcpy(cb_dev, codebooks_bf16_host, cb_bytes, cudaMemcpyHostToDevice); if (err == cudaSuccess) err = cudaMalloc(&modes_dev, (size_t) n_experts); if (err == cudaSuccess) err = cudaMemcpy(modes_dev, modes_host, (size_t) n_experts, cudaMemcpyHostToDevice); - bool any_rotation = false; - for (int i = 0; rotations_host && i < n_experts; ++i) { - any_rotation = any_rotation || rotations_host[i] != 0; - } - if (err == cudaSuccess && any_rotation) { - err = cudaMalloc(&rots_dev, (size_t) n_experts); - if (err == cudaSuccess) err = cudaMemcpy(rots_dev, rotations_host, (size_t) n_experts, cudaMemcpyHostToDevice); - } if (err != cudaSuccess) { if (cb_dev) cudaFree(cb_dev); if (modes_dev) cudaFree(modes_dev); - if (rots_dev) cudaFree(rots_dev); - CUDA_CHECK(err); // report/abort exactly as before, but only after cleanup - return; + GGML_LOG_ERROR("rocmfp3_mix: decode-table upload failed: %s\n", + cudaGetErrorString(err)); + return false; } mix_register_impl(base, nb02, n_experts, out, in, (const nv_bfloat16 *) cb_dev, - (const uint8_t *) modes_dev, (const uint8_t *) rots_dev, - /*owns_device=*/true); + (const uint8_t *) modes_dev, device); + return true; } extern "C" void ggml_cuda_rocmfp3_mix_unregister(const void * base) { - std::lock_guard lk(g_mix_mtx); + std::lock_guard lk(g_mix_mtx); for (size_t i = 0; i < g_mix_registry.size(); ++i) { if (g_mix_registry[i].base == base) { mix_free_entry_device(g_mix_registry[i]); @@ -180,20 +241,40 @@ extern "C" void ggml_cuda_rocmfp3_mix_unregister(const void * base) { } } -static bool mix_lookup(const void * vx, MixEntry & out_e, int & out_expert) { - std::lock_guard lk(g_mix_mtx); - const char * p = (const char *) vx; +static bool mix_lookup( + const void * vx, MixEntry & out_e, int & out_expert, + size_t * out_byte_offset = nullptr) { + std::lock_guard lk(g_mix_mtx); + if (!vx) return false; + const std::uintptr_t address = + reinterpret_cast(vx); for (const auto & e : g_mix_registry) { - const char * b = (const char *) e.base; - if (p >= b && p < b + (size_t) e.n_experts * e.nb02) { + if (!e.base || e.nb02 == 0 || e.n_experts <= 0) continue; + const std::uintptr_t base = + reinterpret_cast(e.base); + if (address < base) continue; + const std::uintptr_t offset = address - base; + const std::uintptr_t expert = offset / e.nb02; + const std::uintptr_t within_expert = offset % e.nb02; + if (expert < static_cast(e.n_experts) && + within_expert < e.expert_bytes) { out_e = e; - out_expert = (int) (((size_t) (p - b)) / e.nb02); + out_expert = (int) expert; + if (out_byte_offset) { + *out_byte_offset = (size_t) within_expert; + } return true; } } return false; } +static bool mix_lookup_expert_base( + const void * vx, MixEntry & out_e, int & out_expert) { + size_t byte_offset = 0; + return mix_lookup(vx, out_e, out_expert, &byte_offset) && byte_offset == 0; +} + __device__ __forceinline__ float mix_ue4m3(uint8_t e) { if (e > 0x7E) return 0.0f; int exp = e >> 3, mant = e & 7; @@ -251,15 +332,33 @@ __global__ void dequantize_rocmfp3_mix_kernel( } void dequantize_rocmfp3_mix_to_fp16_cuda(const void * vx, half * y, int64_t k, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry e; int expert; - if (!mix_lookup(vx, e, expert)) { + size_t byte_offset = 0; + if (!mix_lookup(vx, e, expert, &byte_offset)) { GGML_ABORT("rocmfp3_mix: tensor slice %p not registered", vx); } + const int64_t registered_elements = (int64_t) e.in * e.out; + if (k < 0 || k > registered_elements || (k > 0 && !y)) { + GGML_ABORT("rocmfp3_mix: invalid dequantization range"); + } + if (k == 0) return; + const int64_t requested_blocks = (k + MIX_QK - 1) / MIX_QK; + const size_t available_blocks = + (e.expert_bytes - byte_offset) / MIX_BLOCK_BYTES; + if (byte_offset % MIX_BLOCK_BYTES != 0 || + static_cast(requested_blocks) > available_blocks) { + GGML_ABORT("rocmfp3_mix: invalid dequantization range"); + } const nv_bfloat16 * book = e.codebooks + (size_t) expert * 2 * 8; const uint8_t * mode_ptr = e.modes + expert; const int threads = 256; - const int blocks = (int) ((k + threads - 1) / threads); + const int64_t block_count = (k + threads - 1) / threads; + if (block_count > std::numeric_limits::max()) { + GGML_ABORT("rocmfp3_mix: dequantization grid is too large"); + } + const int blocks = (int) block_count; // Portable launch: triple-chevron compiles under both nvcc and hipcc; the // hipLaunchKernelGGL macro is HIP-only and breaks the default CUDA build, // which still globs this *.cu file. @@ -658,15 +757,17 @@ bool ggml_cuda_rocmfp3_mix_mul_mat_vec_3d( int64_t src1_token_stride, int64_t src1_slice_stride, int64_t dst_token_stride, int64_t dst_slice_stride, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry e; int slice0; - if (!mix_lookup(vx, e, slice0)) { + if (!mix_lookup_expert_base(vx, e, slice0)) { return false; // not registered -> caller keeps the dequant fallback } // The registry must cover every slice this launch will index via blockIdx.y. // Registering a 3-D dense tensor with n_experts < ne02 would silently read a // codebook belonging to another tensor, so refuse rather than corrupt. - if (slice0 != 0 || e.n_experts < nslices) { + if (slice0 != 0 || in != e.in || out != e.out || nslices <= 0 || + e.n_experts < nslices || ntokens <= 0) { return false; } const int warps_per_block = 2; @@ -695,11 +796,16 @@ bool ggml_cuda_rocmfp3_mix_mul_mat_id( int64_t ids_s0, int64_t ids_s1, int64_t src1_s1, int64_t src1_s2, int64_t dst_s1, int64_t dst_s2, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry e; int expert0; - if (!mix_lookup(vx, e, expert0)) { + if (!mix_lookup_expert_base(vx, e, expert0)) { return false; // not registered -> caller falls back to sort + dequant } + if (expert0 != 0 || in != e.in || out != e.out || + n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { + return false; + } const int warps_per_block = 2; // 64 threads (mirror the mmvq path) const int threads = warps_per_block * MIX_WARP; // Two output rows per warp (register-blocked activation reuse), so a workgroup @@ -727,12 +833,16 @@ bool ggml_cuda_rocmfp3_mix_mul_mat_id_glu( int64_t src1_s1, int64_t src1_s2, int64_t dst_s1, int64_t dst_s2, float glu_limit, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry eu, eg; int expert0_u, expert0_g; - if (!mix_lookup(vx_up, eu, expert0_u) || !mix_lookup(vx_gate, eg, expert0_g)) { + if (!mix_lookup_expert_base(vx_up, eu, expert0_u) || + !mix_lookup_expert_base(vx_gate, eg, expert0_g)) { return false; } - if (eu.in != eg.in || eu.out != eg.out || eu.n_experts != eg.n_experts) { + if (expert0_u != 0 || expert0_g != 0 || in != eu.in || out != eu.out || + eu.in != eg.in || eu.out != eg.out || eu.n_experts != eg.n_experts || + n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { return false; } const int warps_per_block = 2; @@ -750,14 +860,27 @@ bool ggml_cuda_rocmfp3_mix_mul_mat_id_glu( bool ggml_cuda_rocmfp3_mix_registered(const void * vx) { MixEntry e; int expert; - return mix_lookup(vx, e, expert); + return mix_lookup_expert_base(vx, e, expert); +} + +void ggml_cuda_rocmfp3_mix_registry_lock() { + g_mix_mtx.lock(); +} + +void ggml_cuda_rocmfp3_mix_registry_unlock() { + g_mix_mtx.unlock(); } bool ggml_cuda_rocmfp3_mix_mmq_info( const void * vx, const void ** codebooks, const uint8_t ** modes) { + if (!codebooks || !modes) { + return false; + } MixEntry e; int expert; - if (!mix_lookup(vx, e, expert) || e.rotations != nullptr) { + size_t byte_offset = 0; + if (!mix_lookup(vx, e, expert, &byte_offset) || + byte_offset % MIX_BLOCK_BYTES != 0) { return false; } *codebooks = e.codebooks + (size_t) expert * 2 * MIX_K; @@ -769,14 +892,15 @@ bool ggml_cuda_rocmfp3_mix_mul_mat_vec( const void * vx, const float * x, float * y, int in, int out, int ncols, int64_t x_col_stride, int64_t y_col_stride, cudaStream_t stream) { + std::lock_guard dispatch_lock(g_mix_mtx); MixEntry e; int expert; - if (!mix_lookup(vx, e, expert)) { + if (!mix_lookup_expert_base(vx, e, expert)) { return false; // not registered -> caller falls back to dequant->cuBLAS } - // TODO(rotation): the current artifact is rotation-free (e.rotations all 0). - // When block-Hadamard-rotated p3 experts land, fold H_32 here (per-expert - // e.rotations[expert]) exactly as the dequant path will — same hook. + if (in != e.in || out != e.out || ncols <= 0) { + return false; + } const nv_bfloat16 * book = e.codebooks + (size_t) expert * 2 * 8; const uint8_t * mode_ptr = e.modes + expert; // Launch-config-only occupancy lever (bit-exact: one warp still owns one diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cuh index 9c8a45cf2..e072ad3f0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cuh @@ -1,30 +1,12 @@ #pragma once // Runtime decode support for GGML_TYPE_Q3_1_ROCMFP3_MIX (105): per-expert mixed // absmax/adaptive ROCmFP3. The 14-byte block wire is identical to q3_0_rocmfpx; -// the per-expert codebook + mode + rotation flag live out-of-band (loaded from a -// sidecar into device buffers) and are attached here via a base-pointer registry, +// the per-expert codebook and mode come from GGUF metadata or a legacy sidecar, +// then are attached here via a base-pointer registry, // because the ggml to_fp16 converter signature carries no expert/codebook context. #include "common.cuh" -extern "C" { - -// Register per-expert decode side-data for one fused down-expert tensor. -// base : device pointer to the tensor's block data (src0->data) -// nb02 : byte stride between experts (src0->nb[2]) -// n_experts : number of experts (ne02) -// out, in : per-expert weight dims (rows, cols) -// codebooks : device buffer, n_experts * 2 * 8 __hip_bfloat16 (2 books x 8 levels) -// modes : device buffer, n_experts uint8 (0 legacy fixed, 1 adaptive 7s1c) -// rotations : device buffer, n_experts uint8 (1 => fold block-Hadamard on cols) -GGML_API void ggml_cuda_rocmfp3_mix_register( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks, const void * modes, const void * rotations); - -GGML_API void ggml_cuda_rocmfp3_mix_unregister(const void * base); - -} - // Dequantize one expert slice (k elements = out*in) starting at vx to half. // Called from ggml_get_to_fp16_cuda for type 105; resolves the expert + codebook // from the registry using the vx pointer. @@ -88,9 +70,16 @@ bool ggml_cuda_rocmfp3_mix_mul_mat_id( // check to confirm the sync-free mul_mat_id path will handle the node). bool ggml_cuda_rocmfp3_mix_registered(const void * vx); +// Hold this lock from side-data lookup through asynchronous kernel launch. +// Unregister/update waits for the lock, synchronizes the owning device, and +// only then frees the registry-owned codebooks and modes. +void ggml_cuda_rocmfp3_mix_registry_lock(); +void ggml_cuda_rocmfp3_mix_registry_unlock(); + // Return the device-resident side data needed by the batched MMQ loader. -// Rotated experts are rejected because that transform is not part of MMQ yet. // Pointers are advanced to the expert containing `vx`, so both whole MoE -// tensors and registered expert slices are safe callers. +// tensors and registered expert slices are safe callers. A caller that uses +// the returned pointers in an asynchronous launch must hold the registry lock +// until that launch has been enqueued. bool ggml_cuda_rocmfp3_mix_mmq_info( const void * vx, const void ** codebooks, const uint8_t ** modes); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 6ed13666e..d54ce8611 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -8425,17 +8425,23 @@ struct ggml_tensor * ggml_ds4_moe_balanced_owner_ids( GGML_ASSERT(local_id_lut->type == GGML_TYPE_I32); GGML_ASSERT(main_candidate_lut->type == GGML_TYPE_F32); GGML_ASSERT(ggml_are_same_shape(global_ids, router_weights)); - GGML_ASSERT(ggml_nelements(local_id_lut) == - ggml_nelements(main_candidate_lut)); - GGML_ASSERT(local_id_lut->ne[0] == 1 && local_id_lut->ne[1] > 0); - GGML_ASSERT(main_candidate_lut->ne[0] == 1 && - main_candidate_lut->ne[1] == local_id_lut->ne[1]); + GGML_ASSERT(global_ids->ne[0] > 0 && global_ids->ne[0] <= INT_MAX / 4 && + global_ids->ne[1] > 0 && global_ids->ne[1] <= INT_MAX && + global_ids->ne[2] == 1 && global_ids->ne[3] == 1); + GGML_ASSERT(local_id_lut->ne[0] == 1 && + local_id_lut->ne[1] > 0 && local_id_lut->ne[1] <= INT_MAX && + local_id_lut->ne[2] == global_ids->ne[1] && + local_id_lut->ne[3] == 1); + GGML_ASSERT(ggml_are_same_shape(local_id_lut, main_candidate_lut)); GGML_ASSERT(ggml_is_contiguous(global_ids)); GGML_ASSERT(ggml_is_contiguous(router_weights)); GGML_ASSERT(ggml_is_contiguous(local_id_lut)); GGML_ASSERT(ggml_is_contiguous(main_candidate_lut)); GGML_ASSERT(main_slots_x4 > 0 && - main_slots_x4 <= 4 * global_ids->ne[0]); + main_slots_x4 <= 4 * (int) global_ids->ne[0]); + const int64_t main_quota = + ((int64_t) main_slots_x4 * global_ids->ne[1] + 2) / 4; + GGML_ASSERT(main_quota > 0 && main_quota <= INT_MAX); struct ggml_tensor * result = ggml_dup_tensor(ctx, global_ids); result->op = GGML_OP_MOE_FUSED; @@ -8444,7 +8450,7 @@ struct ggml_tensor * ggml_ds4_moe_balanced_owner_ids( result->src[2] = local_id_lut; result->src[3] = main_candidate_lut; ggml_set_op_params_i32(result, 0, GGML_MOE_FUSED_BALANCED_OWNER_IDS); - ggml_set_op_params_i32(result, 1, main_slots_x4); + ggml_set_op_params_i32(result, 1, (int32_t) main_quota); ggml_set_op_params_i32(result, 2, main_owner ? 1 : 0); return result; } diff --git a/server/docs/DS4.md b/server/docs/DS4.md index daa14777c..c87641f11 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -64,23 +64,10 @@ options are available: load because the fused graph must reference every expert tensor directly. If that allocation fails, the backend logs the fallback and continues with hybrid expert placement and layered decode. -- **Adaptive (qtype-105/106) artifacts require MONOLITHIC EXPERT RESIDENCY.** Hybrid or - cold expert placement cannot decode them: the per-expert codebooks are registered against - resident tensor bases, so an expert paged to the host has nowhere to read its codebook - from. The backend says so and refuses rather than producing wrong output: - - ``` - [deepseek4] qtype-105 (mixed ROCmFP3) down experts require monolithic residency — - hybrid/cold expert placement cannot decode them. Enable fused decode or provide - enough VRAM to keep all experts resident. - ``` - - Practical floor: the 0731 adaptive artifacts are ~102.3 GB, so they need a device that can - hold that plus KV and activations. Measured 2026-08-05 — an 80 GB H100 fails at load - (`cudaMalloc failed` on a 97161 MiB expert buffer, then hybrid is refused for the reason - above), while a 140 GB H200 and Strix Halo's 125 GiB unified memory both serve them. A - UNIFORM artifact has no such constraint and runs fine under hybrid placement, so this is - specific to the adaptive formats and worth checking before choosing a deployment target. +- Adaptive qtype-105/106 experts work in monolithic mode and across two GPUs + using the same runtime. The loader gives each GPU's compact tensor the + decode-table rows for the experts it owns. CPU expert offload and mixed + CUDA/HIP peers keep the safe monolithic fallback. - `--ds4-expert-top-k N` keeps the highest-ranked `N` routed experts and renormalizes their weights. `0` uses the model default. Reducing this value is an diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index f5f13da0b..b933036c0 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include namespace dflash::common { @@ -131,14 +132,47 @@ const MoeHybridGraphPolicy & moe_hybrid_graph_policy() { return policy; } -static int dynamic_route_balance_main_slots_x4() { - static const int slots_x4 = [] { +int moe_balanced_main_slots_x4(int top_k, double main_to_peer_rate) { + if (top_k <= 0 || top_k > std::numeric_limits::max() / 4 || + !std::isfinite(main_to_peer_rate) || main_to_peer_rate <= 0.0) { + return 0; + } + + const int total = 4 * top_k; + const double peer_exact = (double) total / (main_to_peer_rate + 1.0); + const int peer_lower = std::clamp((int) std::floor(peer_exact), 0, total); + const int peer_upper = std::clamp((int) std::ceil(peer_exact), 0, total); + const auto completion_time = [=](int peer) { + return std::max((double) (total - peer) / main_to_peer_rate, + (double) peer); + }; + const int peer = completion_time(peer_upper) < completion_time(peer_lower) + ? peer_upper : peer_lower; + return total - peer; +} + +// The serial balanced-owner assignment is qualified for the q=5 DSpark +// verifier. Wider batches retain the ordinary parallel owner remap. +constexpr int kDynamicRouteBalanceMaxTokens = 5; + +static int dynamic_route_balance_main_slots_x4( + int n_used, + double * derived_main_to_peer_rate) { + struct DynamicRouteBalanceConfig { + bool enabled = false; + bool valid = true; + long explicit_main_slots_x4 = 0; + double main_to_peer_rate = 0.0; + }; + static const DynamicRouteBalanceConfig config = [] { + DynamicRouteBalanceConfig result; const char * enabled = moe_policy_env( "DFLASH_MOE_TP_DYNAMIC_ROUTE_BALANCE", "DFLASH_DS4_TP_DYNAMIC_ROUTE_BALANCE"); if (!enabled || !*enabled || std::strcmp(enabled, "0") == 0) { - return 0; + return result; } + result.enabled = true; const char * raw_slots_x4 = moe_policy_env( "DFLASH_MOE_TP_DYNAMIC_MAIN_SLOTS_X4", "DFLASH_DS4_TP_DYNAMIC_MAIN_SLOTS_X4"); @@ -148,41 +182,102 @@ static int dynamic_route_balance_main_slots_x4() { const char * raw_slots = moe_policy_env( "DFLASH_MOE_TP_DYNAMIC_MAIN_SLOTS", "DFLASH_DS4_TP_DYNAMIC_MAIN_SLOTS"); - const char * raw = "3"; - long minimum = 1; - long maximum = 6; - int scale = 4; + const int explicit_count = + (raw_slots_x4 && *raw_slots_x4 ? 1 : 0) + + (raw_slots_x2 && *raw_slots_x2 ? 1 : 0) + + (raw_slots && *raw_slots ? 1 : 0); + if (explicit_count > 1) { + result.valid = false; + return result; + } + const char * raw = raw_slots; + long scale = 4L; if (raw_slots_x4 && *raw_slots_x4) { raw = raw_slots_x4; - minimum = 4; - maximum = 24; - scale = 1; + scale = 1L; } else if (raw_slots_x2 && *raw_slots_x2) { raw = raw_slots_x2; - minimum = 2; - maximum = 12; - scale = 2; - } else if (raw_slots && *raw_slots) { - raw = raw_slots; + scale = 2L; + } + if (raw && *raw) { + errno = 0; + char * end = nullptr; + const long value = std::strtol(raw, &end, 10); + if (errno == ERANGE || end == raw || *end != '\0' || value <= 0 || + value > std::numeric_limits::max() / scale) { + result.valid = false; + return result; + } + result.explicit_main_slots_x4 = scale * value; + return result; + } + + // The legacy top-4 default assigned three routes to the main owner. + // Express that as a 3:1 rate so the same policy scales with model top-k. + result.main_to_peer_rate = 3.0; + const char * raw_rate = moe_policy_env( + "DFLASH_MOE_TP_MAIN_TO_PEER_RATE", + "DFLASH_DS4_TP_MAIN_TO_PEER_RATE"); + if (raw_rate && *raw_rate) { + errno = 0; + char * end = nullptr; + const double value = std::strtod(raw_rate, &end); + if (errno == ERANGE || end == raw_rate || *end != '\0' || + !std::isfinite(value) || value <= 0.0) { + result.valid = false; + return result; + } + result.main_to_peer_rate = value; } + return result; + }(); - errno = 0; - char * end = nullptr; - const long requested = std::strtol(raw, &end, 10); - if (errno == ERANGE || end == raw || *end != '\0' || - requested < minimum || requested > maximum) { + if (derived_main_to_peer_rate) { + *derived_main_to_peer_rate = 0.0; + } + if (!config.enabled) return 0; + if (n_used <= 0 || n_used > std::numeric_limits::max() / 4) { + static std::once_flag invalid_top_k_log; + std::call_once(invalid_top_k_log, [n_used] { + std::fprintf(stderr, + "[moe-hybrid] dynamic route balance disabled: invalid " + "top-k=%d\n", n_used); + }); + return 0; + } + + if (!config.valid) { + static std::once_flag malformed_log; + std::call_once(malformed_log, [] { std::fprintf(stderr, "[moe-hybrid] dynamic route balance disabled: " - "main slot quota is malformed or out of range\n"); - return 0; - } - const int requested_x4 = (int) requested * scale; - std::fprintf(stderr, - "[moe-hybrid] dynamic route balance requested: main_slots=%.2f\n", - 0.25 * (double) requested_x4); - return requested_x4; - }(); - return slots_x4; + "main-slot or transfer-rate configuration is malformed or " + "conflicting\n"); + }); + return 0; + } + + long requested_x4 = config.explicit_main_slots_x4; + if (requested_x4 == 0 && config.main_to_peer_rate > 0.0) { + requested_x4 = moe_balanced_main_slots_x4( + n_used, config.main_to_peer_rate); + if (requested_x4 == 0) requested_x4 = -1; + if (derived_main_to_peer_rate) { + *derived_main_to_peer_rate = config.main_to_peer_rate; + } + } + if (requested_x4 < 4 || requested_x4 > 4L * n_used) { + static std::once_flag invalid_log; + std::call_once(invalid_log, [n_used] { + std::fprintf(stderr, + "[moe-hybrid] dynamic route balance disabled: " + "four times the main slot quota must be in [4,%ld] " + "for top-k=%d\n", + 4L * n_used, n_used); + }); + return 0; + } + return (int) requested_x4; } static void add_hybrid_telemetry(MoeHybridFfnTelemetry & dst, @@ -1067,6 +1162,7 @@ bool build_moe_hybrid_ffn_graph( out.output = nullptr; out.main_output = nullptr; out.peer_output = nullptr; + out.dynamic_route_balance = false; if (!ctx || !inp || !global_ids || !router_weights || n_tokens <= 0 || cfg.n_embd <= 0 || cfg.n_ff_exp <= 0 || cfg.n_expert <= 0 || cfg.n_expert_used <= 0) { @@ -1093,40 +1189,53 @@ bool build_moe_hybrid_ffn_graph( &out.cold_local_lut, &out.cold_valid_lut, &out.cold_remap_nodes, &out.cold_nodes}; - const bool secondary_has_all_experts = - secondary_owner.available() && secondary_owner.local_by_global && + double derived_main_to_peer_rate = 0.0; + int dynamic_main_slots_x4 = + route_balance == MoeHybridRouteBalance::Allowed + ? dynamic_route_balance_main_slots_x4( + n_used, &derived_main_to_peer_rate) + : 0; + + const bool complete_secondary_map = + secondary_owner.local_by_global && (int) secondary_owner.local_by_global->size() == cfg.n_expert && - std::all_of( + std::none_of( secondary_owner.local_by_global->begin(), secondary_owner.local_by_global->end(), - [](int32_t local) { return local >= 0; }); - const int requested_main_slots_x4 = - route_balance == MoeHybridRouteBalance::Allowed && - primary_owner.available() - ? dynamic_route_balance_main_slots_x4() - : 0; - if (requested_main_slots_x4 > 0 && !secondary_has_all_experts) { - static std::once_flag warning_once; - std::call_once(warning_once, [] { + [](int32_t local) { return local < 0; }); + if (dynamic_main_slots_x4 > 0 && + (n_tokens > kDynamicRouteBalanceMaxTokens || + !primary_owner.available() || + !secondary_owner.available() || !complete_secondary_map)) { + static std::once_flag fallback_log; + std::call_once(fallback_log, [n_tokens] { std::fprintf(stderr, - "[moe-hybrid] dynamic route balance disabled: " - "the secondary owner does not contain every expert\n"); + "[moe-hybrid] dynamic route balance disabled for an " + "unsupported owner map or batch size (tokens=%d)\n", + n_tokens); }); + dynamic_main_slots_x4 = 0; } - if (secondary_has_all_experts && - requested_main_slots_x4 > 4 * n_used) { - static std::once_flag warning_once; - std::call_once(warning_once, [] { - std::fprintf(stderr, - "[moe-hybrid] dynamic main slot quota exceeds the route " - "width and was clamped\n"); - }); - } - const int dynamic_main_slots_x4 = - secondary_has_all_experts - ? std::min(requested_main_slots_x4, 4 * n_used) - : 0; out.dynamic_route_balance = dynamic_main_slots_x4 > 0; + if (dynamic_main_slots_x4 > 0) { + static std::once_flag active_log; + std::call_once( + active_log, + [dynamic_main_slots_x4, n_used, derived_main_to_peer_rate] { + if (derived_main_to_peer_rate > 0.0) { + std::fprintf(stderr, + "[moe-hybrid] dynamic route balance active: " + "main_slots=%.2f top_k=%d main_to_peer_rate=%.3f\n", + 0.25 * (double) dynamic_main_slots_x4, n_used, + derived_main_to_peer_rate); + } else { + std::fprintf(stderr, + "[moe-hybrid] dynamic route balance active: " + "main_slots=%.2f\n", + 0.25 * (double) dynamic_main_slots_x4); + } + }); + } // Keep graph construction order stable: both remaps, then both optional ID // alignments, then both expert branches. diff --git a/server/src/common/moe_hybrid_ffn_eval.h b/server/src/common/moe_hybrid_ffn_eval.h index 12f6554e4..025f041a7 100644 --- a/server/src/common/moe_hybrid_ffn_eval.h +++ b/server/src/common/moe_hybrid_ffn_eval.h @@ -14,6 +14,10 @@ namespace dflash::common { +// Choose the quarter-route main-owner quota that minimizes the slower owner's +// estimated completion time. Returns zero for invalid inputs. +int moe_balanced_main_slots_x4(int top_k, double main_to_peer_rate); + // GPU-resident residual combine graph: output = residual + hot_out + cold_correction. struct ResidualCombineGraph { ggml_context * ctx = nullptr; @@ -114,11 +118,11 @@ struct MoeHybridFfnTelemetry { // map global router IDs to each backend's compact expert stack and mask the // slots owned by the other backend without a host-side routing round trip. struct MoeHybridGraphInputs { - ggml_tensor * router_weights = nullptr; - // True only when this graph can safely assign routes dynamically. The LUT - // refresh path uses the effective graph decision rather than the raw env - // request, which may be disabled for a sparse secondary expert stack. + // True only when this graph actually uses batch-wide owner balancing. + // The request can fall back to static ownership for unsupported maps or + // widths, so consumers must not infer this from the process environment. bool dynamic_route_balance = false; + ggml_tensor * router_weights = nullptr; std::vector router_nodes; // q>1 decomposes the six selected routes into a four-wide head and a // padded two-wide tail. Keep those derived ID/weight tensors on the main diff --git a/server/src/common/moe_hybrid_storage.cpp b/server/src/common/moe_hybrid_storage.cpp index bf6a96081..57f6715cf 100644 --- a/server/src/common/moe_hybrid_storage.cpp +++ b/server/src/common/moe_hybrid_storage.cpp @@ -28,6 +28,33 @@ namespace dflash::common { +namespace { + +void unregister_mix_tensor(ggml_tensor * tensor) { + if (!tensor || !tensor->data) return; + + if (tensor->type == GGML_TYPE_Q3_1_ROCMFP3_MIX) { + ggml_cuda_rocmfp3_mix_unregister(tensor->data); + } else if (tensor->type == GGML_TYPE_Q2_1_ROCMFP2_MIX) { + ggml_cuda_rocmfp2_mix_unregister(tensor->data); + } +} + +} // namespace + +void MoeHybridStorage::unregister_mix_tensors() { + for (MoeHybridLayerStorage & layer : layers) { + unregister_mix_tensor(layer.gate_hot); + unregister_mix_tensor(layer.up_hot); + unregister_mix_tensor(layer.down_hot); + unregister_mix_tensor(layer.gate_up_hot); + unregister_mix_tensor(layer.gate_cold); + unregister_mix_tensor(layer.up_cold); + unregister_mix_tensor(layer.down_cold); + unregister_mix_tensor(layer.gate_up_cold); + } +} + static bool duplicate_hot_experts_on_cold_gpu() { static const bool enabled = []() { const char * raw = std::getenv("DFLASH_MOE_DUPLICATE_HOT_ON_COLD"); @@ -149,6 +176,9 @@ static ggml_tensor * new_like_with_expert_count(ggml_context * ctx, ggml_tensor } // namespace MoeHybridStorage::~MoeHybridStorage() { + // Registry entries point into the owner buffers, so remove them before the + // buffers can be released or their addresses reused. + unregister_mix_tensors(); if (prefill_route_alloc) { ggml_gallocr_free(prefill_route_alloc); prefill_route_alloc = nullptr; diff --git a/server/src/common/moe_hybrid_storage.h b/server/src/common/moe_hybrid_storage.h index 866b23728..b764ba48e 100644 --- a/server/src/common/moe_hybrid_storage.h +++ b/server/src/common/moe_hybrid_storage.h @@ -220,6 +220,9 @@ struct MoeHybridStorage { // Per-layer file region metadata for streaming (populated when mmap is active). std::vector layer_regions; + // Remove decode-table registrations while their owner tensors are alive. + // Safe to call repeatedly, including during failed partial registration. + void unregister_mix_tensors(); bool matches(const MoeHybridConfig & cfg) const; bool empty() const; bool has_mmap() const { return mmap_data != nullptr && mmap_size > 0; } diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index c50368a80..d24bb282f 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1343,23 +1343,29 @@ bool DeepSeek4Backend::init_hybrid_model() { return true; } - // qtype-105 (Q3_1_ROCMFP3_MIX) is decodable only by the fused CUDA/HIP - // kernel driven by the per-expert sidecar codebooks/modes. Hybrid storage - // slices the stacked down-expert tensor into separate hot/cold buffers that - // carry no sidecar metadata, and the CPU cold path has no type-105 traits — - // so a hybrid placement has no working decode path (the GPU cold path hits - // the unregistered-tensor abort; the CPU cold path has no vec_dot for 105). - // Fail at load with a clear message instead of starting a server that - // crashes or emits garbage at first decode. (The all-hot case above reloads - // the full monolithic model and never reaches here.) - // The same reasoning applies verbatim to qtype-106 (Q2_1_ROCMFP2_MIX) gate/up: hybrid - // slicing drops their registry metadata too, and there is no CPU vec_dot for 106 either. - // Only the down projection was checked, so a model with adaptive gate/up over uniform - // down experts passed this gate and went on to materialise unregistered hot/cold tensors - // -- aborting at first decode, or worse, decoding with whatever the registry lookup - // returned. Check every expert tensor that can carry a mix qtype. + const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); + const bool inprocess_tp = tp.requested && tp.in_process; + const PlacementBackend local_kind = + cfg_.device.backend == PlacementBackend::Auto + ? compiled_placement_backend() : cfg_.device.backend; + if (inprocess_tp && !tp.backend_valid) { + std::fprintf(stderr, + "[deepseek4-moe-tp] invalid DFLASH_DS4_MOE_TP_BACKEND; " + "expected cuda or hip\n"); + return false; + } + const bool same_runtime_tp = inprocess_tp && tp.backend_valid && + tp.secondary_backend == local_kind; + + // Mix qtypes need a learned decode table for every resident tensor. The + // same-runtime GPU path registers the matching expert subset after it + // creates both owner tensors. CPU offload and cross-runtime peers keep the + // existing monolithic fallback. + bool has_mix_experts = false; for (const auto & L : w_.layers) { const struct { ggml_tensor * t; int qtype; const char * what; } mix_experts[] = { + { L.ffn_gate_exps, GGML_TYPE_Q3_1_ROCMFP3_MIX, "qtype-105 (mixed ROCmFP3) gate" }, + { L.ffn_up_exps, GGML_TYPE_Q3_1_ROCMFP3_MIX, "qtype-105 (mixed ROCmFP3) up" }, { L.ffn_down_exps, GGML_TYPE_Q3_1_ROCMFP3_MIX, "qtype-105 (mixed ROCmFP3) down" }, { L.ffn_gate_exps, GGML_TYPE_Q2_1_ROCMFP2_MIX, "qtype-106 (mixed ROCmFP2) gate" }, { L.ffn_up_exps, GGML_TYPE_Q2_1_ROCMFP2_MIX, "qtype-106 (mixed ROCmFP2) up" }, @@ -1369,13 +1375,9 @@ bool DeepSeek4Backend::init_hybrid_model() { if (!m.t || m.t->type != m.qtype) { continue; } - // Mix qtypes decode only from registered resident tensors, so a - // cold-slice placement can never serve them. But refusing outright - // is wrong when the MONOLITHIC footprint fits the device -- the - // 2.5-3.5 bpw mix formats exist precisely so the whole file fits - // where a hybrid split of a larger quant was needed. Mirror the - // all-hot branch above: drop the partial load and attempt the full - // load, failing with the allocator's error if it does not fit. + has_mix_experts = true; + if (same_runtime_tp) continue; + std::fprintf(stderr, "[deepseek4] %s experts cannot decode from hybrid/cold " "placement; falling back to monolithic full load\n", m.what); @@ -1390,22 +1392,38 @@ bool DeepSeek4Backend::init_hybrid_model() { } } +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + if (same_runtime_tp && has_mix_experts) { + const char * mix_mmq = std::getenv("DFLASH_DS4_MIX_MMQ_PREFILL"); + if (mix_mmq && std::strcmp(mix_mmq, "0") == 0) { + std::fprintf(stderr, + "[deepseek4] heterogeneous mixed experts require " + "DFLASH_DS4_MIX_MMQ_PREFILL=1\n"); + return false; + } + if ((!mix_mmq || !mix_mmq[0]) && + ::setenv("DFLASH_DS4_MIX_MMQ_PREFILL", "1", 1) != 0) { + std::fprintf(stderr, + "[deepseek4] failed to enable mixed-expert MMQ prefill\n"); + return false; + } + } +#endif + auto hybrid = std::make_shared(); + const auto fail_hybrid_init = [&]() { + stream_engine_.destroy(); + hybrid.reset(); + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } + return false; + }; 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; 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) { std::fprintf(stderr, "[deepseek4-moe-tp] in-process secondary device must " @@ -1442,11 +1460,15 @@ bool DeepSeek4Backend::init_hybrid_model() { cfg_.model_path, backend_, w_, moe_placement_, &hybrid_cfg, *hybrid, &err, expert_backend_)) { std::fprintf(stderr, "[deepseek4] failed to build hybrid expert storage: %s\n", err.c_str()); - if (expert_backend_) { - ggml_backend_free(expert_backend_); - expert_backend_ = nullptr; - } - return false; + return fail_hybrid_init(); + } + if (same_runtime_tp && has_mix_experts && + !register_deepseek4_moe_hybrid_mix_tables( + cfg_.model_path, w_, *hybrid, &err)) { + std::fprintf(stderr, + "[deepseek4] failed to register hybrid mixed experts: %s\n", + err.c_str()); + return fail_hybrid_init(); } // The physical placement is shared by both phases. Decode may own only a @@ -1457,7 +1479,7 @@ bool DeepSeek4Backend::init_hybrid_model() { w_.n_layer, w_.n_expert, w_.n_expert_used)) { std::fprintf(stderr, "[deepseek4] decode placement dimensions are invalid\n"); - return false; + return fail_hybrid_init(); } for (int il = 0; il < w_.n_layer; ++il) { MoeHybridLayerStorage & layer = hybrid->layers[(size_t) il]; @@ -1471,7 +1493,7 @@ bool DeepSeek4Backend::init_hybrid_model() { "[deepseek4] decode owner expert %d in layer " "%d is not resident\n", (int) expert, il); - return false; + return fail_hybrid_init(); } layer.decode_hot_local_by_global[(size_t) expert] = layer.hot_local_by_global[(size_t) expert]; @@ -1489,12 +1511,12 @@ bool DeepSeek4Backend::init_hybrid_model() { } if (max_expert_bytes == 0) { std::fprintf(stderr, "[deepseek4] failed to compute streaming expert size\n"); - return false; + return fail_hybrid_init(); } if (!stream_engine_.init(backend_, max_expert_bytes, &err)) { std::fprintf(stderr, "[deepseek4] failed to init cold-expert stream engine: %s\n", err.c_str()); - return false; + return fail_hybrid_init(); } std::fprintf(stderr, "[deepseek4] cold-expert stream engine ready: pinned=%.1f MiB scratch=%.1f MiB\n", diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 78a3d9160..eea667423 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -442,6 +442,14 @@ bool build_deepseek4_moe_hybrid_storage_from_file_with_mmap( std::string * err = nullptr, ggml_backend_t cold_gpu_backend = nullptr); +// Attach each compact GPU owner tensor to the learned decode-table rows for +// the global experts stored in that tensor. +bool register_deepseek4_moe_hybrid_mix_tables( + const std::string & path, + const DeepSeek4Weights & w, + MoeHybridStorage & storage, + std::string * err = nullptr); + // Snapshot struct DeepSeek4Snapshot { int cur_pos = 0; diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index 3e832341b..2ff7526c8 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -28,33 +28,15 @@ #include #include #include +#include #include #include #include #include #include -#include extern "C" bool ggml_backend_cuda_buffer_is_managed(ggml_backend_buffer_t buffer); -// Runtime decode registration for GGML_TYPE_Q3_1_ROCMFP3_MIX (105). Defined in -// ggml-cuda/rocmfp3_mix.cu; declared here (no HIP include in the loader) so the -// deepseek4 loader can stage per-expert codebooks/modes to device and register -// each fused down-expert tensor's base pointer + stride with the decoder. -extern "C" void ggml_cuda_rocmfp3_mix_register_host( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host); -extern "C" void ggml_cuda_rocmfp2_mix_register_host( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host); -// Remove a qtype-105/106 tensor's registry entry and free its device side-data. -// Called from free_deepseek4_weights before the GPU buffer is released so -// stale base-pointer ranges can't survive an unload/reload. -extern "C" void ggml_cuda_rocmfp2_mix_unregister(const void * base); -extern "C" void ggml_cuda_rocmfp3_mix_unregister(const void * base); - #if !defined(_WIN32) #include #include @@ -354,25 +336,160 @@ static bool ds4_seek_fwd(std::FILE * f, size_t bytes) { #endif } -// Read the ".p4mix.bin" sidecar (produced on the H100) and register each -// qtype-105 fused down-expert tensor's device base + per-expert codebooks/modes -// with the CUDA/HIP decoder. No-op for uniform (qtype-104) models or when no -// sidecar is present. Called once after weights are uploaded, when the tensor -// data pointers (which the mul_mat_id -> to_fp16 fallback consults) are final. -// Returns false (and the caller must fail the load) if this model carries -// qtype-105 down-experts that cannot be fully and validly registered — an -// unregistered qtype-105 tensor has no working decode path (the to_fp16 shim -// GGML_ABORTs), so a partial/garbage registration would only turn into a -// delayed crash or silently-wrong output at first decode. -// Sanity bounds for a corrupt/hostile sidecar entry, checked BEFORE any buffer -// is sized from its fields. Far above any real DeepSeek MoE (256–512 experts, -// n_ff ~ 2k–18k) but small enough that E*C*K*2 bytes can't overflow size_t. +// Bounds for untrusted decode-table metadata. They are well above supported DS4 +// shapes but prevent corrupt headers from driving allocations or integer casts. constexpr uint32_t DS4_P4MIX_MAX_EXPERTS = 1u << 16; // 65536 constexpr uint32_t DS4_P4MIX_MAX_DIM = 1u << 20; // 1,048,576 constexpr uint32_t DS4_P4MIX_C = 2; // codebooks per expert constexpr uint32_t DS4_P4MIX_K = 8; // levels per codebook constexpr uint32_t DS4_P4MIX_QK = 32; // MIX_QK block width (rocmfp3_mix.cu) constexpr uint8_t DS4_P4MIX_MAX_MODE = 1; // 0 = fixed, 1 = adaptive +constexpr uint32_t DS4_GUMIX_K = 4; // qtype-106 levels per codebook +constexpr uint32_t DS4_ROCMFP2_ROW_ALIGN = 128; // qtype-106 wide-load invariant + +struct Ds4MixTable { + ggml_type type; + uint32_t n_experts; + uint32_t out_dim; + uint32_t in_dim; + const std::vector & books; + const std::vector & modes; +}; + +static bool ds4_register_compact_mix_tensor( + ggml_tensor * target, + const std::vector & global_expert_ids, + const Ds4MixTable & table, + const std::string & label, + std::vector & registered_bases) { + if (!target) { + if (global_expert_ids.empty()) return true; + std::fprintf(stderr, + "[deepseek4] %s has %zu assigned experts but no tensor\n", + label.c_str(), global_expert_ids.size()); + return false; + } + if (global_expert_ids.empty()) { + std::fprintf(stderr, + "[deepseek4] %s has a tensor but no assigned experts\n", + label.c_str()); + return false; + } + if (table.type != GGML_TYPE_Q3_1_ROCMFP3_MIX && + table.type != GGML_TYPE_Q2_1_ROCMFP2_MIX) { + std::fprintf(stderr, "[deepseek4] %s has unsupported mix qtype %d\n", + label.c_str(), (int) table.type); + return false; + } + const size_t books_per_expert = + table.type == GGML_TYPE_Q3_1_ROCMFP3_MIX ? 2u * DS4_P4MIX_K + : 2u * DS4_GUMIX_K; + if ((size_t) table.n_experts > + std::numeric_limits::max() / books_per_expert || + global_expert_ids.size() > (size_t) table.n_experts || + global_expert_ids.size() > + (size_t) std::numeric_limits::max()) { + std::fprintf(stderr, + "[deepseek4] compact mix metadata is too large for %s\n", + label.c_str()); + return false; + } + const size_t expected_books = + (size_t) table.n_experts * books_per_expert; + if (target->type != table.type || !target->data || + target->ne[0] != (int64_t) table.in_dim || + target->ne[1] != (int64_t) table.out_dim || + target->ne[2] != (int64_t) global_expert_ids.size() || + table.modes.size() != table.n_experts || + table.books.size() != expected_books) { + std::fprintf(stderr, + "[deepseek4] invalid compact mix metadata for %s\n", + label.c_str()); + return false; + } + std::vector books( + global_expert_ids.size() * books_per_expert); + std::vector modes(global_expert_ids.size()); + std::vector seen(table.n_experts, false); + for (size_t local = 0; local < global_expert_ids.size(); ++local) { + const int32_t global = global_expert_ids[local]; + if (global < 0 || (uint32_t) global >= table.n_experts) { + std::fprintf(stderr, + "[deepseek4] %s has out-of-range expert id %d\n", + label.c_str(), (int) global); + return false; + } + if (seen[(size_t) global]) { + std::fprintf(stderr, + "[deepseek4] %s assigns expert id %d more than once\n", + label.c_str(), (int) global); + return false; + } + seen[(size_t) global] = true; + if (table.modes[(size_t) global] > DS4_P4MIX_MAX_MODE) { + std::fprintf(stderr, + "[deepseek4] %s expert id %d has unsupported mode %u\n", + label.c_str(), (int) global, + (unsigned) table.modes[(size_t) global]); + return false; + } + std::memcpy( + books.data() + local * books_per_expert, + table.books.data() + (size_t) global * books_per_expert, + books_per_expert * sizeof(uint16_t)); + modes[local] = table.modes[(size_t) global]; + } + + const bool registered = table.type == GGML_TYPE_Q3_1_ROCMFP3_MIX + ? ggml_cuda_rocmfp3_mix_register_host( + target->data, target->nb[2], (int) global_expert_ids.size(), + (int) table.out_dim, (int) table.in_dim, books.data(), + modes.data()) + : ggml_cuda_rocmfp2_mix_register_host( + target->data, target->nb[2], (int) global_expert_ids.size(), + (int) table.out_dim, (int) table.in_dim, books.data(), + modes.data()); + if (!registered) { + std::fprintf(stderr, + "[deepseek4] failed to register mixed tensor for %s\n", + label.c_str()); + return false; + } + registered_bases.push_back(target->data); + return true; +} + +static bool ds4_register_hybrid_mix_tensors( + MoeHybridStorage & hybrid, + uint32_t layer, + uint32_t surface, + const Ds4MixTable & table, + std::vector & registered_bases) { + if (layer >= hybrid.layers.size() || surface > 2) { + std::fprintf(stderr, + "[deepseek4] compact mix target layer/surface is out of range\n"); + return false; + } + + MoeHybridLayerStorage & storage = hybrid.layers[layer]; + ggml_tensor * hot = surface == 0 ? storage.gate_hot + : surface == 1 ? storage.up_hot + : storage.down_hot; + ggml_tensor * cold = surface == 0 ? storage.gate_cold + : surface == 1 ? storage.up_cold + : storage.down_cold; + + const std::string prefix = "layer " + std::to_string(layer) + " " + + (surface == 0 ? "gate" : surface == 1 ? "up" : "down"); + if (!ds4_register_compact_mix_tensor( + hot, storage.hot_expert_ids, table, + prefix + " primary owner", registered_bases)) { + return false; + } + return ds4_register_compact_mix_tensor( + cold, storage.cold_expert_ids, table, + prefix + " secondary owner", registered_bases); +} // ---- learned-codebook sidecars: embedded in the GGUF, or a loose file beside it ---- // The adaptive mix qtypes keep per-expert codebooks out of band because a ggml block has @@ -436,60 +553,106 @@ static inline bool ds4_size_add_overflow(size_t a, size_t b, size_t * out) { #endif } -static FILE * ds4_open_sidecar(const std::string & gguf_path, - const char * kv_key, - const char * suffix, - std::vector & backing) { +// Current embedded tables are below 1 MiB. Keep format headroom without allowing +// an unbounded metadata copy during model load. +constexpr size_t DS4_MAX_EMBEDDED_TABLE_BYTES = 64u * 1024u * 1024u; + +struct Ds4TableInput { + FILE * file = nullptr; + std::string source; +}; + +static Ds4TableInput ds4_open_decode_table( + const std::string & gguf_path, + const char * kv_key, + const char * suffix, + std::vector & backing) { + Ds4TableInput result; struct gguf_init_params gip = { /*no_alloc=*/ true, /*ctx=*/ nullptr }; struct gguf_context * g = gguf_init_from_file(gguf_path.c_str(), gip); if (g) { const int64_t id = gguf_find_key(g, kv_key); - if (id >= 0 && gguf_get_kv_type(g, id) == GGUF_TYPE_ARRAY && - gguf_get_arr_type(g, id) == GGUF_TYPE_UINT8) { + if (id >= 0) { + if (gguf_get_kv_type(g, id) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(g, id) != GGUF_TYPE_UINT8) { + std::fprintf(stderr, + "[deepseek4] embedded decode table %s has the wrong type\n", + kv_key); + gguf_free(g); + return result; + } const size_t n = gguf_get_arr_n(g, id); const uint8_t * d = (const uint8_t *) gguf_get_arr_data(g, id); + if (n == 0 || n > DS4_MAX_EMBEDDED_TABLE_BYTES || !d) { + std::fprintf(stderr, + "[deepseek4] embedded decode table %s has invalid size %zu\n", + kv_key, n); + gguf_free(g); + return result; + } backing.assign(d, d + n); gguf_free(g); - std::fprintf(stderr, "[deepseek4] %s: using codebooks embedded in the GGUF " - "(%zu bytes, no sidecar file needed)\n", kv_key, backing.size()); - return ds4_fopen_memory(backing); + result.file = ds4_fopen_memory(backing); + if (!result.file) { + std::fprintf(stderr, + "[deepseek4] could not open embedded decode table %s\n", + kv_key); + return result; + } + result.source = "embedded GGUF metadata"; + return result; } gguf_free(g); } - const std::string p = gguf_path + suffix; - return std::fopen(p.c_str(), "rb"); + result.source = gguf_path + suffix; + result.file = std::fopen(result.source.c_str(), "rb"); + return result; } static bool ds4_register_p4mix_sidecar(const std::string & gguf_path, const TargetLoadPlan & plan, - DeepSeek4Weights & out) { + const DeepSeek4Weights & out, + MoeHybridStorage * hybrid = nullptr) { // required[layer] == true for a qtype-105 down-expert resident on this shard. std::vector required(out.layers.size(), false); int n_qtype105 = 0; for (size_t li = 0; li < out.layers.size(); ++li) { - const ggml_tensor * dt = out.layers[li].ffn_down_exps; + const DeepSeek4Layer & layer = out.layers[li]; + if ((layer.ffn_gate_exps && + (int) layer.ffn_gate_exps->type == DS4_QTYPE_ROCMFP3_MIX) || + (layer.ffn_up_exps && + (int) layer.ffn_up_exps->type == DS4_QTYPE_ROCMFP3_MIX)) { + std::fprintf(stderr, + "[deepseek4] qtype-105 is supported only for down experts " + "(layer %zu)\n", + li); + return false; + } + const ggml_tensor * dt = layer.ffn_down_exps; if (dt && (int) dt->type == DS4_QTYPE_ROCMFP3_MIX) { required[li] = true; n_qtype105++; } } if (n_qtype105 == 0) return true; // uniform (qtype-104) model — nothing to do - // Metadata-only / expert-skipping shards keep the expert tensor allocated but - // never upload its weights (see should_upload_ds4_tensor), so its device base - // is unfilled. Such shards don't decode experts locally, so skip registration - // rather than register a garbage base. (Hybrid/cold placement of a qtype-105 - // model — which likewise sets skip_expert_tensors — is rejected earlier in - // DeepSeek4Backend::init_hybrid_model, since it has no working 105 decode.) - if (plan.skip_expert_tensors) { + // Metadata-only shards do not own expert data. Hybrid loading also skips + // the original expert tensors, but registers the compact owner tensors + // after both GPU allocations have been materialized. + if (plan.skip_expert_tensors && !hybrid) { std::fprintf(stderr, "[deepseek4] qtype-105 down-experts not resident on this " "shard (skip_expert_tensors) — fused decode disabled here\n"); return true; } - const std::string sc_path = gguf_path + ".p4mix.bin"; - std::vector sc_embedded; // outlives f when the blob came from GGUF KV - FILE * f = ds4_open_sidecar(gguf_path, "deepseek4.p4mix.sidecar", ".p4mix.bin", sc_embedded); + std::vector embedded_table; // keeps an embedded FILE view alive + const Ds4TableInput table_input = ds4_open_decode_table( + gguf_path, "deepseek4.p4mix.sidecar", ".p4mix.bin", embedded_table); + FILE * f = table_input.file; if (!f) { - std::fprintf(stderr, "[deepseek4] qtype-105 down-experts require sidecar but it " - "is missing: %s\n", sc_path.c_str()); + if (!table_input.source.empty()) { + std::fprintf(stderr, + "[deepseek4] qtype-105 decode tables are missing from GGUF " + "metadata and legacy file %s\n", + table_input.source.c_str()); + } return false; } char magic[8]; @@ -498,11 +661,19 @@ static bool ds4_register_p4mix_sidecar(const std::string & gguf_path, std::memcmp(magic, "P4MIXv1\0", 8) != 0 || std::fread(&n_layers, 4, 1, f) != 1 || std::fread(&reserved, 4, 1, f) != 1) { - std::fprintf(stderr, "[deepseek4] bad p4mix sidecar header: %s\n", sc_path.c_str()); + std::fprintf(stderr, "[deepseek4] bad p4mix table header in %s\n", + table_input.source.c_str()); std::fclose(f); return false; } (void) reserved; + if (n_layers == 0 || (size_t) n_layers > out.layers.size()) { + std::fprintf(stderr, + "[deepseek4] p4mix table has invalid entry count %u\n", + n_layers); + std::fclose(f); + return false; + } std::vector done(out.layers.size(), false); // resident layers registered std::vector registered_bases; // for unwind on failure @@ -567,7 +738,7 @@ static bool ds4_register_p4mix_sidecar(const std::string & gguf_path, layer, idim, DS4_P4MIX_QK); ok = false; break; } - if (!dt->data) { + if (!dt->data && !hybrid) { std::fprintf(stderr, "[deepseek4] p4mix layer %u expert data not resident\n", layer); ok = false; break; } @@ -587,9 +758,8 @@ static bool ds4_register_p4mix_sidecar(const std::string & gguf_path, std::fprintf(stderr, "[deepseek4] truncated p4mix entry (layer %u)\n", layer); ok = false; break; } - // Value validation: only modes 0/1 are decoded, and rotation is not yet - // implemented (the CUDA kernel ignores e.rotations), so a nonzero rotation - // would silently produce wrong output — reject it. + // Only modes 0/1 are decoded. Rotation is not implemented, so accepting + // a nonzero value would silently produce incorrect output. for (uint32_t e = 0; e < E && ok; ++e) { if (modes[e] > DS4_P4MIX_MAX_MODE) { std::fprintf(stderr, "[deepseek4] p4mix layer %u expert %u unsupported mode %u\n", @@ -603,10 +773,25 @@ static bool ds4_register_p4mix_sidecar(const std::string & gguf_path, } if (!ok) break; - ggml_cuda_rocmfp3_mix_register_host( - dt->data, dt->nb[2], (int) E, (int) odim, (int) idim, - books.data(), modes.data(), rots.data()); - registered_bases.push_back(dt->data); + if (hybrid) { + const Ds4MixTable table{ + GGML_TYPE_Q3_1_ROCMFP3_MIX, E, odim, idim, + books, modes}; + ok = ds4_register_hybrid_mix_tensors( + *hybrid, layer, 2, table, registered_bases); + } else { + if (!ggml_cuda_rocmfp3_mix_register_host( + dt->data, dt->nb[2], (int) E, (int) odim, (int) idim, + books.data(), modes.data())) { + std::fprintf(stderr, + "[deepseek4] failed to register p4mix layer %u\n", + layer); + ok = false; + break; + } + registered_bases.push_back(dt->data); + } + if (!ok) break; done[layer] = true; } std::fclose(f); @@ -629,17 +814,16 @@ static bool ds4_register_p4mix_sidecar(const std::string & gguf_path, for (const void * b : registered_bases) ggml_cuda_rocmfp3_mix_unregister(b); return false; } - std::fprintf(stderr, "[deepseek4] registered %d qtype-105 down-expert layer(s) " - "from %s\n", (int) registered_bases.size(), sc_path.c_str()); + std::fprintf(stderr, "[deepseek4] registered %d qtype-105 down-expert tensor(s) " + "from %s\n", (int) registered_bases.size(), + table_input.source.c_str()); return true; } constexpr int DS4_QTYPE_ROCMFP2_MIX = 106; // GGML_TYPE_Q2_1_ROCMFP2_MIX -// Read the ".gumix.bin" sidecar and register each qtype-106 gate/up expert -// tensor's device base + per-expert codebooks/modes with the CUDA/HIP decoder. -// Mirrors ds4_register_p4mix_sidecar's contract and validation discipline; three -// things genuinely differ, and each is why this is a separate function rather than -// a parameterised one: +// Read the ".gumix.bin" sidecar and register each qtype-106 expert tensor's +// device base and per-expert decode tables. This format differs from p4mix in +// three ways: // // 1. TWO tensors per layer. The serving GGUF stores gate and up separately // (ffn_gate_exps / ffn_up_exps), so an entry carries a `surface` selector and @@ -650,9 +834,7 @@ constexpr int DS4_QTYPE_ROCMFP2_MIX = 106; // GGML_TYPE_Q2_1_ROCMFP2_MIX // is keyed by base pointer and frees what it owns, so sharing one buffer // between two entries would double-free on unregister. // 3. No rotation field. The qtype-106 encoder never emits rotation, so the wire -// omits it entirely rather than carrying a byte that must always be zero. The -// kernel's register_host still takes a rotations argument, so zeros are -// passed explicitly. +// omits it entirely rather than carrying a byte that must always be zero. // ---- ".dmix.bin": per-tensor codebooks for DENSE mix-qtype tensors ---- // Registers dense (non-MoE) qtype-105/106 tensors -- the attention stack -- with the // device decoder. Motivated by measurement rather than symmetry: on attention the learned @@ -679,12 +861,6 @@ constexpr int DS4_QTYPE_ROCMFP2_MIX = 106; // GGML_TYPE_Q2_1_ROCMFP2_MIX // 3. Classes, not surfaces. Five attention weight classes per layer, any subset of // which may be a mix qtype in a given artifact. constexpr uint32_t DS4_DMIX_CLASSES = 5; // q_a, q_b, kv, output_a, output_b -constexpr uint32_t DS4_DMIX_QK = 32; // MIX_QK block width (rocmfp{2,3}_mix.cu) -// The wide load-from-floor staging in mix_block_accum reads a 16 B window per 10/14 B -// block, which only stays in bounds on the final block when the row length is a multiple -// of 128. register_host aborts on violation; checking here gives a message naming the -// tensor instead of a bare abort. -constexpr int64_t DS4_DMIX_ROW_ALIGN = 128; static const char * ds4_dmix_class_name(uint32_t cls) { switch (cls) { @@ -766,12 +942,17 @@ static bool ds4_register_dmix_sidecar(const std::string & gguf_path, } if (n_dense_mix == 0) return true; // uniform dense (101/104/107) -- nothing to do - const std::string sc_path = gguf_path + ".dmix.bin"; - std::vector sc_embedded; // outlives f when the blob came from GGUF KV - FILE * f = ds4_open_sidecar(gguf_path, "deepseek4.dmix.sidecar", ".dmix.bin", sc_embedded); + std::vector embedded_table; // keeps an embedded FILE view alive + const Ds4TableInput table_input = ds4_open_decode_table( + gguf_path, "deepseek4.dmix.sidecar", ".dmix.bin", embedded_table); + FILE * f = table_input.file; if (!f) { - std::fprintf(stderr, "[deepseek4] %d dense mix-qtype attention tensors require " - "sidecar but it is missing: %s\n", n_dense_mix, sc_path.c_str()); + if (!table_input.source.empty()) { + std::fprintf(stderr, + "[deepseek4] decode tables for %d dense mix-qtype attention " + "tensors are missing from GGUF metadata and legacy file %s\n", + n_dense_mix, table_input.source.c_str()); + } return false; } char magic[8]; @@ -780,7 +961,17 @@ static bool ds4_register_dmix_sidecar(const std::string & gguf_path, std::memcmp(magic, "DMIXs1\0\0", 8) != 0 || std::fread(&n_entries, 4, 1, f) != 1 || std::fread(&reserved, 4, 1, f) != 1) { - std::fprintf(stderr, "[deepseek4] bad dmix sidecar header: %s\n", sc_path.c_str()); + std::fprintf(stderr, "[deepseek4] bad dmix table header in %s\n", + table_input.source.c_str()); + std::fclose(f); + return false; + } + const uint64_t max_entries = + (uint64_t) n_layers_out * DS4_DMIX_CLASSES; + if (n_entries == 0 || (uint64_t) n_entries > max_entries) { + std::fprintf(stderr, + "[deepseek4] dmix table has invalid entry count %u\n", + n_entries); std::fclose(f); return false; } @@ -833,13 +1024,6 @@ static bool ds4_register_dmix_sidecar(const std::string & gguf_path, (int) t->type); ok = false; break; } - if (t->ne[0] % DS4_DMIX_ROW_ALIGN != 0) { - std::fprintf(stderr, "[deepseek4] dmix L%u %s row length %lld is not a multiple " - "of %lld; the fused kernel's wide block load would read past the " - "tensor\n", layer, ds4_dmix_class_name(cls), - (long long) t->ne[0], (long long) DS4_DMIX_ROW_ALIGN); - ok = false; break; - } if (t->ne[1] % (int64_t) nslices != 0) { std::fprintf(stderr, "[deepseek4] dmix L%u %s: nslices=%u does not divide " "ne[1]=%lld\n", layer, ds4_dmix_class_name(cls), nslices, @@ -855,6 +1039,19 @@ static bool ds4_register_dmix_sidecar(const std::string & gguf_path, } const size_t nb02 = total_bytes / nslices; + const int64_t out_dim64 = t->ne[1] / (int64_t) nslices; + if (t->ne[0] <= 0 || + t->ne[0] > (int64_t) std::numeric_limits::max() || + out_dim64 <= 0 || + out_dim64 > (int64_t) std::numeric_limits::max()) { + std::fprintf(stderr, + "[deepseek4] dmix L%u %s dimensions are out of range\n", + layer, ds4_dmix_class_name(cls)); + ok = false; + break; + } + const int in_dim = (int) t->ne[0]; + const int out_dim = (int) out_dim64; // Replicate the single per-tensor codebook/mode across slices, so the kernels' // `codebooks + slice*C*K` and `modes[slice]` striding resolves to the same values // for every slice with no dense-specific branch in device code. @@ -863,18 +1060,21 @@ static bool ds4_register_dmix_sidecar(const std::string & gguf_path, std::memcpy(books.data() + (size_t) sl * C * K, book_one.data(), book_one.size() * 2); } - std::vector modes(nslices, mode), rots(nslices, 0); + std::vector modes(nslices, mode); - const int in_dim = (int) t->ne[0]; - const int out_dim = (int) (t->ne[1] / (int64_t) nslices); - if (qtype == (uint32_t) DS4_QTYPE_ROCMFP3_MIX) { - ggml_cuda_rocmfp3_mix_register_host(t->data, nb02, (int) nslices, - out_dim, in_dim, books.data(), - modes.data(), rots.data()); - } else { - ggml_cuda_rocmfp2_mix_register_host(t->data, nb02, (int) nslices, - out_dim, in_dim, books.data(), - modes.data(), rots.data()); + const bool registered = qtype == (uint32_t) DS4_QTYPE_ROCMFP3_MIX + ? ggml_cuda_rocmfp3_mix_register_host( + t->data, nb02, (int) nslices, out_dim, in_dim, + books.data(), modes.data()) + : ggml_cuda_rocmfp2_mix_register_host( + t->data, nb02, (int) nslices, out_dim, in_dim, + books.data(), modes.data()); + if (!registered) { + std::fprintf(stderr, + "[deepseek4] failed to register dmix L%u %s\n", + layer, ds4_dmix_class_name(cls)); + ok = false; + break; } registered_bases.push_back(t->data); covered[layer][cls] = true; @@ -904,19 +1104,18 @@ static bool ds4_register_dmix_sidecar(const std::string & gguf_path, return false; } std::fprintf(stderr, "[deepseek4] dmix: registered %d dense mix-qtype attention " - "tensors from %s\n", n_dense_mix, sc_path.c_str()); + "tensors from %s\n", n_dense_mix, table_input.source.c_str()); return true; } constexpr uint32_t DS4_GUMIX_C = 2; // codebooks per expert -constexpr uint32_t DS4_GUMIX_K = 4; // levels per codebook (2-bit codes) constexpr uint32_t DS4_GUMIX_SURFACES = 3; // 0 = gate, 1 = up, 2 = down -constexpr uint32_t DS4_GUMIX_QK = 32; // MIX_QK block width (rocmfp2_mix.cu) static bool ds4_register_gumix_sidecar(const std::string & gguf_path, const TargetLoadPlan & plan, - DeepSeek4Weights & out) { - // required[layer][surface]: a qtype-106 gate/up expert resident on this shard. + const DeepSeek4Weights & out, + MoeHybridStorage * hybrid = nullptr) { + // required[layer][surface]: a resident qtype-106 expert tensor on this shard. const size_t n_layers_out = out.layers.size(); std::vector> required( n_layers_out, std::array{false, false, false}); @@ -931,20 +1130,25 @@ static bool ds4_register_gumix_sidecar(const std::string & gguf_path, } } } - if (n_qtype106 == 0) return true; // uniform (qtype-107) gate/up — nothing to do + if (n_qtype106 == 0) return true; // uniform (qtype-107) experts — nothing to do - if (plan.skip_expert_tensors) { - std::fprintf(stderr, "[deepseek4] qtype-106 gate/up experts not resident on this " + if (plan.skip_expert_tensors && !hybrid) { + std::fprintf(stderr, "[deepseek4] qtype-106 experts not resident on this " "shard (skip_expert_tensors) — fused decode disabled here\n"); return true; } - const std::string sc_path = gguf_path + ".gumix.bin"; - std::vector sc_embedded; // outlives f when the blob came from GGUF KV - FILE * f = ds4_open_sidecar(gguf_path, "deepseek4.gumix.sidecar", ".gumix.bin", sc_embedded); + std::vector embedded_table; // keeps an embedded FILE view alive + const Ds4TableInput table_input = ds4_open_decode_table( + gguf_path, "deepseek4.gumix.sidecar", ".gumix.bin", embedded_table); + FILE * f = table_input.file; if (!f) { - std::fprintf(stderr, "[deepseek4] qtype-106 gate/up experts require sidecar but it " - "is missing: %s\n", sc_path.c_str()); + if (!table_input.source.empty()) { + std::fprintf(stderr, + "[deepseek4] qtype-106 decode tables are missing from GGUF " + "metadata and legacy file %s\n", + table_input.source.c_str()); + } return false; } char magic[8]; @@ -957,11 +1161,20 @@ static bool ds4_register_gumix_sidecar(const std::string & gguf_path, // and would silently mis-register against the split GGUF, so its magic is // rejected here rather than tolerated. std::fprintf(stderr, "[deepseek4] bad gumix sidecar header (need split-form " - "GUMIXs1): %s\n", sc_path.c_str()); + "GUMIXs1) in %s\n", table_input.source.c_str()); std::fclose(f); return false; } (void) reserved; + const uint64_t max_entries = + (uint64_t) n_layers_out * DS4_GUMIX_SURFACES; + if (n_entries == 0 || (uint64_t) n_entries > max_entries) { + std::fprintf(stderr, + "[deepseek4] gumix table has invalid entry count %u\n", + n_entries); + std::fclose(f); + return false; + } std::vector> done( n_layers_out, std::array{false, false, false}); @@ -1028,12 +1241,13 @@ static bool ds4_register_gumix_sidecar(const std::string & gguf_path, (long long) gt->ne[1], (long long) gt->ne[0]); ok = false; break; } - if (idim % DS4_GUMIX_QK != 0) { + if (idim % DS4_ROCMFP2_ROW_ALIGN != 0) { std::fprintf(stderr, "[deepseek4] gumix layer %u surface %u in=%u not a multiple " - "of %u\n", layer, surface, idim, DS4_GUMIX_QK); + "of %u\n", layer, surface, idim, + DS4_ROCMFP2_ROW_ALIGN); ok = false; break; } - if (!gt->data) { + if (!gt->data && !hybrid) { std::fprintf(stderr, "[deepseek4] gumix layer %u surface %u expert data not " "resident\n", layer, surface); ok = false; break; @@ -1045,7 +1259,6 @@ static bool ds4_register_gumix_sidecar(const std::string & gguf_path, } std::vector modes(E); - std::vector rots(E, 0); // wire carries none; kernel ignores it std::vector books(book_elems); if (std::fread(modes.data(), 1, E, f) != E || std::fread(books.data(), sizeof(uint16_t), books.size(), f) != books.size()) { @@ -1062,10 +1275,26 @@ static bool ds4_register_gumix_sidecar(const std::string & gguf_path, } if (!ok) break; - ggml_cuda_rocmfp2_mix_register_host( - gt->data, gt->nb[2], (int) E, (int) odim, (int) idim, - books.data(), modes.data(), rots.data()); - registered_bases.push_back(gt->data); + if (hybrid) { + const Ds4MixTable table{ + GGML_TYPE_Q2_1_ROCMFP2_MIX, E, odim, idim, + books, modes}; + ok = ds4_register_hybrid_mix_tensors( + *hybrid, layer, surface, table, registered_bases); + } else { + if (!ggml_cuda_rocmfp2_mix_register_host( + gt->data, gt->nb[2], (int) E, (int) odim, (int) idim, + books.data(), modes.data())) { + std::fprintf(stderr, + "[deepseek4] failed to register gumix layer %u " + "surface %u\n", + layer, surface); + ok = false; + break; + } + registered_bases.push_back(gt->data); + } + if (!ok) break; done[layer][surface] = true; } std::fclose(f); @@ -1090,13 +1319,59 @@ static bool ds4_register_gumix_sidecar(const std::string & gguf_path, for (const void * b : registered_bases) ggml_cuda_rocmfp2_mix_unregister(b); return false; } - std::fprintf(stderr, "[deepseek4] registered %d qtype-106 gate/up expert tensor(s) " - "from %s\n", (int) registered_bases.size(), sc_path.c_str()); + std::fprintf(stderr, "[deepseek4] registered %d qtype-106 expert tensor(s) " + "from %s\n", (int) registered_bases.size(), + table_input.source.c_str()); return true; } } // namespace +bool register_deepseek4_moe_hybrid_mix_tables( + const std::string & path, + const DeepSeek4Weights & w, + MoeHybridStorage & storage, + std::string * err) { + bool has_mix_experts = false; + for (const DeepSeek4Layer & layer : w.layers) { + const ggml_tensor * experts[] = { + layer.ffn_gate_exps, layer.ffn_up_exps, layer.ffn_down_exps, + }; + for (const ggml_tensor * tensor : experts) { + has_mix_experts = has_mix_experts || + (tensor && + (tensor->type == GGML_TYPE_Q2_1_ROCMFP2_MIX || + tensor->type == GGML_TYPE_Q3_1_ROCMFP3_MIX)); + } + } + if (!has_mix_experts) return true; + + if (storage.layers.size() != w.layers.size() || + storage.cold_backend_kind != MoeHybridColdBackend::Gpu || + !storage.materialized_hot_experts || + !storage.materialized_cold_experts) { + if (err) *err = "mixed expert qtypes require materialized GPU owners"; + return false; + } + + TargetLoadPlan plan; + plan.skip_expert_tensors = true; + if (!ds4_register_gumix_sidecar(path, plan, w, &storage)) { + storage.unregister_mix_tensors(); + if (err) *err = "failed to register compact qtype-106 expert tensors"; + return false; + } + if (!ds4_register_p4mix_sidecar(path, plan, w, &storage)) { + storage.unregister_mix_tensors(); + if (err) *err = "failed to register compact qtype-105 expert tensors"; + return false; + } + + std::fprintf(stderr, + "[deepseek4] compact mixed-expert decode tables registered\n"); + return true; +} + // Exported with C linkage purely so the unit test can reach the dmix entry rules: they live // in the anonymous namespace above (internal linkage), which is right for the parser but // makes them unlinkable from outside the TU. A thin forwarding wrapper keeps one definition @@ -1629,13 +1904,9 @@ bool load_deepseek4_gguf_partial(const std::string & path, gguf_free(gctx); // Note: meta_ctx is now owned by out.ctx — do NOT free it here. - // qtype-105 mixed-policy down-experts: register per-expert codebooks now - // that tensor data pointers are final (no-op for uniform models). A qtype-105 - // model with a missing/invalid/mismatched sidecar has no working decode path, - // so fail the load here rather than aborting at first decode. - // qtype-106 mixed-policy gate/up experts: same contract as the qtype-105 call - // below, on the other MoE surface. Registered first so a gate/up failure unwinds - // before any down-expert entries exist. + // Register mixed-expert decode tables after tensor data pointers are final. + // A missing or invalid table makes the corresponding qtype undecodable, so + // fail the load here instead of crashing on the first request. // Dense (attention) mix-qtype tensors first: a dense failure unwinds before any MoE // entries exist, keeping the teardown order the reverse of registration. if (!ds4_register_dmix_sidecar(path, out)) { @@ -1646,7 +1917,7 @@ bool load_deepseek4_gguf_partial(const std::string & path, } if (!ds4_register_gumix_sidecar(path, plan, out)) { - std::fprintf(stderr, "[deepseek4] qtype-106 gate/up sidecar registration failed " + std::fprintf(stderr, "[deepseek4] qtype-106 sidecar registration failed " "for %s\n", path.c_str()); free_deepseek4_weights(out); return false; @@ -1892,27 +2163,17 @@ bool build_deepseek4_moe_hybrid_storage_from_file( void free_deepseek4_weights(DeepSeek4Weights & w) { deepseek4_release_runtime_graphs(w); - // Drop qtype-105 registry entries (and free their device side-data) while the - // tensors are still valid — BEFORE ggml_free destroys them and the GPU buffer - // is released. Otherwise the entries leak device memory across park/unpark and - // reload, and a later load reusing the same device address would resolve to a - // stale range. Keyed by dt->data, so read it here. + // Drop expert registry entries before their GPU buffers. Otherwise reloads can + // resolve a reused device address to stale decode tables. for (auto & L : w.layers) { - ggml_tensor * dt = L.ffn_down_exps; - if (dt && dt->type == GGML_TYPE_Q3_1_ROCMFP3_MIX && dt->data) { - ggml_cuda_rocmfp3_mix_unregister(dt->data); - } - } - // Same contract for qtype-106 gate/up, and it needs both halves: the gumix - // sidecar registers the SPLIT form, one entry per tensor, and each entry owns - // its own device copy of the (shared) codebooks — duplicated deliberately so - // that unregister can free per entry without double-freeing. Dropping only one - // half would still leak the other's codebooks and leave a stale range behind. - for (auto & L : w.layers) { - ggml_tensor * const gu[3] = { L.ffn_gate_exps, L.ffn_up_exps, - L.ffn_down_exps }; - for (ggml_tensor * t : gu) { - if (t && t->type == GGML_TYPE_Q2_1_ROCMFP2_MIX && t->data) { + ggml_tensor * const experts[] = { + L.ffn_gate_exps, L.ffn_up_exps, L.ffn_down_exps, + }; + for (ggml_tensor * t : experts) { + if (!t || !t->data) continue; + if (t->type == GGML_TYPE_Q3_1_ROCMFP3_MIX) { + ggml_cuda_rocmfp3_mix_unregister(t->data); + } else if (t->type == GGML_TYPE_Q2_1_ROCMFP2_MIX) { ggml_cuda_rocmfp2_mix_unregister(t->data); } } diff --git a/server/test/bench_rocmfp_mix_gateup_glu.cpp b/server/test/bench_rocmfp_mix_gateup_glu.cpp index 26264fac0..b7fdf2a5f 100644 --- a/server/test/bench_rocmfp_mix_gateup_glu.cpp +++ b/server/test/bench_rocmfp_mix_gateup_glu.cpp @@ -15,6 +15,8 @@ // relative delta. Single stream and back-to-back launches on purpose: that is the decode // dependency chain, where launch overhead cannot be hidden. +#include "ggml-cuda.h" + #include #include @@ -24,12 +26,6 @@ #include #include -extern "C" void ggml_cuda_rocmfp2_mix_register_host( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host); -extern "C" void ggml_cuda_rocmfp2_mix_unregister(const void * base); - bool ggml_cuda_rocmfp2_mix_mul_mat_id( const void * vx, const float * src1, const int32_t * ids, float * dst, int in, int out, int n_expert_used, int n_tokens, int ne11, @@ -83,7 +79,7 @@ int main() { } std::vector books((size_t) n_experts * 2 * K); for (size_t i = 0; i < books.size(); ++i) books[i] = f32_to_bf16(-0.5f + 0.2f * (float) (i % 5)); - std::vector modes(n_experts, 1), rots(n_experts, 0); + std::vector modes(n_experts, 1); void * d_up = nullptr, * d_gate = nullptr; float * d_x = nullptr, * d_a = nullptr, * d_b = nullptr; @@ -104,10 +100,19 @@ int main() { for (int i = 0; i < n_used * ntok; ++i) idsh[i] = i % n_experts; HIP_OK(hipMemcpy(d_ids, idsh.data(), sizeof(int32_t) * idsh.size(), hipMemcpyHostToDevice)); - ggml_cuda_rocmfp2_mix_register_host(d_up, rows_bytes, n_experts, out, in, - books.data(), modes.data(), rots.data()); - ggml_cuda_rocmfp2_mix_register_host(d_gate, rows_bytes, n_experts, out, in, - books.data(), modes.data(), rots.data()); + if (!ggml_cuda_rocmfp2_mix_register_host( + d_up, rows_bytes, n_experts, out, in, + books.data(), modes.data())) { + std::fprintf(stderr, "FAIL: mixed-tensor registration failed\n"); + return 1; + } + if (!ggml_cuda_rocmfp2_mix_register_host( + d_gate, rows_bytes, n_experts, out, in, + books.data(), modes.data())) { + ggml_cuda_rocmfp2_mix_unregister(d_up); + std::fprintf(stderr, "FAIL: mixed-tensor registration failed\n"); + return 1; + } const int64_t ids_s0 = 1, ids_s1 = n_used; const int64_t src1_s1 = 0, src1_s2 = in; diff --git a/server/test/test_ds4_mix_registry_teardown.cpp b/server/test/test_ds4_mix_registry_teardown.cpp index fa6a70f27..07d9bfa2f 100644 --- a/server/test/test_ds4_mix_registry_teardown.cpp +++ b/server/test/test_ds4_mix_registry_teardown.cpp @@ -23,6 +23,8 @@ // on them and never dereferences, exactly as test_rocmfp3_mix_registry relies on. #include "deepseek4_internal.h" +#include "common/moe_hybrid_storage.h" +#include "ggml-cuda.h" #include #include @@ -30,16 +32,9 @@ using dflash::common::DeepSeek4Layer; using dflash::common::DeepSeek4Weights; +using dflash::common::MoeHybridStorage; using dflash::common::free_deepseek4_weights; -extern "C" void ggml_cuda_rocmfp3_mix_register_host( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host); -extern "C" void ggml_cuda_rocmfp2_mix_register_host( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host); bool ggml_cuda_rocmfp3_mix_registered(const void * vx); bool ggml_cuda_rocmfp2_mix_registered(const void * vx); @@ -66,17 +61,16 @@ constexpr int OUT = 32; constexpr int IN = 128; constexpr size_t NB02 = 4096; -void register_as(bool is105, const void * base) { +bool register_as(bool is105, const void * base) { const int K = is105 ? 8 : 4; std::vector books((size_t) E * 2 * (size_t) K, 0x3f80); // bf16 ~1.0 - std::vector modes(E, 1), rots(E, 0); + std::vector modes(E, 1); if (is105) { - ggml_cuda_rocmfp3_mix_register_host(base, NB02, E, OUT, IN, - books.data(), modes.data(), rots.data()); - } else { - ggml_cuda_rocmfp2_mix_register_host(base, NB02, E, OUT, IN, - books.data(), modes.data(), rots.data()); + return ggml_cuda_rocmfp3_mix_register_host( + base, NB02, E, OUT, IN, books.data(), modes.data()); } + return ggml_cuda_rocmfp2_mix_register_host( + base, NB02, E, OUT, IN, books.data(), modes.data()); } } // namespace @@ -122,7 +116,7 @@ int main() { if (!t) { std::fprintf(stderr, "SKIP: tensor alloc failed\n"); return 0; } t->data = (void *) fake_base(slot++); const bool is105 = (t->type == GGML_TYPE_Q3_1_ROCMFP3_MIX); - register_as(is105, t->data); + CHECK(register_as(is105, t->data), "dense registration succeeds"); (is105 ? dense105 : dense106).push_back(t->data); } @@ -135,7 +129,7 @@ int main() { if (!t) { std::fprintf(stderr, "SKIP: tensor alloc failed\n"); return 0; } t->data = (void *) fake_base(slot++); const bool is105 = (t->type == GGML_TYPE_Q3_1_ROCMFP3_MIX); - register_as(is105, t->data); + CHECK(register_as(is105, t->data), "expert registration succeeds"); (is105 ? expert105 : expert106).push_back(t->data); } w.layers.push_back(L); @@ -170,6 +164,67 @@ int main() { CHECK(stale_dense == 0, "NO dense attention mix entry survives free_deepseek4_weights"); CHECK(stale_expert == 0, "no expert mix entry survives free_deepseek4_weights"); + // Hybrid loading owns separate compact tensors on the primary and + // secondary GPUs. MoeHybridStorage must unregister all of them before it + // releases either owner's buffers. + std::vector hybrid105, hybrid106; + { + ggml_context * hybrid_ctx = ggml_init(ip); + if (!hybrid_ctx) { + std::fprintf(stderr, "SKIP: hybrid ggml_init failed\n"); + return g_fails ? 1 : 0; + } + + MoeHybridStorage hybrid; + hybrid.layers.resize(1); + auto & layer = hybrid.layers[0]; + layer.hot_ctx = hybrid_ctx; + layer.gate_hot = ggml_new_tensor_2d( + hybrid_ctx, GGML_TYPE_Q2_1_ROCMFP2_MIX, IN, OUT); + layer.down_hot = ggml_new_tensor_2d( + hybrid_ctx, GGML_TYPE_Q3_1_ROCMFP3_MIX, IN, OUT); + layer.gate_cold = ggml_new_tensor_2d( + hybrid_ctx, GGML_TYPE_Q2_1_ROCMFP2_MIX, IN, OUT); + layer.down_cold = ggml_new_tensor_2d( + hybrid_ctx, GGML_TYPE_Q3_1_ROCMFP3_MIX, IN, OUT); + + ggml_tensor * compact[] = { + layer.gate_hot, layer.down_hot, layer.gate_cold, layer.down_cold, + }; + for (ggml_tensor * tensor : compact) { + if (!tensor) { + std::fprintf(stderr, "SKIP: hybrid tensor alloc failed\n"); + return g_fails ? 1 : 0; + } + tensor->data = (void *) fake_base(slot++); + const bool is105 = + tensor->type == GGML_TYPE_Q3_1_ROCMFP3_MIX; + CHECK(register_as(is105, tensor->data), + "hybrid registration succeeds"); + (is105 ? hybrid105 : hybrid106).push_back(tensor->data); + } + + size_t live_hybrid = 0; + for (const void * base : hybrid105) { + live_hybrid += ggml_cuda_rocmfp3_mix_registered(base) ? 1 : 0; + } + for (const void * base : hybrid106) { + live_hybrid += ggml_cuda_rocmfp2_mix_registered(base) ? 1 : 0; + } + CHECK(live_hybrid == hybrid105.size() + hybrid106.size(), + "every compact owner registration resolves before teardown"); + } + + size_t stale_hybrid = 0; + for (const void * base : hybrid105) { + stale_hybrid += ggml_cuda_rocmfp3_mix_registered(base) ? 1 : 0; + } + for (const void * base : hybrid106) { + stale_hybrid += ggml_cuda_rocmfp2_mix_registered(base) ? 1 : 0; + } + CHECK(stale_hybrid == 0, + "no compact mix entry survives MoeHybridStorage teardown"); + std::fprintf(stderr, g_fails ? "MIX REGISTRY TEARDOWN TEST FAILED (%d)\n" : "MIX REGISTRY TEARDOWN TEST OK\n", g_fails); return g_fails ? 1 : 0; diff --git a/server/test/test_moe_hybrid_storage.cpp b/server/test/test_moe_hybrid_storage.cpp index b09948e3d..b57ceb15a 100644 --- a/server/test/test_moe_hybrid_storage.cpp +++ b/server/test/test_moe_hybrid_storage.cpp @@ -1,7 +1,9 @@ #include "CppUnitTestFramework.hpp" +#include "../src/common/moe_hybrid_ffn_eval.h" #include "../src/common/moe_hybrid_storage.h" #include +#include #include using namespace dflash::common; @@ -35,3 +37,42 @@ TEST_CASE(MoeHybridStorageFixture, expert_residency_tracks_model_sized_expert_se REQUIRE(!storage.is_expert_hot(256)); REQUIRE(!storage.all_routed_are_hot(all_hot.data(), (int)all_hot.size())); } + +TEST_CASE(MoeHybridStorageFixture, heterogeneous_route_balance_scales_with_model_top_k) { + REQUIRE(moe_balanced_main_slots_x4(4, 4.4) == 13); + REQUIRE(moe_balanced_main_slots_x4(6, 4.4) == 20); + REQUIRE(moe_balanced_main_slots_x4(0, 4.4) == 0); + REQUIRE(moe_balanced_main_slots_x4(6, 0.0) == 0); +} + +TEST_CASE(MoeHybridStorageFixture, fractional_route_quota_rounds_over_the_batch) { + ggml_init_params params{ + /*mem_size=*/1024 * 1024, + /*mem_buffer=*/nullptr, + /*no_alloc=*/true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 6, 5); + ggml_tensor * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 6, 5); + ggml_tensor * local_lut = + ggml_new_tensor_4d(ctx, GGML_TYPE_I32, 1, 8, 5, 1); + ggml_tensor * candidate_lut = + ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 8, 5, 1); + REQUIRE(ids && weights && local_lut && candidate_lut); + + // top-k 6 at a 3:1 owner rate is 4.5 main routes per token. Across a + // five-token verifier batch the exact quota is 22.5, which rounds to 23. + ggml_tensor * owner_ids = ggml_ds4_moe_balanced_owner_ids( + ctx, ids, weights, local_lut, candidate_lut, + /*main_slots_x4=*/18, /*main_owner=*/true); + REQUIRE(owner_ids != nullptr); + int32_t main_quota = 0; + std::memcpy(&main_quota, + owner_ids->op_params + sizeof(int32_t), + sizeof(main_quota)); + REQUIRE(main_quota == 23); + + ggml_free(ctx); +} diff --git a/server/test/test_rocmfp3_mix_registry.cpp b/server/test/test_rocmfp3_mix_registry.cpp index 8cd9af5ad..09947982b 100644 --- a/server/test/test_rocmfp3_mix_registry.cpp +++ b/server/test/test_rocmfp3_mix_registry.cpp @@ -4,25 +4,21 @@ // - register_host makes a resolvable entry; range lookup is correct; // - unregister removes it (no stale base range survives an "unload"); // - update-in-place and repeated register/unregister cycles do not leak the -// device side-data buffers (codebooks/modes/rotations). +// device side-data buffers (codebooks/modes). // The pre-fix code (unregister only erased the vector entry, never cudaFree'd // the register_host allocations) fails the leak assertion below. #include "ds4_test_gpu_runtime.h" +#include "ggml-cuda.h" +#include "rocmfp3_mix.cuh" +#include +#include #include #include +#include #include -// Registry entry points (register_host/unregister are extern "C"; registered is -// C++ linkage — declared to match the definitions in rocmfp3_mix.cu). -extern "C" void ggml_cuda_rocmfp3_mix_register_host( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host); -extern "C" void ggml_cuda_rocmfp3_mix_unregister(const void * base); -bool ggml_cuda_rocmfp3_mix_registered(const void * vx); - static int g_fails = 0; #define CHECK(cond, msg) \ do { \ @@ -31,9 +27,10 @@ static int g_fails = 0; int main() { const int E = 8, out = 64, in = 64; - const size_t nb02 = 4096; // per-expert byte stride (opaque to the registry) + const size_t expert_bytes = (size_t) out * (in / 32) * 14; + const size_t nb02 = 4096; // includes alignment padding after each payload std::vector books((size_t) E * 2 * 8, 0x3f80); // bf16 ~1.0 - std::vector modes(E, 1), rots(E, 0); + std::vector modes(E, 1); // registered() only does pointer-range arithmetic on the base key — it never // dereferences it — so opaque, distinct, aligned values stand in for two @@ -41,45 +38,109 @@ int main() { const void * b0 = reinterpret_cast(0x100000000ull); const void * b1 = reinterpret_cast(0x200000000ull); - // 1. register + range lookup - ggml_cuda_rocmfp3_mix_register_host(b0, nb02, E, out, in, - books.data(), modes.data(), rots.data()); + // 1. Invalid metadata fails without aborting or leaving a registry entry. + CHECK(!ggml_cuda_rocmfp3_mix_register_host( + nullptr, nb02, E, out, in, books.data(), modes.data()), + "null registration base is rejected"); + CHECK(!ggml_cuda_rocmfp3_mix_register_host( + b1, expert_bytes - 1, E, out, in, books.data(), modes.data()), + "undersized expert stride is rejected"); + std::vector invalid_modes = modes; + invalid_modes[0] = 2; + CHECK(!ggml_cuda_rocmfp3_mix_register_host( + b1, nb02, E, out, in, books.data(), invalid_modes.data()), + "unsupported mode is rejected"); + CHECK(!ggml_cuda_rocmfp3_mix_registered(b1), + "invalid registration leaves no entry"); + + // 2. register + range lookup + CHECK(ggml_cuda_rocmfp3_mix_register_host( + b0, nb02, E, out, in, books.data(), modes.data()), + "valid registration succeeds"); CHECK(ggml_cuda_rocmfp3_mix_registered(b0), "b0 resolves after register"); CHECK(ggml_cuda_rocmfp3_mix_registered(static_cast(b0) + nb02), "expert-1 slice resolves (in range)"); + CHECK(!ggml_cuda_rocmfp3_mix_registered(static_cast(b0) + 14), + "an interior block is not a registered tensor base"); + CHECK(!ggml_cuda_rocmfp3_mix_registered( + static_cast(b0) + expert_bytes), + "padding after an expert payload does not resolve"); CHECK(!ggml_cuda_rocmfp3_mix_registered(static_cast(b0) + (size_t) E * nb02), "just past the last expert does not resolve"); CHECK(!ggml_cuda_rocmfp3_mix_registered(b1), "unrelated base does not resolve"); + const void * codebooks = nullptr; + const uint8_t * registered_modes = nullptr; + CHECK(ggml_cuda_rocmfp3_mix_mmq_info( + static_cast(b0) + 14, + &codebooks, ®istered_modes), + "MMQ accepts a block-aligned offset inside an expert"); + CHECK(codebooks != nullptr && registered_modes != nullptr, + "MMQ returns the registered side data"); + CHECK(!ggml_cuda_rocmfp3_mix_mmq_info( + static_cast(b0) + 1, + &codebooks, ®istered_modes), + "MMQ rejects an unaligned offset inside an expert"); - // 2. unregister leaves no stale range (the reload/address-reuse hazard) + // 3. unregister leaves no stale range (the reload/address-reuse hazard) ggml_cuda_rocmfp3_mix_unregister(b0); CHECK(!ggml_cuda_rocmfp3_mix_registered(b0), "b0 gone after unregister"); - // 3. update-in-place then unregister - ggml_cuda_rocmfp3_mix_register_host(b0, nb02, E, out, in, - books.data(), modes.data(), rots.data()); - ggml_cuda_rocmfp3_mix_register_host(b0, nb02, E, out, in, // update same base - books.data(), modes.data(), rots.data()); + // 4. update-in-place then unregister + CHECK(ggml_cuda_rocmfp3_mix_register_host( + b0, nb02, E, out, in, books.data(), modes.data()), + "registration before update succeeds"); + CHECK(ggml_cuda_rocmfp3_mix_register_host( + b0, nb02, E, out, in, books.data(), modes.data()), + "in-place registration update succeeds"); CHECK(ggml_cuda_rocmfp3_mix_registered(b0), "b0 resolves after in-place update"); ggml_cuda_rocmfp3_mix_unregister(b0); CHECK(!ggml_cuda_rocmfp3_mix_registered(b0), "b0 gone after update+unregister"); - // 4. no device-memory leak across many register/unregister cycles. A missing - // cudaFree in unregister (or on update) leaks ~E*(2*8*2 + 1 + 1) bytes per - // cycle; 4000 cycles would drop free VRAM by tens of MB. + // 5. Teardown cannot free side data between lookup and asynchronous kernel + // enqueue. Dispatchers hold this lock across both operations; unregister + // must wait until the launch has been handed to the device. + CHECK(ggml_cuda_rocmfp3_mix_register_host( + b0, nb02, E, out, in, books.data(), modes.data()), + "registration before dispatch lock test succeeds"); + std::atomic teardown_started{false}; + std::atomic teardown_finished{false}; + ggml_cuda_rocmfp3_mix_registry_lock(); + std::thread teardown([&] { + teardown_started.store(true, std::memory_order_release); + ggml_cuda_rocmfp3_mix_unregister(b0); + teardown_finished.store(true, std::memory_order_release); + }); + while (!teardown_started.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + CHECK(!teardown_finished.load(std::memory_order_acquire), + "unregister waits for an in-flight dispatch"); + ggml_cuda_rocmfp3_mix_registry_unlock(); + teardown.join(); + CHECK(teardown_finished.load(std::memory_order_acquire), + "unregister completes after dispatch releases the registry"); + CHECK(!ggml_cuda_rocmfp3_mix_registered(b0), + "dispatch lock test leaves no registry entry"); + + // 6. no device-memory leak across many register/unregister cycles. A missing + // cudaFree in unregister (or on update) leaks ~E*(2*8*2 + 1) bytes per + // cycle; 4000 cycles would produce a measurable free-VRAM drop. cudaDeviceSynchronize(); size_t free_warm = 0, total = 0; // warm the allocator first so pool growth isn't counted as a leak for (int i = 0; i < 64; ++i) { - ggml_cuda_rocmfp3_mix_register_host(b0, nb02, E, out, in, - books.data(), modes.data(), rots.data()); + CHECK(ggml_cuda_rocmfp3_mix_register_host( + b0, nb02, E, out, in, books.data(), modes.data()), + "warmup registration succeeds"); ggml_cuda_rocmfp3_mix_unregister(b0); } cudaDeviceSynchronize(); (void) cudaMemGetInfo(&free_warm, &total); for (int i = 0; i < 4000; ++i) { - ggml_cuda_rocmfp3_mix_register_host(b0, nb02, E, out, in, - books.data(), modes.data(), rots.data()); + CHECK(ggml_cuda_rocmfp3_mix_register_host( + b0, nb02, E, out, in, books.data(), modes.data()), + "cycle registration succeeds"); ggml_cuda_rocmfp3_mix_unregister(b0); } cudaDeviceSynchronize(); diff --git a/server/test/test_rocmfp_mix_gateup_glu.cpp b/server/test/test_rocmfp_mix_gateup_glu.cpp index 3ebe4da32..9f1b35a5e 100644 --- a/server/test/test_rocmfp_mix_gateup_glu.cpp +++ b/server/test/test_rocmfp_mix_gateup_glu.cpp @@ -33,6 +33,7 @@ // tensor with an undecoded one. #include "ds4_test_gpu_runtime.h" +#include "ggml-cuda.h" #include #include @@ -40,12 +41,6 @@ #include #include -extern "C" void ggml_cuda_rocmfp2_mix_register_host( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host); -extern "C" void ggml_cuda_rocmfp2_mix_unregister(const void * base); - bool ggml_cuda_rocmfp2_mix_mul_mat_id( const void * vx, const float * src1, const int32_t * ids, float * dst, int in, int out, int n_expert_used, int n_tokens, int ne11, @@ -144,7 +139,6 @@ int main() { books_gate[i] = f32_to_bf16( 0.5f - 0.21f * (float) (i % 5)); } std::vector modes_up(n_experts, 1), modes_gate(n_experts, 1); // 1 = adaptive - std::vector rots(n_experts, 0); void * d_up = nullptr, * d_gate = nullptr; float * d_x = nullptr, * d_up_out = nullptr, * d_gate_out = nullptr, * d_fused = nullptr; @@ -170,10 +164,15 @@ int main() { for (size_t i = 0; i < idsh.size(); ++i) idsh[i] = (int32_t) ((i * 2 + 1) % n_experts); HIP_OK(cudaMemcpy(d_ids, idsh.data(), sizeof(int32_t) * idsh.size(), cudaMemcpyHostToDevice)); - ggml_cuda_rocmfp2_mix_register_host(d_up, rows_bytes, n_experts, out, in, - books_up.data(), modes_up.data(), rots.data()); - ggml_cuda_rocmfp2_mix_register_host(d_gate, rows_bytes, n_experts, out, in, - books_gate.data(), modes_gate.data(), rots.data()); + CHECK(!ggml_cuda_rocmfp2_mix_register_host( + d_up, rows_bytes, n_experts, out, in - 32, + books_up.data(), modes_up.data())); + CHECK(ggml_cuda_rocmfp2_mix_register_host( + d_up, rows_bytes, n_experts, out, in, + books_up.data(), modes_up.data())); + CHECK(ggml_cuda_rocmfp2_mix_register_host( + d_gate, rows_bytes, n_experts, out, in, + books_gate.data(), modes_gate.data())); const int64_t ids_s0 = 1, ids_s1 = n_used; const int64_t src1_s1 = 0, src1_s2 = in; // ne11 == 1 -> slot broadcast @@ -284,8 +283,9 @@ int main() { limit, nullptr)); // Re-register with a DIFFERENT out: a shape-mismatched pair must be refused too, because the // grid is sized from one half and would index past the other. - ggml_cuda_rocmfp2_mix_register_host(d_gate, rows_bytes, n_experts, out / 2, in, - books_gate.data(), modes_gate.data(), rots.data()); + CHECK(ggml_cuda_rocmfp2_mix_register_host( + d_gate, rows_bytes, n_experts, out / 2, in, + books_gate.data(), modes_gate.data())); CHECK(!ggml_cuda_rocmfp2_mix_mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_fused, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, diff --git a/server/test/test_rocmfp_mix_slice_matvec.cpp b/server/test/test_rocmfp_mix_slice_matvec.cpp index 178ae1368..19e33e71e 100644 --- a/server/test/test_rocmfp_mix_slice_matvec.cpp +++ b/server/test/test_rocmfp_mix_slice_matvec.cpp @@ -21,6 +21,7 @@ // neighbouring tensor's codebook. That failure mode is wrong numbers, not a crash. #include "ds4_test_gpu_runtime.h" +#include "ggml-cuda.h" #include #include @@ -28,12 +29,6 @@ #include #include -extern "C" void ggml_cuda_rocmfp3_mix_register_host( - const void * base, size_t nb02, int n_experts, int out, int in, - const void * codebooks_bf16_host, const uint8_t * modes_host, - const uint8_t * rotations_host); -extern "C" void ggml_cuda_rocmfp3_mix_unregister(const void * base); - bool ggml_cuda_rocmfp3_mix_mul_mat_vec( const void * vx, const float * x, float * y, int in, int out, int ncols, @@ -114,7 +109,6 @@ int main() { books[i] = (uint16_t) (bits >> 16); // fp32 -> bf16 (truncate) } std::vector modes(nslices, 1); // 1 = adaptive: exercises the codebook path - std::vector rots(nslices, 0); uint8_t * d_w = nullptr; float * d_x = nullptr; @@ -139,18 +133,31 @@ int main() { const int64_t d_tok = out, d_sl = (int64_t) out * ntokens; // --- 1. the guard: too few registered slices must be REFUSED, not decoded --- - ggml_cuda_rocmfp3_mix_register_host(d_w, slice_bytes, /*n_experts=*/nslices - 1, - out, in, books.data(), modes.data(), rots.data()); + CHECK(ggml_cuda_rocmfp3_mix_register_host( + d_w, slice_bytes, /*n_experts=*/nslices - 1, + out, in, books.data(), modes.data()), + "partial slice registration succeeds"); CHECK(!ggml_cuda_rocmfp3_mix_mul_mat_vec_3d( d_w, d_x, d_y3, in, out, nslices, ntokens, s1_tok, s1_sl, d_tok, d_sl, nullptr), "3d matvec refuses a tensor registered with fewer slices than the grid indexes"); ggml_cuda_rocmfp3_mix_unregister(d_w); - // --- 2. correctness: 3-D launch vs the validated 2-D entry point, per slice --- - ggml_cuda_rocmfp3_mix_register_host(d_w, slice_bytes, nslices, out, in, - books.data(), modes.data(), rots.data()); + // --- 2. registered metadata must match the requested launch shape ----------- + CHECK(ggml_cuda_rocmfp3_mix_register_host( + d_w, slice_bytes, nslices, out, in, + books.data(), modes.data()), + "full slice registration succeeds"); + CHECK(!ggml_cuda_rocmfp3_mix_mul_mat_vec_3d( + d_w, d_x, d_y3, in, out + 1, nslices, ntokens, + s1_tok, s1_sl, d_tok, d_sl, nullptr), + "3d matvec refuses a shape that differs from its registration"); + CHECK(!ggml_cuda_rocmfp3_mix_mul_mat_vec( + d_w + BLOCK_BYTES, d_x, d_y2, in, out, 1, + s1_tok, d_tok, nullptr), + "2d matvec refuses an interior block pointer as a full slice"); + // --- 3. correctness: 3-D launch vs the validated 2-D entry point, per slice --- CHECK(ggml_cuda_rocmfp3_mix_mul_mat_vec_3d( d_w, d_x, d_y3, in, out, nslices, ntokens, s1_tok, s1_sl, d_tok, d_sl, nullptr),