diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 99550d2d6..d79af1092 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -59,6 +59,12 @@ set_property(CACHE LUCE_GPU_BACKEND PROPERTY STRINGS cuda hip) option(LUCE_ENABLE_MIXED_CUDA_HIP "Primary GPU runtime plus an isolated in-process CUDA/HIP peer (Linux only)" OFF) +# Expert-parallel cluster (after maikzz32/lucebox-halo-cluster): one +# luce_server rank per box, routed experts split across the ranks, one RCCL +# all-reduce per MoE layer over RoCE. The GPU-free cluster sources (config, +# protocol, control channel, placement, decisions) are always compiled; ON +# adds the RCCL communicator. +option(LUCE_CLUSTER "Multi-box RCCL expert-parallel cluster (HIP only)" OFF) string(TOLOWER "${LUCE_GPU_BACKEND}" LUCE_GPU_BACKEND) if(LUCE_ENABLE_MIXED_CUDA_HIP) # Reject unsupported targets before project() tries to discover both GPU @@ -199,6 +205,9 @@ endif() # No BLAS, no Metal, no Vulkan, no examples/tests/tools. if(LUCE_GPU_BACKEND STREQUAL "cuda") + if(LUCE_CLUSTER) + message(FATAL_ERROR "LUCE_CLUSTER requires LUCE_GPU_BACKEND=hip (RCCL)") + endif() set(GGML_CUDA ON CACHE BOOL "" FORCE) if(LUCE_ENABLE_MIXED_CUDA_HIP) set(GGML_HIP ON CACHE BOOL "" FORCE) @@ -234,7 +243,11 @@ elseif(LUCE_GPU_BACKEND STREQUAL "hip") set(GGML_CUDA_MODULE OFF CACHE BOOL "" FORCE) endif() set(GGML_HIP ON CACHE BOOL "" FORCE) - set(GGML_HIP_RCCL OFF CACHE BOOL "" FORCE) + if(LUCE_CLUSTER) + set(GGML_HIP_RCCL ON CACHE BOOL "" FORCE) + else() + set(GGML_HIP_RCCL OFF CACHE BOOL "" FORCE) + endif() set(GGML_HIP_GRAPHS OFF CACHE BOOL "Enable HIP graph replay (held off: first-capture quality variance, see experiment log E30/T12)") set(LUCE_GGML_BACKEND_TARGET ggml-hip) set(LUCE_HIP_ARCHITECTURES "" CACHE STRING "HIP GPU targets, e.g. gfx906;gfx1100") @@ -279,6 +292,31 @@ if(LUCE_ROCMFP2_AFFINE) add_compile_definitions(ROCMFP2_AFFINE=1) endif() +# Ranks exchange this build identity in the control-channel Hello and refuse +# to form a cluster from different binaries. +if(LUCE_CLUSTER) + add_compile_definitions(LUCE_CLUSTER=1 LUCE_CLUSTER_RCCL=1) + find_package(Git QUIET) + set(_luce_build_sha "unknown") + if(GIT_FOUND) + execute_process( + COMMAND ${GIT_EXECUTABLE} describe --always --dirty --tags + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + OUTPUT_VARIABLE _luce_build_sha_out + RESULT_VARIABLE _luce_build_sha_rc + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET) + if(_luce_build_sha_rc EQUAL 0 AND _luce_build_sha_out) + set(_luce_build_sha "${_luce_build_sha_out}") + endif() + unset(_luce_build_sha_out) + unset(_luce_build_sha_rc) + endif() + add_compile_definitions(LUCE_BUILD_SHA="${_luce_build_sha}") + message(STATUS "luce cluster: LUCE_CLUSTER=ON (build sha ${_luce_build_sha})") + unset(_luce_build_sha) +endif() + # HIP Phase 2 — opt-in rocWMMA flashprefill kernels (Strix Halo / gfx1151). # Default OFF (Phase 1 = ggml q8 fallback). Set ON for the 1.7-2.7× compress # speedup at 8K-32K context. Requires rocwmma headers installed under @@ -582,8 +620,10 @@ add_library(luce_common STATIC src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp src/deepseek4/deepseek4_dspark.cpp src/deepseek4/deepseek4_dspark_spec.cpp + src/deepseek4/deepseek4_cluster.cpp src/deepseek4/deepseek4_vision.cpp src/common/gpu_page_pool.cpp + src/common/host_stream_flag.cpp src/common/vision/image_decode.cpp src/common/vision/image_resize.cpp src/common/vision/mmproj_file.cpp @@ -671,6 +711,18 @@ add_library(luce_common STATIC src/common/feature_gate.cpp src/placement/placement_config.cpp src/placement/device_select.cpp + # Expert-parallel cluster runtime. GPU-free parts are always compiled; + # cluster_comm.cpp's RCCL half and fast_reduce.cpp's RDMA half are under + # LUCE_CLUSTER. + src/cluster/cluster_config.cpp + src/cluster/cluster_protocol.cpp + src/cluster/cluster_control.cpp + src/cluster/cluster_comm.cpp + src/cluster/fast_reduce.cpp + src/cluster/cluster_expert_placement.cpp + src/cluster/cluster_decision_hooks.cpp + src/cluster/cluster_head_backend.cpp + src/cluster/cluster_worker_main.cpp src/common/layer_split_utils.cpp src/common/ddtree.cpp src/common/peer_access.cpp @@ -748,6 +800,19 @@ elseif(LUCE_GPU_BACKEND STREQUAL "hip") # hip_compat shim is needed by ALL luce_common sources (peer_access.cpp, # dflash_feature_ring.cpp, flashprefill.cpp), not just the SM80_EQUIV path. target_include_directories(luce_common PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) + if(LUCE_CLUSTER) + # RCCL communicator plus the f32<->bf16 kernels behind + # allreduce_sum_bf16_compressed. + find_package(rccl REQUIRED) + target_sources(luce_common PRIVATE src/cluster/cluster_comm_kernels.cu) + set_source_files_properties(src/cluster/cluster_comm_kernels.cu PROPERTIES LANGUAGE HIP) + target_link_libraries(luce_common PUBLIC roc::rccl) + # Lean decode all-reduce over raw RDMA (stream wait/write flags, no RCCL kernel). + find_library(LUCE_IBVERBS_LIB ibverbs REQUIRED) + target_sources(luce_common PRIVATE src/cluster/fast_reduce_kernels.cu) + set_source_files_properties(src/cluster/fast_reduce_kernels.cu PROPERTIES LANGUAGE HIP) + target_link_libraries(luce_common PUBLIC ${LUCE_IBVERBS_LIB}) + endif() endif() if(LUCE_GPU_BACKEND STREQUAL "cuda") @@ -1075,6 +1140,7 @@ if(LUCE_TESTS) test_ds4_mix_fitter) set(_luce_cpu_cppunit_targets test_feature_gate + test_cluster_unit test_platform_compat test_paged_kv_pool test_blocking_row_pool @@ -1375,6 +1441,12 @@ if(LUCE_TESTS) target_include_directories(test_qwen35_tensor_parallel PRIVATE ${LUCE_SRC_INCLUDE_DIRS}) target_link_libraries(test_qwen35_tensor_parallel PRIVATE luce_common) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench/bench_ds41_q8_mmvq.cpp") + # Benchmark, not a test: DS4.1 dense Q8_0 projections through the ggml + # mul_mat dispatch at 1-4 columns, with the column-identity check. + add_executable(bench_ds41_q8_mmvq test/bench/bench_ds41_q8_mmvq.cpp) + target_link_libraries(bench_ds41_q8_mmvq PRIVATE ggml ggml-base ${LUCE_GGML_BACKEND_TARGET}) + endif() add_executable(test_rope_tail test/unit/test_rope_tail.cpp) target_compile_definitions(test_rope_tail PRIVATE LUCE_ROPE_TEST_GPU=1) target_link_libraries(test_rope_tail PRIVATE ggml ggml-base ggml-cpu ${LUCE_GGML_BACKEND_TARGET}) @@ -1683,6 +1755,27 @@ if(LUCE_TESTS) target_link_libraries(test_ggml_rmsnorm_batch PRIVATE CUDA::cudart) endif() endif() + # Cluster unit tests (maikzz32/lucebox-halo-cluster): protocol + # round-trips, placement invariants, decision hooks over a socketpair, + # control-channel handshake. GPU-free at run time. + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/unit/test_cluster_unit.cpp") + add_executable(test_cluster_unit test/unit/test_cluster_unit.cpp) + target_include_directories(test_cluster_unit PRIVATE ${LUCE_SRC_INCLUDE_DIRS}) + target_link_libraries(test_cluster_unit PRIVATE luce_common ggml ${LUCE_GGML_BACKEND_TARGET} ggml-base) + if(NOT WIN32) + target_link_libraries(test_cluster_unit PRIVATE pthread) + else() + target_link_libraries(test_cluster_unit PRIVATE ws2_32) + endif() + if(LUCE_GPU_BACKEND STREQUAL "hip") + target_compile_definitions(test_cluster_unit PRIVATE LUCE_BACKEND_HIP=1 GGML_USE_HIP) + target_link_libraries(test_cluster_unit PRIVATE hip::host) + else() + target_compile_definitions(test_cluster_unit PRIVATE LUCE_BACKEND_CUDA=1) + find_package(CUDAToolkit REQUIRED) + target_link_libraries(test_cluster_unit PRIVATE CUDA::cudart) + endif() + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/unit/test_ggml_cluster_allreduce.cpp") add_executable(test_ggml_cluster_allreduce test/unit/test_ggml_cluster_allreduce.cpp) target_include_directories(test_ggml_cluster_allreduce PRIVATE ${LUCE_SRC_INCLUDE_DIRS}) @@ -2586,6 +2679,7 @@ if(LUCE_TESTS) test_ds4_dmix_entry_validation test_ggml_rmsnorm_batch test_ggml_cluster_allreduce + test_cluster_unit test_qwen35_tensor_parallel test_deepseek4_hc_cuda test_capture_safe_copies diff --git a/server/deps/llama.cpp/ggml/include/ggml-backend.h b/server/deps/llama.cpp/ggml/include/ggml-backend.h index fdfee8076..b7c8cf731 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-backend.h +++ b/server/deps/llama.cpp/ggml/include/ggml-backend.h @@ -390,6 +390,13 @@ extern "C" { GGML_API void ggml_backend_sched_set_batch_split_copies( ggml_backend_sched_t sched, bool enabled); + // Split-input copies get storage of their own (never shared with another + // tensor of the graph), so filling one cannot overwrite data a split on + // its backend still reads: the scheduler then skips the blocking wait for + // the destination backend before each split's copies, and the host keeps + // queueing while the devices run. Must be set before allocation. + GGML_API void ggml_backend_sched_set_dedicated_copies( + ggml_backend_sched_t sched, bool enabled); // // Meta backend diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index e1b800f99..a8b862869 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -162,6 +162,13 @@ GGML_BACKEND_API int ggml_backend_cuda_set_mmvq_max_ncols_override(int max_ncols // Returns the previous setting. GGML_BACKEND_API bool ggml_backend_cuda_set_mmvq_batch_invariant(bool enabled); +// A MUL_MAT node whose src0 is a column range of wider Q8_0 rows (a view with +// the rows' stride, e.g. one rank's slice of a contraction): on the RDNA4 Q8_0 +// matvec every block keeps the lane it takes in the whole rows, so the product +// is bit-identical to the whole rows times src1 zero-padded outside the range, +// at the range's reads. Other kernels compute the range's product as usual. +GGML_BACKEND_API void ggml_backend_cuda_mul_mat_whole_row_lanes(struct ggml_tensor * mul_mat); + // Calling-thread DS4 mixed-expert dispatch ceiling, scoped to a graph compute. // Accepts 0 (the default of five) or 1..16; returns the previous ceiling. GGML_BACKEND_API int ggml_backend_cuda_set_ds4_mix_mmv_max_tokens_override(int max_tokens); diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 28015e65f..b85fe1599 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -1880,6 +1880,17 @@ extern "C" { struct ggml_tensor * sinks, float scale); + // ggml_soft_max_ext_sink_col with an additive mask (scale, then the mask, + // rounded as separate scale and add launches round them); a is + // [n_cols, rows, ...] and mask [n_cols, >= rows, ...] broadcast over + // dims 2 and 3. + GGML_API struct ggml_tensor * ggml_soft_max_ext_sink_col_mask( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + struct ggml_tensor * sinks, + float scale); + GGML_API void ggml_soft_max_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks); @@ -2568,6 +2579,21 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * positions); + // Run several single-token query lanes of one raw ring in one op, as if + // each lane wrote its own ring row and then attended before the next lane + // wrote: query token t reads lane_kv row j in place of ring row + // ring_rows[j] for every j <= t, and the ring's own contents elsewhere. + // ring_rows is I32 [n_query] (row indices into the raw segment), lane_kv + // has K's type and is contiguous [D, n_query]. Every other row, the mask + // and the selection keep their per-token meaning, so each token's result + // equals the single-token op over its own ring state bit for bit. Only the + // D=512 split-KV decode kernel consumes it (segmented K/V, indexed mask, + // n_query <= 8, no fused RoPE: ring_rows takes the RoPE positions slot). + GGML_API void ggml_flash_attn_ext_set_ds4_lane_rows( + struct ggml_tensor * a, + struct ggml_tensor * ring_rows, + struct ggml_tensor * lane_kv); + // True when flash_attn_ext carries the DS4 sparse-layout or fused-RoPE // contract. Backends must implement that complete contract or reject it. GGML_API bool ggml_flash_attn_ext_is_ds4( @@ -2685,6 +2711,7 @@ extern "C" { GGML_MOE_FUSED_HOST_POST = -8, GGML_MOE_FUSED_HOST_WAIT = -9, GGML_MOE_FUSED_PROTECTED_ROUTES = -10, + GGML_MOE_FUSED_COMBINE_MASKED = -11, }; // Word offsets in ggml_tensor::op_params for the host mailbox pair. All @@ -2746,11 +2773,34 @@ extern "C" { ggml_cluster_allreduce_fn fn, void * user); + // The same collective in place: the result is a view of `a`, so the sum + // overwrites `a` and no device copy runs ahead of the collective. Only + // for an `a` that no other node reads. + GGML_API struct ggml_tensor * ggml_cluster_allreduce_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + ggml_cluster_allreduce_fn fn, + void * user); + GGML_API struct ggml_tensor * ggml_laguna_moe_combine( struct ggml_context * ctx, struct ggml_tensor * experts, struct ggml_tensor * expert_weights); + // One owner's routed-expert sum: out[:, t] = sum_e experts[:, e, t] * + // (weights[e, t] * valid[0, ids[e, t], t]), the masked weight, product and + // sum rounded as mul(weights, get_rows(valid, ids)), mul and repeat_back + // round them (routes with a zero masked weight add nothing). + // experts [n_embd, n_used, T], weights [n_used, T] F32, valid [1, n_expert, T] + // F32, ids [n_used, T] I32 (global expert ids; a route masked to -1 adds + // nothing). + GGML_API struct ggml_tensor * ggml_moe_combine_masked( + struct ggml_context * ctx, + struct ggml_tensor * experts, + struct ggml_tensor * weights, + struct ggml_tensor * valid, + struct ggml_tensor * ids); + // Coarse DeepSeek-V4 routed-owner op. gate_up contains concatenated gate // and up output rows; the backend performs fused gate/up MMVQ + clamped // SwiGLU, down MMVQ, route weighting, and local reduction. @@ -2917,6 +2967,23 @@ extern "C" { // ggml_ds4_hc_collapse: dst[d, t] = sum_h hc[d + h*n_embd, t] * pre[h, t], // rounded as mul, permute and sum_rows round it. // hc: F32 [n_embd*n_hc, T] (nb0 = 4), pre: F32 [n_hc, T] (nb0 = 4) -> F32 [n_embd, T]. + // ggml_ds4_argmax_pair: the argmax of each column of logits [n_rows, T] + // (rows row_offset.. of a vocabulary split across n_slots ranks) as one + // slot of an F32 [8 * n_slots, T] exchange row: the value's 32 bits and the + // global index as eight byte-valued floats (exact through a bf16 sum), + // zeros in the other slots. ggml_ds4_argmax_pick reads the summed rows + // back: the largest value, ties to the lowest index (I32 [T]). + GGML_API struct ggml_tensor * ggml_ds4_argmax_pair( + struct ggml_context * ctx, + struct ggml_tensor * logits, + int row_offset, + int slot, + int n_slots); + GGML_API struct ggml_tensor * ggml_ds4_argmax_pick( + struct ggml_context * ctx, + struct ggml_tensor * pairs, + int n_slots); + GGML_API struct ggml_tensor * ggml_ds4_hc_collapse( struct ggml_context * ctx, struct ggml_tensor * hc, diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp index f4643547c..2c3d64e7a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp @@ -883,6 +883,7 @@ struct ggml_backend_sched { int deferred_peer_copies_capacity; bool split_deferred_peer_copies; bool batch_split_copies; + bool dedicated_copies; ggml_backend_buffer_t batch_staging_buffers[GGML_SCHED_MAX_BACKENDS]; uint8_t * batch_staging_bases[GGML_SCHED_MAX_BACKENDS]; size_t batch_staging_sizes[GGML_SCHED_MAX_BACKENDS]; @@ -1522,7 +1523,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra for (int c = 0; c < sched->n_copies; c++) { struct ggml_tensor * tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); - if (sched->n_copies > 1 || sched->copy_events) { + if (sched->n_copies > 1 || sched->copy_events || sched->dedicated_copies) { ggml_set_input(tensor_copy); ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor } @@ -2370,7 +2371,11 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_tensor_copy(input, input_cpy); } else { // wait for the split backend to finish using the input before overwriting it - copy_state.wait_for_destination_generation(); + // (a dedicated copy is not shared with anything the backend still reads, + // and the previous evaluation ended with every backend synchronized) + if (!sched->dedicated_copies) { + copy_state.wait_for_destination_generation(); + } // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used ggml_tensor * node = split->graph.nodes[0]; @@ -2970,6 +2975,13 @@ void ggml_backend_sched_set_batch_split_copies( sched->batch_split_copies = enabled; } +void ggml_backend_sched_set_dedicated_copies( + ggml_backend_sched_t sched, bool enabled) { + GGML_ASSERT(sched); + GGML_ASSERT(!sched->is_alloc); + sched->dedicated_copies = enabled; +} + int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); return sched->n_splits; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp index ca3c22619..093d2dfe9 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp @@ -5362,11 +5362,12 @@ static void ggml_compute_forward_soft_max_f32( // sinks const float * sk = src2 ? (float *)((char *) src2->data) : nullptr; - // Lucebox sink-as-virtual-column mode (see ggml_soft_max_ext_sink_col) + // Lucebox sink-as-virtual-column mode (see ggml_soft_max_ext_sink_col and + // its _mask form: the mask is added below like any other). int32_t sink_col = 0; memcpy(&sink_col, (int32_t *) dst->op_params + 2, sizeof(int32_t)); if (sink_col) { - GGML_ASSERT(sk && !src1); + GGML_ASSERT(sk); } for (int64_t i03 = 0; i03 < ne03; i03++) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh index a6b134044..083f262f0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh @@ -2,6 +2,9 @@ #include "ggml.h" #include "ggml-impl.h" + +// MUL_MAT op_params[1]: see ggml_backend_cuda_mul_mat_whole_row_lanes. +#define GGML_CUDA_WHOLE_ROW_LANES 0x574C524B #include "ggml-cuda.h" #include @@ -1492,6 +1495,13 @@ struct ggml_backend_cuda_context { }; std::vector luce_q8_memo; + // DS4 HC boundary launches (ds4-hc.cu) synchronize their blocks on this + // device barrier: an arrival count and a generation. + unsigned int * ds4_hc_barrier = nullptr; + // The split boundary's mix rows ([64][32] floats) and per-token block + // counts (64 uints), allocated with the barrier. + void * ds4_hc_split_scratch = nullptr; + cudaStream_t streams[GGML_CUDA_MAX_DEVICES][GGML_CUDA_MAX_STREAMS] = { { nullptr } }; cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES] = {nullptr}; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu index 3d7a698b9..e3586763b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu @@ -1,4 +1,5 @@ #include "cpy.cuh" +#include "quantize.cuh" #include "dequantize.cuh" #include "cpy-utils.cuh" #include "../../rocmfp4/rocmfp4_hip_scale.cuh" @@ -627,6 +628,14 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg const bool can_be_transposed = nb01 == (int64_t)ggml_element_size(src0) && src0->ne[3] == 1 && nb02 == ne00 * ne01 * (int64_t)ggml_element_size(src0); + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_Q8_1) { + // The q8_1 rows MMVQ reads, from the kernel it quantizes its own + // activations with: an MoE owner on another GPU takes its input this way. + GGML_ASSERT(contiguous_srcs && ne00 % MATRIX_ROW_PADDING == 0); + quantize_row_q8_1_cuda((const float *) src0_ddc, nullptr, src1_ddc, GGML_TYPE_F32, + ne00, ne00, ne00*ne01, ne00*ne01*ne02, ne00, ne01, ne02, src0->ne[3], main_stream); + return; + } if (src0->type == src1->type && contiguous_srcs) { GGML_ASSERT(ggml_nbytes(src0) == ggml_nbytes(src1)); #if defined(GGML_USE_MUSA) && defined(GGML_MUSA_MUDNN_COPY) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu index 0ee122822..d78048808 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu @@ -1,5 +1,7 @@ #include "ds4-hc.cuh" +#include + // Fused DeepSeek4 hyper-connection ops. // // mode 0 (pre): src0 = mix [mix_dim,n_tokens] (f32, from fn @ rms_norm(hc_state)) @@ -490,11 +492,46 @@ static __device__ void ds4_router_sort_desc(int * idx, float * val, int n, int c } } -template +// The weights of one token's selected routes: ds4_router_weights_kernel's +// arithmetic for token t (one thread). +static __device__ __forceinline__ void ds4_router_weights_token( + const float * __restrict__ logits, size_t ld, const int32_t * __restrict__ ids, + int k, float clamp_min, float scale, int apply_scale, float * __restrict__ dst, int t) { + // One thread per token: k <= 32 values, reduced in reduce_rows' order. + float w[32]; + float v[32]; +#pragma unroll + for (int l = 0; l < 32; ++l) { + w[l] = l < k ? ds4_router_prob(logits[(size_t) t * ld + ids[(size_t) t * k + l]]) : 0.0f; + v[l] = w[l]; + } + // reduce_rows_f32 on a k-wide row: lane l holds element l, then the + // 32-lane xor butterfly (the second block stage adds zeros only). +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + float n[32]; +#pragma unroll + for (int l = 0; l < 32; ++l) n[l] = v[l] + v[l ^ off]; +#pragma unroll + for (int l = 0; l < 32; ++l) v[l] = n[l]; + } + float sum = v[0]; + sum = fminf(fmaxf(sum, clamp_min), INFINITY); + for (int i = 0; i < k; ++i) { + float r = w[i] / sum; + if (apply_scale) r = __fadd_rn(__fmul_rn(scale, r), 0.0f); // scale_f32: scale * x + 0 + dst[(size_t) t * k + i] = r; + } +} + +// WITH_WEIGHTS: thread 0 then writes the routes' weights as +// ds4_router_weights_kernel would (the router's two launches in one). +template static __global__ void ds4_router_select_kernel( const float * __restrict__ logits, size_t ld, const float * __restrict__ bias, const float * __restrict__ native_bias, const int32_t * __restrict__ protected_mask, - int n_expert, int k, int32_t * __restrict__ dst) { + int n_expert, int k, int32_t * __restrict__ dst, + float clamp_min = 0.0f, float scale = 1.0f, float * __restrict__ weights_dst = nullptr) { __shared__ int idx[NPAD]; __shared__ float val[NPAD]; __shared__ float probs[NPAD]; @@ -522,36 +559,43 @@ static __global__ void ds4_router_select_kernel( __syncthreads(); ds4_router_sort_desc(idx, val, n_expert, col); if (col < k) dst[(size_t) t * k + col] = keep_native ? native_top[col] : idx[col]; + if constexpr (WITH_WEIGHTS) { + __syncthreads(); + __threadfence_block(); + if (col == 0) { + ds4_router_weights_token(logits, ld, dst, k, clamp_min, scale, scale != 1.0f ? 1 : 0, + weights_dst, t); + } + } } static __global__ void ds4_router_weights_kernel( const float * __restrict__ logits, size_t ld, const int32_t * __restrict__ ids, int k, float clamp_min, float scale, int apply_scale, float * __restrict__ dst) { - // One thread per token: k <= 32 values, reduced in reduce_rows' order. - const int t = blockIdx.x; - float w[32]; - float v[32]; -#pragma unroll - for (int l = 0; l < 32; ++l) { - w[l] = l < k ? ds4_router_prob(logits[(size_t) t * ld + ids[(size_t) t * k + l]]) : 0.0f; - v[l] = w[l]; - } - // reduce_rows_f32 on a k-wide row: lane l holds element l, then the - // 32-lane xor butterfly (the second block stage adds zeros only). -#pragma unroll - for (int off = 16; off > 0; off >>= 1) { - float n[32]; -#pragma unroll - for (int l = 0; l < 32; ++l) n[l] = v[l] + v[l ^ off]; -#pragma unroll - for (int l = 0; l < 32; ++l) v[l] = n[l]; - } - float sum = v[0]; - sum = fminf(fmaxf(sum, clamp_min), INFINITY); - for (int i = 0; i < k; ++i) { - float r = w[i] / sum; - if (apply_scale) r = __fadd_rn(__fmul_rn(scale, r), 0.0f); // scale_f32: scale * x + 0 - dst[(size_t) t * k + i] = r; + ds4_router_weights_token(logits, ld, ids, k, clamp_min, scale, apply_scale, dst, (int) blockIdx.x); +} + +// The router's select and weights ops (modes 4 and 5) over the same logits, +// in one launch. +void ggml_cuda_op_ds4_router_fused(ggml_backend_cuda_context & ctx, ggml_tensor * select, ggml_tensor * weights) { + const ggml_tensor * logits = select->src[0]; + const int k = ggml_get_op_params_i32(select, 1); + const int n_expert = (int) logits->ne[0]; + const int n_tokens = (int) logits->ne[1]; + const float * nb = select->src[2] ? (const float *) select->src[2]->data : nullptr; + const int32_t * pm = select->src[3] ? (const int32_t *) select->src[3]->data : nullptr; + const float clamp_min = ggml_get_op_params_f32(weights, 4); + const float scale = ggml_get_op_params_f32(weights, 5); + if (n_expert <= 512) { + ds4_router_select_kernel<512, true><<>>( + (const float *) logits->data, logits->nb[1] / sizeof(float), + (const float *) select->src[1]->data, nb, pm, n_expert, k, (int32_t *) select->data, + clamp_min, scale, (float *) weights->data); + } else { + ds4_router_select_kernel<1024, true><<>>( + (const float *) logits->data, logits->nb[1] / sizeof(float), + (const float *) select->src[1]->data, nb, pm, n_expert, k, (int32_t *) select->data, + clamp_min, scale, (float *) weights->data); } } @@ -593,9 +637,662 @@ static __global__ void ds4_hc_collapse_kernel( } } +// ── One HC sub-block boundary per launch ──────────────────────────────── +// The graph builds the end of a staggered (V4.1) sub-block X and the start of +// the next one as six launches: +// n1 = rms_norm(R) R: the sub-block's HC input streams +// mix = fn @ n1 F16 weights, MMVF +// pre = ds4_hc_pre(mix, base, R) Sinkhorn split (+ working vector) +// R' = ds4_hc_post(R, block_out, split) the next HC streams +// in = ds4_hc_collapse(R', pre[0:n_hc]) the next sub-block's input +// out = rms_norm(in) * w ... normalized for its projections +// ds4_hc_boundary_kernel does all six for one token per block. Each phase +// repeats the arithmetic of the kernel it replaces in the same order: +// rms_norm_f32<1024> (thread col mapping and block reduction), MMVF +// (per-thread half2 partial sums, warp then +// cross-warp reduction), ds4_hc_pre_kernel_t, ds4_hc_post_kernel, +// ds4_hc_collapse_kernel (contraction off) and rms_norm_f32<1024, mul>, so +// the outputs are bit-identical (test_ds4_fused_ops_cuda). + +// The collapse of one embedding column: ds4_hc_collapse_kernel's arithmetic. +template +static __device__ __forceinline__ float ds4_hc_collapse_value(const float * v, const float * p) { +#pragma clang fp contract(off) + if constexpr (NHC == 4) { + const float v0 = __fmul_rn(v[0], p[0]); + const float v1 = __fmul_rn(v[1], p[1]); + const float v2 = __fmul_rn(v[2], p[2]); + const float v3 = __fmul_rn(v[3], p[3]); + return (v0 + v2) + (v1 + v3); + } else { + float w[32]; +#pragma unroll + for (int l = 0; l < 32; ++l) w[l] = l < NHC ? __fmul_rn(v[l], p[l]) : 0.0f; +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { + float n[32]; +#pragma unroll + for (int l = 0; l < 32; ++l) n[l] = w[l] + w[l ^ off]; +#pragma unroll + for (int l = 0; l < 32; ++l) w[l] = n[l]; + } + return w[0]; + } +} + +#define DS4_HC_BOUNDARY_MAX_COLS 8 // n_embd <= 8 * 1024 + +// An F32 value the compiler must materialize as it stands: the replaced +// kernels store it to memory before the next one reads it, so it must not be +// fused into a following conversion (v_fma_mix) or FMA. +static __device__ __forceinline__ float ds4_hc_f32(float x) { +#if defined(GGML_USE_HIP) + asm volatile("" : "+v"(x)); +#else + asm volatile("" : "+f"(x)); +#endif + return x; +} + +// Every block of the grid waits here for the others (one block per token +// reaches it). The barrier keeps its state on the device only: the last block +// to arrive zeroes the count and advances the generation the others wait on, +// so it needs no per-launch target from the host and a CUDA graph can replay +// the launch. All the grid's blocks must be resident at once. +static __device__ __forceinline__ void ds4_hc_grid_barrier(unsigned int * barrier, const unsigned int n_blocks) { + __syncthreads(); + if (threadIdx.x == 0) { + unsigned int * count = barrier; + unsigned int * generation = barrier + 1; + const unsigned int gen = atomicAdd(generation, 0u); + __threadfence(); + if (atomicAdd(count, 1u) == n_blocks - 1) { + atomicExch(count, 0u); + __threadfence(); + atomicAdd(generation, 1u); + } else { + while (atomicAdd(generation, 0u) == gen) { +#if defined(GGML_USE_HIP) + __builtin_amdgcn_s_sleep(1); +#endif + } + } + __threadfence(); + } + __syncthreads(); +} + +template +static __global__ void __launch_bounds__(1024) ds4_hc_boundary_kernel( + const float * __restrict__ R, int64_t R_stride, + const half * __restrict__ fn, int64_t fn_stride, + const float * __restrict__ base, + const float * __restrict__ block_out, int64_t bo_stride, + const float * __restrict__ w, + float * __restrict__ pre_dst, int64_t pre_stride, + float * __restrict__ R_next, int64_t Rn_stride, + float * out, int64_t out_stride, + const int n_embd, const float eps_hc, const float eps_out, const int iters, + const float pre_scale, const float post_scale, const float comb_scale, + const int write_pre, + unsigned int * barrier, + block_q8_1 * __restrict__ q8, const int64_t q8_row_blocks) { + constexpr int block_size = 1024; + constexpr int mix_dim = 2 * NHC + NHC * NHC; + constexpr int n_groups = block_size / MMV_BLOCK; + constexpr int warps_per_group = MMV_BLOCK / WARP_SIZE; + const int token = (int) blockIdx.x; + const int tid = (int) threadIdx.x; + const int lane = tid % WARP_SIZE; + const int ncols_hc = n_embd * NHC; + + R += (int64_t) token * R_stride; + block_out += (int64_t) token * bo_stride; + R_next += (int64_t) token * Rn_stride; + out += (int64_t) token * out_stride; + if (write_pre) pre_dst += (int64_t) token * pre_stride; + + extern __shared__ half2 s_y2[]; // normalized R, as MMVF reads it + __shared__ float s_sum[32]; + __shared__ float s_sum_out[32]; + __shared__ float s_part[mix_dim][warps_per_group]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + __shared__ float split[DS4_HC_MAX_MIX]; + + // n1 = rms_norm(R): rms_norm_f32<1024>. + float tmp = 0.0f; + for (int col = tid; col < ncols_hc; col += block_size) { + const float xi = R[col]; + tmp += xi * xi; + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / ncols_hc; + const float scale = rsqrtf(mean + eps_hc); + // MMVF reads F32 pairs and rounds each to half (make_half2 of two floats). + const float2 * R2 = (const float2 *) R; + for (int c2 = tid; c2 < ncols_hc / 2; c2 += block_size) { + const float2 x = R2[c2]; + // rms_norm_f32 stores scale * x as F32; MMVF then rounds it to half. + const float y0 = ds4_hc_f32(scale * x.x); + const float y1 = ds4_hc_f32(scale * x.y); + s_y2[c2] = make_half2(y0, y1); + } + __syncthreads(); + + // mix = fn @ n1: each group of MMV_BLOCK threads is one MMVF block. + { + const int g = tid / MMV_BLOCK; + const int lt = tid % MMV_BLOCK; + const int ncols2 = ncols_hc / 2; + for (int r = g; r < mix_dim; r += n_groups) { + const half2 * x2 = (const half2 *) (fn + (int64_t) r * fn_stride); + half2 sumh2 = make_half2(0.0f, 0.0f); +#pragma unroll 8 + for (int col2 = lt; col2 < ncols2; col2 += MMV_BLOCK) { + const half2 tmpx = x2[col2]; + sumh2 += tmpx * s_y2[col2]; + } + float sumf = __low2float(sumh2) + __high2float(sumh2); + sumf = warp_reduce_sum(sumf); + if (lane == 0) { + s_part[r][lt / WARP_SIZE] = sumf; + } + } + } + __syncthreads(); + if (tid / WARP_SIZE < mix_dim) { + // MMVF's cross-warp stage: the warp sums in lanes 0..warps-1, zeros above. + const int r = tid / WARP_SIZE; + float v = lane < warps_per_group ? s_part[r][lane] : 0.0f; + v = warp_reduce_sum(v); + if (lane == 0) { + s_mix[r] = v; + } + } + if (tid < mix_dim) { + s_base[tid] = base[tid]; + } + __syncthreads(); + + // pre = ds4_hc_pre(mix, base, R): ds4_hc_pre_kernel_t. + if (tid == 0) { + ds4_hc_sinkhorn_split_t(s_mix, s_base, pre_scale, post_scale, comb_scale, iters, split); + if (write_pre) { +#pragma unroll + for (int i = 0; i < mix_dim; ++i) { + pre_dst[n_embd + i] = split[i]; + } + } + } + __syncthreads(); + if (write_pre) { + for (int d = tid; d < n_embd; d += block_size) { + float acc = 0.0f; +#pragma unroll + for (int h = 0; h < NHC; ++h) { + acc += split[h] * R[(int64_t) h * n_embd + d]; + } + pre_dst[d] = acc; + } + } + + // R' = ds4_hc_post(R, block_out, split); in = collapse(R', pre). The + // column mapping is rms_norm_f32<1024>'s over n_embd, and in stays in + // registers: out may share memory with R or block_out (the unfused + // allocation frees both before out), so it is written only after every + // block has finished reading them. + const float * post = split + NHC; + const float * comb = split + 2 * NHC; + float xin[DS4_HC_BOUNDARY_MAX_COLS]; + float tmp_out = 0.0f; +#pragma unroll + for (int k = 0; k < DS4_HC_BOUNDARY_MAX_COLS; ++k) { + const int d = tid + k * block_size; + xin[k] = 0.0f; + if (d < n_embd) { + float rn[NHC]; +#pragma unroll + for (int h = 0; h < NHC; ++h) { + // ds4_hc_post_kernel compiles to one multiply, then an FMA + // per source copy accumulating into it: spelled out here so + // the unrolled loop cannot contract the other way. + float acc = __fmul_rn(block_out[d], post[h]); + for (int src = 0; src < NHC; ++src) { + acc = __fmaf_rn(comb[h + src * NHC], R[(int64_t) src * n_embd + d], acc); + } + rn[h] = acc; + R_next[(int64_t) h * n_embd + d] = acc; + } + const float xi = ds4_hc_collapse_value(rn, split); + xin[k] = xi; + tmp_out += xi * xi; + } + } + // out = rms_norm(in) * w: rms_norm_f32<1024, mul>. + tmp_out = block_reduce(tmp_out, s_sum_out); + const float mean_out = tmp_out / n_embd; + const float scale_out = rsqrtf(mean_out + eps_out); + + ds4_hc_grid_barrier(barrier, gridDim.x); +#pragma unroll + for (int k = 0; k < DS4_HC_BOUNDARY_MAX_COLS; ++k) { + const int d = tid + k * block_size; + if (d < n_embd) { + const float xo = scale_out * xin[k] * w[d]; + out[d] = xo; + if (q8) { + // quantize_q8_1 of out: a warp holds one 32-wide block (d is + // warp-contiguous), reduced in the same butterfly order. + float amax = fabsf(xo); + float sum = xo; + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + const float dq = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(xo / dq); + block_q8_1 * yb = q8 + (int64_t) token * q8_row_blocks + d / QK8_1; + const int iqs = d % QK8_1; + yb->qs[iqs] = q; + if (iqs == 0) { + yb->ds = make_half2(dq, sum); + } + } + } + } +} + +template +static __global__ void __launch_bounds__(1024) ds4_hc_boundary_split_kernel( + const float * __restrict__ R, int64_t R_stride, + const half * __restrict__ fn, int64_t fn_stride, + const float * __restrict__ base, + const float * __restrict__ block_out, int64_t bo_stride, + const float * __restrict__ w, + float * __restrict__ pre_dst, int64_t pre_stride, + float * __restrict__ R_next, int64_t Rn_stride, + float * out, int64_t out_stride, + const int n_embd, const float eps_hc, const float eps_out, const int iters, + const float pre_scale, const float post_scale, const float comb_scale, + const int write_pre, + unsigned int * barrier, + block_q8_1 * __restrict__ q8, const int64_t q8_row_blocks, + float * __restrict__ mix_scratch, unsigned int * __restrict__ rows_done) { + constexpr int block_size = 1024; + constexpr int mix_dim = 2 * NHC + NHC * NHC; + constexpr int n_groups = block_size / MMV_BLOCK; + constexpr int warps_per_group = MMV_BLOCK / WARP_SIZE; + const int token = (int) blockIdx.x; + const int tid = (int) threadIdx.x; + const int lane = tid % WARP_SIZE; + const int ncols_hc = n_embd * NHC; + + R += (int64_t) token * R_stride; + block_out += (int64_t) token * bo_stride; + R_next += (int64_t) token * Rn_stride; + out += (int64_t) token * out_stride; + if (write_pre) pre_dst += (int64_t) token * pre_stride; + + extern __shared__ half2 s_y2[]; // normalized R, as MMVF reads it + __shared__ float s_sum[32]; + __shared__ float s_sum_out[32]; + __shared__ float s_part[n_groups][warps_per_group]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + __shared__ float split[DS4_HC_MAX_MIX]; + + // n1 = rms_norm(R): rms_norm_f32<1024>. + float tmp = 0.0f; + for (int col = tid; col < ncols_hc; col += block_size) { + const float xi = R[col]; + tmp += xi * xi; + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / ncols_hc; + const float scale = rsqrtf(mean + eps_hc); + // MMVF reads F32 pairs and rounds each to half (make_half2 of two floats). + const float2 * R2 = (const float2 *) R; + for (int c2 = tid; c2 < ncols_hc / 2; c2 += block_size) { + const float2 x = R2[c2]; + // rms_norm_f32 stores scale * x as F32; MMVF then rounds it to half. + const float y0 = ds4_hc_f32(scale * x.x); + const float y1 = ds4_hc_f32(scale * x.y); + s_y2[c2] = make_half2(y0, y1); + } + __syncthreads(); + + // mix = fn @ n1: block (token, split) computes rows split*n_groups + g, + // one per MMVF group, with MMVF's per-row arithmetic. + static_assert(NSPLIT * n_groups >= mix_dim, "the split blocks must cover every mix row"); + const int split_id = (int) blockIdx.y; + { + const int g = tid / MMV_BLOCK; + const int lt = tid % MMV_BLOCK; + const int ncols2 = ncols_hc / 2; + const int r = split_id * n_groups + g; + if (r < mix_dim) { + const half2 * x2 = (const half2 *) (fn + (int64_t) r * fn_stride); + half2 sumh2 = make_half2(0.0f, 0.0f); +#pragma unroll 8 + for (int col2 = lt; col2 < ncols2; col2 += MMV_BLOCK) { + const half2 tmpx = x2[col2]; + sumh2 += tmpx * s_y2[col2]; + } + float sumf = __low2float(sumh2) + __high2float(sumh2); + sumf = warp_reduce_sum(sumf); + if (lane == 0) { + s_part[g][lt / WARP_SIZE] = sumf; + } + } + } + __syncthreads(); + if (tid / WARP_SIZE < n_groups) { + // MMVF's cross-warp stage for this block's rows. + const int g = tid / WARP_SIZE; + const int r = split_id * n_groups + g; + if (r < mix_dim) { + float v = lane < warps_per_group ? s_part[g][lane] : 0.0f; + v = warp_reduce_sum(v); + if (lane == 0) { + mix_scratch[(int64_t) token * 32 + r] = v; + __threadfence(); // the row is visible before this block counts itself done + } + } + } + // The token's last block to finish its rows runs the rest of the boundary. + // Every launch adds NSPLIT to each of its tokens' counts. + __shared__ int s_last; + __syncthreads(); + if (tid == 0) { + __threadfence(); + const unsigned int done = atomicAdd(&rows_done[token], 1u) + 1u; + s_last = (done % NSPLIT) == 0; + __threadfence(); + } + __syncthreads(); + if (!s_last) { + return; + } + if (tid < mix_dim) { + s_mix[tid] = ((volatile const float *) mix_scratch)[(int64_t) token * 32 + tid]; + s_base[tid] = base[tid]; + } + __syncthreads(); + + // pre = ds4_hc_pre(mix, base, R): ds4_hc_pre_kernel_t. + if (tid == 0) { + ds4_hc_sinkhorn_split_t(s_mix, s_base, pre_scale, post_scale, comb_scale, iters, split); + if (write_pre) { +#pragma unroll + for (int i = 0; i < mix_dim; ++i) { + pre_dst[n_embd + i] = split[i]; + } + } + } + __syncthreads(); + if (write_pre) { + for (int d = tid; d < n_embd; d += block_size) { + float acc = 0.0f; +#pragma unroll + for (int h = 0; h < NHC; ++h) { + acc += split[h] * R[(int64_t) h * n_embd + d]; + } + pre_dst[d] = acc; + } + } + + // R' = ds4_hc_post(R, block_out, split); in = collapse(R', pre). The + // column mapping is rms_norm_f32<1024>'s over n_embd, and in stays in + // registers: out may share memory with R or block_out (the unfused + // allocation frees both before out), so it is written only after every + // block has finished reading them. + const float * post = split + NHC; + const float * comb = split + 2 * NHC; + float xin[DS4_HC_BOUNDARY_MAX_COLS]; + float tmp_out = 0.0f; +#pragma unroll + for (int k = 0; k < DS4_HC_BOUNDARY_MAX_COLS; ++k) { + const int d = tid + k * block_size; + xin[k] = 0.0f; + if (d < n_embd) { + float rn[NHC]; +#pragma unroll + for (int h = 0; h < NHC; ++h) { + // ds4_hc_post_kernel compiles to one multiply, then an FMA + // per source copy accumulating into it: spelled out here so + // the unrolled loop cannot contract the other way. + float acc = __fmul_rn(block_out[d], post[h]); + for (int src = 0; src < NHC; ++src) { + acc = __fmaf_rn(comb[h + src * NHC], R[(int64_t) src * n_embd + d], acc); + } + rn[h] = acc; + R_next[(int64_t) h * n_embd + d] = acc; + } + const float xi = ds4_hc_collapse_value(rn, split); + xin[k] = xi; + tmp_out += xi * xi; + } + } + // out = rms_norm(in) * w: rms_norm_f32<1024, mul>. + tmp_out = block_reduce(tmp_out, s_sum_out); + const float mean_out = tmp_out / n_embd; + const float scale_out = rsqrtf(mean_out + eps_out); + + ds4_hc_grid_barrier(barrier, gridDim.x); +#pragma unroll + for (int k = 0; k < DS4_HC_BOUNDARY_MAX_COLS; ++k) { + const int d = tid + k * block_size; + if (d < n_embd) { + const float xo = scale_out * xin[k] * w[d]; + out[d] = xo; + if (q8) { + // quantize_q8_1 of out: a warp holds one 32-wide block (d is + // warp-contiguous), reduced in the same butterfly order. + float amax = fabsf(xo); + float sum = xo; + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + const float dq = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(xo / dq); + block_q8_1 * yb = q8 + (int64_t) token * q8_row_blocks + d / QK8_1; + const int iqs = d % QK8_1; + yb->qs[iqs] = q; + if (iqs == 0) { + yb->ds = make_half2(dq, sum); + } + } + } + } +} + +// Blocks of the split boundary kernel for `mmv_block` that `device` holds at +// once. Its finishing blocks wait for each other at the output barrier, so a +// grid wider than this would never finish. Cached per device and block size. +static int ds4_hc_split_resident_blocks(int device, int mmv_block, size_t shmem) { + struct Entry { size_t shmem; int blocks; }; + static Entry cache[GGML_CUDA_MAX_DEVICES][3] = {}; + const int b = mmv_block == 64 ? 0 : mmv_block == 128 ? 1 : 2; + Entry & e = cache[device][b]; + if (e.shmem != shmem) { + int per_cu = 0; + cudaError_t err; + switch (b) { + case 0: err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &per_cu, ds4_hc_boundary_split_kernel<4, 64, 2>, 1024, shmem); break; + case 1: err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &per_cu, ds4_hc_boundary_split_kernel<4, 128, 3>, 1024, shmem); break; + default: err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &per_cu, ds4_hc_boundary_split_kernel<4, 256, 6>, 1024, shmem); break; + } + if (err != cudaSuccess) { + (void) cudaGetLastError(); + per_cu = 0; + } + e = { shmem, per_cu * ggml_cuda_info().devices[device].nsm }; + } + return e.blocks; +} + +bool ggml_cuda_ds4_hc_boundary_supported(int n_hc, int n_embd, int mmv_block) { + // n_embd % 1024: every warp's columns are one whole q8_1 block of out. + return n_hc == 4 && n_embd >= 1024 && n_embd % 1024 == 0 && + n_embd <= DS4_HC_BOUNDARY_MAX_COLS * 1024 && + (n_embd * n_hc) % 2 == 0 && + (size_t) n_embd * n_hc / 2 * sizeof(half2) <= 48u * 1024u && + (mmv_block == 64 || mmv_block == 128 || mmv_block == 256); +} + +void ggml_cuda_ds4_hc_boundary( + ggml_backend_cuda_context & ctx, int n_tokens, int n_hc, int mmv_block, + const float * R, int64_t R_stride, const half * fn, int64_t fn_stride, const float * base, + const float * block_out, int64_t bo_stride, const float * w, + float * pre_dst, int64_t pre_stride, float * R_next, int64_t Rn_stride, + float * out, int64_t out_stride, + int n_embd, float eps_hc, float eps_out, int iters, + float pre_scale, float post_scale, float comb_scale, bool write_pre, + void * q8_out) { + GGML_ASSERT(ggml_cuda_ds4_hc_boundary_supported(n_hc, n_embd, mmv_block)); + GGML_ASSERT(n_tokens >= 1 && n_tokens <= 64); + // The output barrier's count and generation (ds4_hc_grid_barrier). + if (!ctx.ds4_hc_barrier) { + ggml_cuda_set_device(ctx.device); + CUDA_CHECK(cudaMalloc((void **) &ctx.ds4_hc_barrier, 2 * sizeof(unsigned int))); + CUDA_CHECK(cudaMemsetAsync(ctx.ds4_hc_barrier, 0, 2 * sizeof(unsigned int), ctx.stream())); + } + const size_t shmem = (size_t) n_embd * n_hc / 2 * sizeof(half2); + cudaStream_t stream = ctx.stream(); + // LUCE_DS4_HC_BOUNDARY_SPLIT=1: the mix over several blocks per token + // (one token's block reads the whole fn matrix alone otherwise). Only + // while every block fits at once beside the others, since the finishing + // blocks wait for each other at the output barrier. The unsplit grid is + // one block per token, at most MMVF_MAX_BATCH_SIZE of them. + static const bool split_mix = [] { + const char * v = getenv("LUCE_DS4_HC_BOUNDARY_SPLIT"); + return v && *v && strcmp(v, "0") != 0; + }(); + const int n_split = mmv_block == 64 ? 2 : mmv_block == 128 ? 3 : 6; + if (split_mix && n_tokens * n_split <= 48 && + n_tokens * n_split <= ds4_hc_split_resident_blocks(ctx.device, mmv_block, shmem)) { + if (!ctx.ds4_hc_split_scratch) { + ggml_cuda_set_device(ctx.device); + CUDA_CHECK(cudaMalloc(&ctx.ds4_hc_split_scratch, 64 * 32 * sizeof(float) + 64 * sizeof(unsigned int))); + CUDA_CHECK(cudaMemsetAsync(ctx.ds4_hc_split_scratch, 0, 64 * 32 * sizeof(float) + 64 * sizeof(unsigned int), stream)); + } + float * mix_scratch = (float *) ctx.ds4_hc_split_scratch; + unsigned int * rows_done = (unsigned int *) (mix_scratch + 64 * 32); +#define DS4_HC_BOUNDARY_SPLIT_LAUNCH(B, NS) \ + ds4_hc_boundary_split_kernel<4, B, NS><<>>( \ + R, R_stride, fn, fn_stride, base, block_out, bo_stride, w, \ + pre_dst, pre_stride, R_next, Rn_stride, out, out_stride, \ + n_embd, eps_hc, eps_out, iters, pre_scale, post_scale, comb_scale, write_pre ? 1 : 0, \ + ctx.ds4_hc_barrier, \ + (block_q8_1 *) q8_out, (int64_t) GGML_PAD(n_embd, MATRIX_ROW_PADDING) / QK8_1, \ + mix_scratch, rows_done) + switch (mmv_block) { + case 64: DS4_HC_BOUNDARY_SPLIT_LAUNCH(64, 2); break; + case 128: DS4_HC_BOUNDARY_SPLIT_LAUNCH(128, 3); break; + default: DS4_HC_BOUNDARY_SPLIT_LAUNCH(256, 6); break; + } +#undef DS4_HC_BOUNDARY_SPLIT_LAUNCH + CUDA_CHECK(cudaGetLastError()); + return; + } +#define DS4_HC_BOUNDARY_LAUNCH(B) \ + ds4_hc_boundary_kernel<4, B><<>>( \ + R, R_stride, fn, fn_stride, base, block_out, bo_stride, w, \ + pre_dst, pre_stride, R_next, Rn_stride, out, out_stride, \ + n_embd, eps_hc, eps_out, iters, pre_scale, post_scale, comb_scale, write_pre ? 1 : 0, \ + ctx.ds4_hc_barrier, \ + (block_q8_1 *) q8_out, (int64_t) GGML_PAD(n_embd, MATRIX_ROW_PADDING) / QK8_1) + switch (mmv_block) { + case 64: DS4_HC_BOUNDARY_LAUNCH(64); break; + case 128: DS4_HC_BOUNDARY_LAUNCH(128); break; + default: DS4_HC_BOUNDARY_LAUNCH(256); break; + } +#undef DS4_HC_BOUNDARY_LAUNCH + CUDA_CHECK(cudaGetLastError()); +} + // GGML_OP_DS4_HC modes (op_params[0]): 0 hc_pre, 1 hc_post, 2 hc_out, // 3 hc_post_split, 4 router_select, 5 router_weights, 6 hc_collapse. +// Mode 7: the column argmax of a vocabulary slice (argmax_f32's scan: first +// occurrence per thread, then the butterfly) as one rank's exchange slot. +static __global__ void ds4_argmax_pair_kernel(const float * __restrict__ x, float * __restrict__ dst, + const int64_t nrows, const int row_offset, + const int slot, const int n_slots) { + const int64_t t = blockIdx.x; + const float * col = x + t * nrows; + float maxval = -FLT_MAX; + int argmax = -1; + for (int64_t r = threadIdx.x; r < nrows; r += blockDim.x) { + const float v = col[r]; + if (v > maxval) { maxval = v; argmax = (int) r; } + } +#pragma unroll + for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) { + const float v = __shfl_xor_sync(0xFFFFFFFF, maxval, offset, WARP_SIZE); + const int c = __shfl_xor_sync(0xFFFFFFFF, argmax, offset, WARP_SIZE); + if (v > maxval || (v == maxval && c >= 0 && (argmax < 0 || c < argmax))) { maxval = v; argmax = c; } + } + __shared__ float s_val[32]; + __shared__ int s_arg[32]; + const int lane = threadIdx.x % WARP_SIZE, warp = threadIdx.x / WARP_SIZE; + if (lane == 0) { s_val[warp] = maxval; s_arg[warp] = argmax; } + __syncthreads(); + float * out = dst + t * (int64_t) (8 * n_slots); + if (threadIdx.x < 8 * n_slots) out[threadIdx.x] = 0.0f; + __syncthreads(); + if (threadIdx.x == 0) { + const int n_warps = blockDim.x / WARP_SIZE; + float bv = s_val[0]; int bi = s_arg[0]; + for (int w = 1; w < n_warps; ++w) { + if (s_val[w] > bv || (s_val[w] == bv && s_arg[w] >= 0 && (bi < 0 || s_arg[w] < bi))) { bv = s_val[w]; bi = s_arg[w]; } + } + const unsigned int vb = __float_as_uint(bv); + const unsigned int ib = (unsigned int) (bi < 0 ? 0 : bi + row_offset); + float * o = out + 8 * slot; + for (int b = 0; b < 4; ++b) { + o[b] = (float) ((vb >> (8 * b)) & 0xFFu); + o[4 + b] = (float) ((ib >> (8 * b)) & 0xFFu); + } + } +} + +// Mode 8: the winner of the summed slots, ties to the lower index. +static __global__ void ds4_argmax_pick_kernel(const float * __restrict__ pairs, int32_t * __restrict__ dst, + const int n_tokens, const int n_slots) { + const int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= n_tokens) return; + const float * p = pairs + (int64_t) t * (8 * n_slots); + float bv = 0.0f; int bi = -1; + for (int s = 0; s < n_slots; ++s) { + unsigned int vb = 0, ib = 0; + for (int b = 0; b < 4; ++b) { + vb |= ((unsigned int) p[8 * s + b]) << (8 * b); + ib |= ((unsigned int) p[8 * s + 4 + b]) << (8 * b); + } + const float v = __uint_as_float(vb); + if (bi < 0 || v > bv || (v == bv && (int) ib < bi)) { bv = v; bi = (int) ib; } + } + dst[t] = bi; +} + void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + if (ggml_get_op_params_i32(dst, 0) == 7) { + const ggml_tensor * logits = dst->src[0]; + ds4_argmax_pair_kernel<<<(unsigned) logits->ne[1], 1024, 0, ctx.stream()>>>( + (const float *) logits->data, (float *) dst->data, logits->ne[0], + ggml_get_op_params_i32(dst, 1), ggml_get_op_params_i32(dst, 2), ggml_get_op_params_i32(dst, 3)); + CUDA_CHECK(cudaGetLastError()); + return; + } + if (ggml_get_op_params_i32(dst, 0) == 8) { + const ggml_tensor * pairs = dst->src[0]; + const int n_tokens = (int) pairs->ne[1]; + ds4_argmax_pick_kernel<<<(n_tokens + 63) / 64, 64, 0, ctx.stream()>>>( + (const float *) pairs->data, (int32_t *) dst->data, n_tokens, ggml_get_op_params_i32(dst, 3)); + CUDA_CHECK(cudaGetLastError()); + return; + } const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; const ggml_tensor * src2 = dst->src[2]; @@ -735,3 +1432,77 @@ void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ABORT("ds4_hc: unknown mode"); } } + +// ── RMS norm + HC mix product for prefill bands ───────────────────────── +// Each block handles kRnmTokens tokens: the mix rows are read once per block +// for all of them, and each thread keeps the partial sums of its K slice. +static constexpr int kRnmThreads = 256; +static constexpr int kRnmTokens = 4; +static constexpr int kRnmMaxM = 32; + +template +static __global__ void __launch_bounds__(kRnmThreads) ds4_rms_norm_mix_kernel( + const float * __restrict__ x, const int64_t x_stride, const half * __restrict__ w, const int64_t w_stride, + float * __restrict__ out, const int64_t out_stride, const int k, const int n_tokens, const float eps) { + const int t0 = blockIdx.x * kRnmTokens; + const int nt = min(kRnmTokens, n_tokens - t0); + float acc[kRnmTokens][M]; + float ss[kRnmTokens]; +#pragma unroll + for (int j = 0; j < kRnmTokens; ++j) { + ss[j] = 0.0f; +#pragma unroll + for (int r = 0; r < M; ++r) acc[j][r] = 0.0f; + } + for (int i = threadIdx.x; i < k; i += kRnmThreads) { + float wv[M]; +#pragma unroll + for (int r = 0; r < M; ++r) wv[r] = __half2float(w[(int64_t) r * w_stride + i]); +#pragma unroll + for (int j = 0; j < kRnmTokens; ++j) { + if (j >= nt) break; + const float v = x[(int64_t) (t0 + j) * x_stride + i]; + ss[j] += v * v; +#pragma unroll + for (int r = 0; r < M; ++r) acc[j][r] += wv[r] * v; + } + } + __shared__ float s_part[kRnmThreads / WARP_SIZE][kRnmTokens][M + 1]; + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; +#pragma unroll + for (int j = 0; j < kRnmTokens; ++j) { + const float s = warp_reduce_sum(ss[j]); + if (lane == 0) s_part[warp][j][M] = s; +#pragma unroll + for (int r = 0; r < M; ++r) { + const float a = warp_reduce_sum(acc[j][r]); + if (lane == 0) s_part[warp][j][r] = a; + } + } + __syncthreads(); + for (int idx = threadIdx.x; idx < nt * M; idx += kRnmThreads) { + const int j = idx / M, r = idx % M; + float a = 0.0f, s = 0.0f; +#pragma unroll + for (int q = 0; q < kRnmThreads / WARP_SIZE; ++q) { + a += s_part[q][j][r]; + s += s_part[q][j][M]; + } + out[(int64_t) (t0 + j) * out_stride + r] = a * rsqrtf(s / (float) k + eps); + } +} + +bool ggml_cuda_ds4_rms_norm_mix_supported(int64_t k, int64_t m) { + return k > 0 && k <= INT_MAX && m == 24; +} + +void ggml_cuda_ds4_rms_norm_mix(ggml_backend_cuda_context & ctx, const float * x, int64_t x_stride, + const half * w, int64_t w_stride, float * out, int64_t out_stride, + int64_t k, int64_t m, int64_t n_tokens, float eps) { + GGML_ASSERT(ggml_cuda_ds4_rms_norm_mix_supported(k, m)); + const int blocks = (int) ((n_tokens + kRnmTokens - 1) / kRnmTokens); + ds4_rms_norm_mix_kernel<24><<>>( + x, x_stride, w, w_stride, out, out_stride, (int) k, (int) n_tokens, eps); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh index fc03b57cc..f9b52f651 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh @@ -8,3 +8,30 @@ // but not bit-identical to the CPU HC reference (expf ULP differences // amplified by the sinkhorn iterations). void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// One HC sub-block boundary (rms_norm, fn mix, Sinkhorn pre, post, collapse, +// rms_norm * w) per token in one launch; see ds4-hc.cu. mmv_block is the MMVF +// block size the unfused fn product would launch with (it fixes the +// partial-sum order). +bool ggml_cuda_ds4_hc_boundary_supported(int n_hc, int n_embd, int mmv_block); +void ggml_cuda_ds4_hc_boundary( + ggml_backend_cuda_context & ctx, int n_tokens, int n_hc, int mmv_block, + const float * R, int64_t R_stride, const half * fn, int64_t fn_stride, const float * base, + const float * block_out, int64_t bo_stride, const float * w, + float * pre_dst, int64_t pre_stride, float * R_next, int64_t Rn_stride, + float * out, int64_t out_stride, + int n_embd, float eps_hc, float eps_out, int iters, + float pre_scale, float post_scale, float comb_scale, bool write_pre, + void * q8_out); // optional: out quantized to q8_1 as MMVQ would (rows padded to MATRIX_ROW_PADDING) + +// An RMS norm read only by a product against a few F16 rows (the HC mix of a +// prefill band): out[m, t] = rsqrt(mean_k x[k, t]^2 + eps) * sum_k w[k, m] x[k, t] +// in one pass over x, F32 throughout. m < 32. +bool ggml_cuda_ds4_rms_norm_mix_supported(int64_t k, int64_t m); +void ggml_cuda_ds4_rms_norm_mix(ggml_backend_cuda_context & ctx, const float * x, int64_t x_stride, + const half * w, int64_t w_stride, float * out, int64_t out_stride, + int64_t k, int64_t m, int64_t n_tokens, float eps); + +// The DS4 router's select (mode 4) and weights (mode 5) ops on the same +// logits in one launch, bit-identical to the two. +void ggml_cuda_op_ds4_router_fused(ggml_backend_cuda_context & ctx, ggml_tensor * select, ggml_tensor * weights); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index a412b72f0..cf924cfb6 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -285,17 +285,15 @@ __global__ static void ds4_fa_mean_comp_blocks_kernel( // token. Attention blocks for all query-head groups reuse these four bounds. // Internal masked rows remain inside the envelope and are still evaluated as // masked, so this changes storage only, not attention semantics. +// One query token's visible envelope, computed by one warp (lane = lane id). template -__global__ static void ds4_fa_visibility_bounds_kernel( +static __device__ __forceinline__ void ds4_fa_visibility_bounds_token( const Mask * mask, int * bounds, - int n_tokens, + int t, int n_kv, - int raw_rows) { - const int t = (int) blockIdx.x; - const int lane = (int) threadIdx.x; - if (t >= n_tokens || lane >= warpSize) return; - + int raw_rows, + int lane) { int raw_first = raw_rows; int raw_last = -1; int comp_first = n_kv; @@ -347,6 +345,19 @@ __global__ static void ds4_fa_visibility_bounds_kernel( } } +template +__global__ static void ds4_fa_visibility_bounds_kernel( + const Mask * mask, + int * bounds, + int n_tokens, + int n_kv, + int raw_rows) { + const int t = (int) blockIdx.x; + const int lane = (int) threadIdx.x; + if (t >= n_tokens || lane >= warpSize) return; + ds4_fa_visibility_bounds_token(mask, bounds, t, n_kv, raw_rows, lane); +} + // Exact DS4 ratio-4 layer-major visibility without a materialized mask. // Physical raw rows are [prior chronological SWA | current chunk]. A // compressed row becomes visible only after its four source tokens complete. @@ -581,7 +592,9 @@ __global__ static void ds4_fa_indexed_rows_parallel_kernel( // A shared-memory bitonic sort restores ascending physical-row order, matching // the old top-k -> mask -> physical scan path and therefore preserving each // reduction lane's accumulation order exactly. -template +// WITH_BOUNDS: the first warp then also writes the token's visibility +// envelope (ds4_fa_visibility_bounds_kernel's), saving that launch. +template __global__ static void ds4_fa_indexed_rows_topk_kernel( const Mask * mask, const int32_t * topk, @@ -594,7 +607,8 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel( int raw_rows, int capacity, int kv_start = 0, - int causal_ratio = 4) { + int causal_ratio = 4, + int * bounds = nullptr) { const int t = (int) blockIdx.x; const int tid = (int) threadIdx.x; if (t >= n_tokens) return; @@ -687,25 +701,44 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel( } } } + + if constexpr (WITH_BOUNDS) { + if (tid < warpSize) { + ds4_fa_visibility_bounds_token(mask, bounds, t, n_kv, raw_rows, tid); + } + } } +// bounds: when set, the same launch writes the visibility envelope too. template static void ds4_launch_indexed_rows_topk( const Mask * mask, const int32_t * topk, int * selected_rows, int * selected_counts, int * owner_offsets, int * owner_ranks, int n_tokens, int n_kv, int raw_rows, int capacity, - cudaStream_t stream) { + cudaStream_t stream, int * bounds = nullptr) { // The learned top-512 stays on its original launch. A batched verifier // appends a small saved-raw suffix and needs the next sorting bucket. if (capacity <= 512) { - ds4_fa_indexed_rows_topk_kernel<<>>( - mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks, - n_tokens, n_kv, raw_rows, capacity); + if (bounds) { + ds4_fa_indexed_rows_topk_kernel<<>>( + mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks, + n_tokens, n_kv, raw_rows, capacity, 0, 4, bounds); + } else { + ds4_fa_indexed_rows_topk_kernel<<>>( + mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks, + n_tokens, n_kv, raw_rows, capacity); + } } else { - ds4_fa_indexed_rows_topk_kernel<<>>( - mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks, - n_tokens, n_kv, raw_rows, capacity); + if (bounds) { + ds4_fa_indexed_rows_topk_kernel<<>>( + mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks, + n_tokens, n_kv, raw_rows, capacity, 0, 4, bounds); + } else { + ds4_fa_indexed_rows_topk_kernel<<>>( + mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks, + n_tokens, n_kv, raw_rows, capacity); + } } } @@ -2487,6 +2520,28 @@ static __device__ __forceinline__ const KV * ds4_fa_segmented_row( return preserved_tail + (size_t) (row - compressed_rows) * D; } +// Lane rows (ggml_flash_attn_ext_set_ds4_lane_rows): query token t reads lane +// j's own K/V row in place of ring row lane_row_index[j] for every j <= t. +// Lanes write distinct ring rows, so at most one lane matches. +template +static __device__ __forceinline__ const KV * ds4_fa_lane_row( + const KV * row_ptr, + const KV * lane_kv, + const int * lane_row_index, + int token, + int row, + int raw_rows) { + constexpr int D = 512; + if (row < raw_rows) { + for (int j = 0; j <= token; ++j) { + if (lane_row_index[j] == row) { + return lane_kv + (size_t) j * D; + } + } + } + return row_ptr; +} + // Split-KV decode for indexed MLA. A single grouped block leaves most of a // wide RDNA GPU idle when q is small: DS4 has 64 heads, so the four-head // kernel exposes only 16 blocks per layer. Split the bounded raw+top-k row @@ -2494,7 +2549,8 @@ static __device__ __forceinline__ const KV * ds4_fa_segmented_row( // the states in a second kernel. The implementation is expressed in terms of // the D512 latent-attention contract rather than model weights, so future MLA // models using the same layout can reuse it. -template +template __global__ static void ds4_flash_attn_d512_indexed_split_stage1_kernel( float * partial, const float * q, @@ -2519,10 +2575,13 @@ __global__ static void ds4_flash_attn_d512_indexed_split_stage1_kernel( const int * indexed_counts, int indexed_capacity, ds4_inverse_rope_params inverse_rope, - const float * forward_rope_coefficients) { + const float * forward_rope_coefficients, + const int * lane_rows, + const KV * lane_kv) { constexpr int D = 512; constexpr int N_THREADS = 256; constexpr int VALUES_PER_THREAD = 4; + constexpr int MAX_LANES = 8; const int token = (int) blockIdx.x; const int head_begin = (int) blockIdx.y * HEADS_PER_BLOCK; @@ -2535,6 +2594,14 @@ __global__ static void ds4_flash_attn_d512_indexed_split_stage1_kernel( float * reduction = scores + (size_t) HEADS_PER_BLOCK * split_stride; float * q_rope_tail = reduction + (size_t) HEADS_PER_BLOCK * N_THREADS; + __shared__ int lane_row_index[MAX_LANES]; + if constexpr (LANE_ROWS) { + if (tid <= token && tid < MAX_LANES) { + lane_row_index[tid] = lane_rows[tid]; + } + __syncthreads(); + } + const int raw_first = visibility_bounds ? visibility_bounds[(size_t) token * 4 + 0] : 0; const int raw_last = visibility_bounds @@ -2594,6 +2661,9 @@ __global__ static void ds4_flash_attn_d512_indexed_split_stage1_kernel( k, compressed_kv, preserved_tail_kv, row, materialized_kv_rows, compressed_kv_rows) : k + (size_t) row * D; + if constexpr (LANE_ROWS) { + kr = ds4_fa_lane_row(kr, lane_kv, lane_row_index, token, row, raw_rows); + } float dot[HEADS_PER_BLOCK] = {}; #pragma unroll for (int d = 0; d < D; ++d) { @@ -2690,6 +2760,9 @@ __global__ static void ds4_flash_attn_d512_indexed_split_stage1_kernel( v, compressed_kv, preserved_tail_kv, row, materialized_kv_rows, compressed_kv_rows) : v + (size_t) row * D; + if constexpr (LANE_ROWS) { + vr = ds4_fa_lane_row(vr, lane_kv, lane_row_index, token, row, raw_rows); + } ds4_fa_load_quad(vr + d0, vv0, vv1, vv2, vv3); const int score_index = slot - split_begin; @@ -2819,7 +2892,9 @@ static void ds4_launch_flash_attn_d512_indexed_split( ds4_inverse_rope_params inverse_rope, const float * inverse_rope_coefficients, const float * forward_rope_coefficients, - cudaStream_t stream) { + cudaStream_t stream, + const int * lane_rows = nullptr, + const KV * lane_kv = nullptr) { constexpr int HEADS_PER_BLOCK = 4; const size_t shmem = ((size_t) HEADS_PER_BLOCK * split_stride + @@ -2830,15 +2905,29 @@ static void ds4_launch_flash_attn_d512_indexed_split( (unsigned) n_tokens, (unsigned) (n_heads / HEADS_PER_BLOCK), (unsigned) N_SPLITS); - ds4_flash_attn_d512_indexed_split_stage1_kernel< - KV, Mask, HEADS_PER_BLOCK, N_SPLITS> - <<>>( - partial, q, q_stride_token, q_stride_head, k, v, - compressed_kv, preserved_tail_kv, mask, sinks, - n_tokens, n_heads, n_kv, materialized_kv_rows, - compressed_kv_rows, scale, raw_rows, split_stride, - visibility_bounds, indexed_rows, indexed_counts, - indexed_capacity, inverse_rope, forward_rope_coefficients); + if (lane_kv) { + ds4_flash_attn_d512_indexed_split_stage1_kernel< + KV, Mask, HEADS_PER_BLOCK, N_SPLITS, true> + <<>>( + partial, q, q_stride_token, q_stride_head, k, v, + compressed_kv, preserved_tail_kv, mask, sinks, + n_tokens, n_heads, n_kv, materialized_kv_rows, + compressed_kv_rows, scale, raw_rows, split_stride, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, forward_rope_coefficients, + lane_rows, lane_kv); + } else { + ds4_flash_attn_d512_indexed_split_stage1_kernel< + KV, Mask, HEADS_PER_BLOCK, N_SPLITS> + <<>>( + partial, q, q_stride_token, q_stride_head, k, v, + compressed_kv, preserved_tail_kv, mask, sinks, + n_tokens, n_heads, n_kv, materialized_kv_rows, + compressed_kv_rows, scale, raw_rows, split_stride, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, forward_rope_coefficients, + nullptr, nullptr); + } const dim3 stage2_grid((unsigned) n_tokens, (unsigned) n_heads, 1); ds4_flash_attn_d512_indexed_split_stage2_kernel <<>>( @@ -3121,7 +3210,26 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) const int rope_flags = (int) (packed_rope_flags & 0xffffu); const int causal_ratio = (int) (packed_rope_flags >> 16); const ggml_tensor * rope_positions = dst->src[6]; - if (rope_positions && + const ggml_tensor * lane_kv = dst->src[9]; + if ((rope_flags & 4) != 0) { + // Lane rows: src[6] holds the lanes' ring rows (I32), src[9] their + // K/V rows. Only the split-KV decode schedule implements them. + if ((rope_flags & 3) != 0 || causal_ratio != 0 || + !rope_positions || !lane_kv || + rope_positions->type != GGML_TYPE_I32 || + rope_positions->ne[0] != n_tokens || rope_positions->ne[1] != 1 || + rope_positions->ne[2] != 1 || rope_positions->ne[3] != 1 || + !ggml_is_contiguous(rope_positions) || + !segmented_kv || !indexer_topk || !mask || + mask->type != GGML_TYPE_F16 || !kv_f16 || n_tokens > 8 || + lane_kv->type != K->type || lane_kv->ne[0] != K->ne[0] || + lane_kv->ne[1] != n_tokens || lane_kv->ne[2] != 1 || + lane_kv->ne[3] != 1 || !ggml_is_contiguous(lane_kv)) { + return false; + } + } else if (lane_kv) { + return false; + } else if (rope_positions && ((rope_flags & 1) == 0 || causal_ratio != 0 || rope_positions->type != GGML_TYPE_I32 || rope_positions->ne[0] != n_tokens || rope_positions->ne[1] != 1 || @@ -3141,7 +3249,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) (ds4_layout != 0 && (raw_window <= 0 || sparse_block_size <= 0)) || (sparse_keep_rows != 0 && !mask && !ratio4_causal && causal_ratio == 0) || - (rope_flags & ~3) != 0 || + (rope_flags & ~7) != 0 || ((rope_flags & 2) != 0 && (rope_flags & 1) == 0)) { return false; } @@ -3260,6 +3368,10 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( const int rope_flags = (int) (packed_rope_flags & 0xffffu); const int causal_ratio = (int) (packed_rope_flags >> 16); const int causal_kv_start = ggml_get_op_params_i32(dst, 8); + // Lane rows (bit 2) carry the lanes' ring rows in src[6] and their K/V + // rows in src[9]; supported() admits them only on the split-KV schedule. + const ggml_tensor * lane_rows = (rope_flags & 4) ? dst->src[6] : nullptr; + const ggml_tensor * lane_kv = lane_rows ? dst->src[9] : nullptr; const int prior_rows = raw_rows - n_tokens; const bool contiguous_causal = !mask && !indexer_topk && causal_ratio > 0 && n_tokens > raw_window && @@ -3293,7 +3405,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( ds4_inverse_rope_params inverse_rope{}; inverse_rope.enabled = rope_flags & 1; inverse_rope.forward_q_enabled = (rope_flags & 2) != 0; - if (rope_flags != 0) { + if ((rope_flags & 3) != 0) { inverse_rope.kv_start = ggml_get_op_params_i32(dst, 8); inverse_rope.positions = dst->src[6] ? static_cast(dst->src[6]->data) : nullptr; @@ -3452,6 +3564,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( int * indexed_counts = nullptr; int * indexed_owner_offsets = nullptr; int * indexed_owner_ranks = nullptr; + bool bounds_with_rows = false; if (indexed_mask) { indexed_rows = indexed_rows_alloc.alloc( (size_t) n_tokens * indexed_capacity); @@ -3472,12 +3585,15 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( inverse_rope.kv_start, indexed_causal_ratio); } else if (mask->type == GGML_TYPE_F16) { if (indexer_topk) { + // The envelope scan below reads the same mask: one launch. + bounds_with_rows = !contiguous_causal; ds4_launch_indexed_rows_topk( (const half *) mask->data, (const int32_t *) indexer_topk->data, indexed_rows, indexed_counts, indexed_owner_offsets, indexed_owner_ranks, - n_tokens, n_kv, raw_rows, indexed_capacity, stream); + n_tokens, n_kv, raw_rows, indexed_capacity, stream, + bounds_with_rows ? visibility_bounds : nullptr); } else if (parallel_index_scan) { ds4_fa_indexed_rows_parallel_kernel<<>>( (const half *) mask->data, indexed_rows, indexed_counts, @@ -3493,12 +3609,14 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( } } else { if (indexer_topk) { + bounds_with_rows = !contiguous_causal; ds4_launch_indexed_rows_topk( (const float *) mask->data, (const int32_t *) indexer_topk->data, indexed_rows, indexed_counts, indexed_owner_offsets, indexed_owner_ranks, - n_tokens, n_kv, raw_rows, indexed_capacity, stream); + n_tokens, n_kv, raw_rows, indexed_capacity, stream, + bounds_with_rows ? visibility_bounds : nullptr); } else if (parallel_index_scan) { ds4_fa_indexed_rows_parallel_kernel<<>>( (const float *) mask->data, indexed_rows, indexed_counts, @@ -3525,6 +3643,8 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( (n_tokens + 255) / 256, 256, 0, stream>>>( visibility_bounds, n_tokens, n_kv, raw_rows, raw_window, causal_kv_start, causal_ratio); + } else if (bounds_with_rows) { + // written by the indexed-rows launch above } else if (mask->type == GGML_TYPE_F16) { ds4_fa_visibility_bounds_kernel<<>>( (const half *) mask->data, visibility_bounds, @@ -3826,7 +3946,11 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( visibility_bounds, indexed_rows, indexed_counts, indexed_capacity, inverse_rope, inverse_rope_coefficients, forward_rope_coefficients, - stream); + stream, + lane_rows ? (const int *) lane_rows->data : nullptr, + lane_rows ? (const half *) lane_kv->data : nullptr); + } else if (lane_rows) { + return false; } else if (kv_f32 && mask->type == GGML_TYPE_F32) { ds4_launch_flash_attn_d512_indexed_split< float, float, split_count>( diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index a745a0df3..73d11ac27 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -175,6 +175,11 @@ extern "C" bool ggml_backend_cuda_set_mmvq_batch_invariant(bool enabled) { return previous; } +extern "C" void ggml_backend_cuda_mul_mat_whole_row_lanes(ggml_tensor * mul_mat) { + GGML_ASSERT(mul_mat != nullptr && mul_mat->op == GGML_OP_MUL_MAT); + ggml_set_op_params_i32(mul_mat, 1, GGML_CUDA_WHOLE_ROW_LANES); +} + bool ggml_cuda_mmvq_batch_invariant() { return ggml_cuda_mmvq_batch_invariant_enabled; } @@ -972,6 +977,16 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { while (!luce_q8_memo.empty()) { luce_q8_memo.pop_back(); } + if (ds4_hc_barrier != nullptr) { + ggml_cuda_set_device(device); + CUDA_CHECK(cudaFree(ds4_hc_barrier)); + ds4_hc_barrier = nullptr; + } + if (ds4_hc_split_scratch != nullptr) { + ggml_cuda_set_device(device); + CUDA_CHECK(cudaFree(ds4_hc_split_scratch)); + ds4_hc_split_scratch = nullptr; + } #if defined(GGML_HIP_DS4V_VISION) if (vision_bias_workspace) { @@ -3060,6 +3075,11 @@ static bool ggml_cuda_try_fuse_mul_mat_glu( const ggml_tensor * src1 = up->src[1]; const ggml_tensor * ids = up->src[2]; + // The fused kernels read F32 activations; q8_1 rows (an MoE owner's input + // quantized on another device) take the unfused MMVQ dispatch. + if (src1->type != GGML_TYPE_F32) { + return false; + } #if defined(GGML_USE_HIP) if (ids && ggml_cuda_mmb_supported_glu(gate->src[0], up->src[0], src1, ids, glu)) { ggml_cuda_mul_mat_id_mmb_glu(ctx, gate->src[0], up->src[0], src1, ids, glu); @@ -3145,6 +3165,19 @@ static bool ggml_cuda_try_fuse_mul_mat_glu( return true; } + // RDNA4 Q8_0 gate/up pairs at several columns (the cluster shared expert in + // a verify): one launch reads both matrices and writes the SwiGLU-DS4. + if (!ids && direct_vector_layout && src1->ne[1] <= ggml_cuda_mmvq_max_ncols() && + !ggml_backend_buft_is_cuda_split(src0->buffer->buft) && + !ggml_backend_buft_is_cuda_split(gate->src[0]->buffer->buft) && + ggml_cuda_mmvq_rdna4_glu_pair(gate, up, glu)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, glu, &fusion_data); + return true; + } + if (direct_vector_layout && ggml_cuda_should_fuse_mul_mat_vec_q(up)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate->src[0]; @@ -3240,8 +3273,60 @@ static bool ggml_cuda_mmb_cublas_shape_ok(const ggml_tensor * w) { (w->ne[0] == 10240 && w->ne[1] == 320) || (w->ne[0] == 320 && w->ne[1] == 10240)); } +// The BF16 shadow of src0 when ggml_cuda_mul_mat runs the product as a BF16 +// cuBLAS GEMM on it (qualified qwen4exp shapes), else null. +static const void * ggml_cuda_mul_mat_bf16_shadow(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { + if (!ggml_cuda_qwen4exp_enabled() || ggml_backend_buft_is_cuda_split(src0->buffer->buft) || + ggml_mul_mat_is_grouped_src(dst) || src0->ne[2] != 1 || src0->ne[3] != 1 || + (src0->type != GGML_TYPE_IQ4_NL && src0->type != GGML_TYPE_Q6_K && src0->type != GGML_TYPE_Q5_K) || + src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + !ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || + !ggml_cuda_mmb_cublas_shape_ok(src0) || src1->ne[1] * src1->ne[2] * src1->ne[3] < 512) { + return nullptr; + } + return ggml_cuda_mmb_shadow_ptr(src0); +} #endif // defined(GGML_USE_HIP) +// If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. +// But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. +// Therefore, in such cases use cuBLAS. +static bool ggml_cuda_src0_bad_padding_clear(const ggml_tensor * src0) { + return ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE + && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; +} + +// Whether ggml_cuda_mul_mat may run this product on MMVQ, before a device +// chooses between MMVQ and MMQ (ggml_cuda_mul_mat_prefers_mmq). The mix +// qtypes have no generic MMVQ path because their per-expert codebooks live in +// an out-of-band registry, and MXFP8 has its own matvec. +static bool ggml_cuda_mul_mat_vec_q_eligible(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { + if (!ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_MXFP8 || + src0->type == GGML_TYPE_Q3_1_ROCMFP3_MIX || src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX || + ggml_cuda_src0_bad_padding_clear(src0) || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + src1->ne[1] > ggml_cuda_mmvq_max_ncols()) { + return false; + } +#ifdef ROCMFP2_AFFINE + // The affine MMVQ dot carries the offset correction (vecdotq.cuh); it + // stays behind its A/B switch. + if (src0->type == GGML_TYPE_Q2_0_ROCMFP2 && std::getenv("LUCE_CUDA_MMVQ_FP2_AFFINE") == nullptr) { + return false; + } +#endif // ROCMFP2_AFFINE + return true; +} + +// Rows from which RDNA4 runs a Q4_K product of two or more columns on MMQ: +// MMQ streams such tall matrices faster than MMVQ, which keeps shorter ones. +#define GGML_CUDA_RDNA4_Q4_K_MMQ_MIN_ROWS 2048 + +// Whether a product that both MMVQ and MMQ can run takes MMQ on device cc. +static bool ggml_cuda_mul_mat_prefers_mmq(const ggml_tensor * src0, const ggml_tensor * src1, const int cc) { + return src0->type == GGML_TYPE_Q4_K && GGML_CUDA_CC_IS_RDNA4(cc) && src1->ne[1] > 1 && + src0->ne[1] >= GGML_CUDA_RDNA4_Q4_K_MMQ_MIN_ROWS; +} + static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); const bool grouped_src = ggml_mul_mat_is_grouped_src(dst); @@ -3261,13 +3346,8 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor } #if defined(GGML_USE_HIP) - const bool cublas_shape_ok = ggml_cuda_mmb_cublas_shape_ok(src0); - if (ggml_cuda_qwen4exp_enabled() && !split && !grouped_src && src0->ne[2] == 1 && src0->ne[3] == 1 && - (src0->type == GGML_TYPE_IQ4_NL || src0->type == GGML_TYPE_Q6_K || src0->type == GGML_TYPE_Q5_K) && - src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 && - ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && - cublas_shape_ok && src1->ne[1] * src1->ne[2] * src1->ne[3] >= 512) { - const void * sh = ggml_cuda_mmb_shadow_ptr(src0); + { + const void * sh = ggml_cuda_mul_mat_bf16_shadow(src0, src1, dst); if (sh) { ggml_tensor tmp = *src0; tmp.type = GGML_TYPE_BF16; @@ -3295,11 +3375,7 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor } #endif // defined(GGML_USE_HIP) - // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. - // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. - // Therefore, in such cases use cuBLAS. - const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE - && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; + const bool bad_padding_clear = ggml_cuda_src0_bad_padding_clear(src0); bool use_mul_mat_vec_f = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; @@ -3319,9 +3395,7 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor const bool is_rocmfp3_mix = src0->type == GGML_TYPE_Q3_1_ROCMFP3_MIX; const bool is_rocmfp2_mix = src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX; const bool is_mix_qtype = is_rocmfp3_mix || is_rocmfp2_mix; - bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !is_mix_qtype && !bad_padding_clear - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 - && src1->ne[1] <= luce_mmvq_max_ncols; + bool use_mul_mat_vec_q = ggml_cuda_mul_mat_vec_q_eligible(src0, src1, dst); bool use_mul_mat_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; @@ -3332,9 +3406,6 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor // throughput. The phase-scoped master switch still enables the qualified // paired gate/up fusion independently. if (src0->type == GGML_TYPE_Q2_0_ROCMFP2) { - if (std::getenv("LUCE_CUDA_MMVQ_FP2_AFFINE") == nullptr) { - use_mul_mat_vec_q = false; - } if (std::getenv("LUCE_CUDA_MMQ_FP2_AFFINE") == nullptr || std::getenv("LUCE_CUDA_MMQ_FP2_AFFINE_GENERAL") == nullptr) { use_mul_mat_q = false; @@ -3367,6 +3438,9 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); + if (use_mul_mat_vec_q && use_mul_mat_q && ggml_cuda_mul_mat_prefers_mmq(src0, src1, cc)) { + use_mul_mat_vec_q = false; + } } // debug helpers @@ -3501,7 +3575,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * const ggml_tensor * src1 = dst->src[1]; const ggml_tensor * ids = dst->src[2]; - GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_Q8_1); GGML_ASSERT(dst->type == GGML_TYPE_F32); GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft) && "mul_mat_id does not support split buffers"); @@ -3525,6 +3599,14 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * } }; + if (src1->type == GGML_TYPE_Q8_1) { + // Activations already quantized into MMVQ's rows (by the main owner). + GGML_ASSERT(ggml_is_quantized(src0->type) && ne2 <= mmvq_mmid_max); + log_dispatch("mmvq-q8"); + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); + return; + } + // qtype-105/106 fused routed-expert decode: run the whole // mul_mat_id on device with the routing ids read in-kernel, avoiding the // generic sort-based fallback below (which needs a host cudaStreamSynchronize @@ -4308,8 +4390,21 @@ struct ggml_cuda_pending_peer_copy_batch { ggml_backend_cuda_context * dst = nullptr; size_t bytes = 0; int copies = 0; + // GGML_CUDA_PEER_COPY_KERNEL=1: the copies are collected here and leave + // as one kernel on the producer stream at the flush (one launch instead + // of a runtime blit per tensor, the per-layer MoE handoff carries three). + int n_descs = 0; + ggml_cuda_copy_desc descs[ggml_cuda_copy_batch_max]; }; +static bool ggml_cuda_peer_copy_kernel_enabled() { + static const bool enabled = [] { + const char * value = getenv("GGML_CUDA_PEER_COPY_KERNEL"); + return value && *value && strcmp(value, "0") != 0; + }(); + return enabled; +} + static thread_local ggml_cuda_pending_peer_copy_batch ggml_cuda_pending_peer_copies; @@ -4334,6 +4429,12 @@ static void ggml_cuda_flush_peer_copy_batch(const char * reason) { GGML_ASSERT(batch.src && batch.dst); ggml_cuda_set_device(batch.src->device); + if (batch.n_descs > 0) { + // Peer access is enabled, so a kernel on the producer writes the + // consumer's memory directly; one launch for every collected copy. + ggml_cuda_copy_batch(batch.descs, batch.n_descs, batch.src->stream()); + batch.n_descs = 0; + } if (!batch.src->copy_event) { CUDA_CHECK(cudaEventCreateWithFlags( &batch.src->copy_event, cudaEventDisableTiming)); @@ -4504,10 +4605,24 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ ggml_cuda_flush_peer_copy_batch("peer-pair-change"); } ggml_cuda_set_device(cuda_ctx_src->device); - CUDA_CHECK(cudaMemcpyPeerAsync( - dst->data, cuda_ctx_dst->device, - src->data, cuda_ctx_src->device, - ggml_nbytes(dst), cuda_ctx_src->stream())); + if (ggml_cuda_peer_copy_kernel_enabled() && + ggml_is_contiguous(src) && ggml_is_contiguous(dst)) { + if (batch.n_descs == ggml_cuda_copy_batch_max) { + ggml_cuda_copy_batch(batch.descs, batch.n_descs, cuda_ctx_src->stream()); + batch.n_descs = 0; + } + batch.descs[batch.n_descs++] = {src->data, dst->data, ggml_nbytes(dst)}; + } else { + if (batch.n_descs > 0) { + // Keep the producer-stream order of earlier copies. + ggml_cuda_copy_batch(batch.descs, batch.n_descs, cuda_ctx_src->stream()); + batch.n_descs = 0; + } + CUDA_CHECK(cudaMemcpyPeerAsync( + dst->data, cuda_ctx_dst->device, + src->data, cuda_ctx_src->device, + ggml_nbytes(dst), cuda_ctx_src->stream())); + } batch.src = cuda_ctx_src; batch.dst = cuda_ctx_dst; batch.bytes += ggml_nbytes(dst); @@ -5350,6 +5465,390 @@ static bool ggml_cuda_node_is_noop(const ggml_tensor * node) { return ggml_cuda_node_launches_nothing(node) || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0; } +// ── DS4 HC sub-block boundary in one launch ────────────────────────────── +// Matches, at node i, the six launches the graph builds between a staggered +// HC sub-block and the next one (views between them launch nothing): +// n1 = RMS_NORM(R) mm = MUL_MAT(fn F16, n1) pre = DS4_HC pre(mm, base, R) +// post = DS4_HC post(R, block_out, split = pre[n_embd:]) coll = DS4_HC collapse(post, pre[n_embd:n_embd+n_hc]) +// n2 = RMS_NORM(coll) out = MUL(n2, w) +// and runs ggml_cuda_ds4_hc_boundary, which writes post and out (and pre +// when anything outside the group reads it) bit-identically. n1, mm, coll +// and n2 must have no reader outside the group. + +static int32_t ggml_cuda_graph_use_count(const ggml_cgraph * cgraph, const ggml_tensor * t) { + const size_t pos = ggml_hash_find(&cgraph->visited_hash_set, t); + if (pos == GGML_HASHSET_FULL || !ggml_bitset_get(cgraph->visited_hash_set.used, pos) || + cgraph->visited_hash_set.keys[pos] != t) { + return -1; + } + return cgraph->use_counts[pos]; +} + +// A node computed inside a fused group whose value nothing else reads. +static bool ggml_cuda_elidable(const ggml_cgraph * cgraph, const ggml_tensor * t, int32_t uses) { + return t && !t->view_src && (t->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && + ggml_cuda_graph_use_count(cgraph, t) == uses; +} + +// Row stride in floats of a [ne0, T] F32 tensor whose rows are dense. +static int64_t ggml_cuda_f32_row_stride(const ggml_tensor * t) { + return t->ne[1] > 1 ? (int64_t) (t->nb[1] / sizeof(float)) : t->ne[0]; +} + +// Whether ggml_cuda_mul_mat runs this MUL_MAT node on its plain single-device +// MMVQ branch, by the same tests it dispatches with. A fusion that runs MMVQ +// in the node's place, or hands MMVQ a q8_1 activation, must decline +// otherwise: another kernel would give other bits than the unfused graph. +static bool ggml_cuda_mul_mat_takes_mmvq(const ggml_backend_cuda_context & ctx, const ggml_tensor * mm) { + const ggml_tensor * src0 = mm->src[0]; + const ggml_tensor * src1 = mm->src[1]; + if (mm->op != GGML_OP_MUL_MAT || ggml_backend_buft_is_cuda_split(src0->buffer->buft) || + !ggml_cuda_mul_mat_vec_q_eligible(src0, src1, mm)) { + return false; + } +#if defined(GGML_USE_HIP) + if (ggml_cuda_mul_mat_bf16_shadow(src0, src1, mm) != nullptr) { + return false; + } +#endif // defined(GGML_USE_HIP) + // MMQ's own conditions reduce to this one for an MMVQ-eligible product. + const int cc = ggml_cuda_info().devices[ctx.device].cc; + return !(ggml_cuda_mul_mat_prefers_mmq(src0, src1, cc) && + ggml_cuda_should_use_mmq(mm, cc, src1->ne[1], /*n_experts=*/0)); +} + +// When the node at `idx` (contiguous F32 rows) is read next by a quantized +// matmul, its producer can write the q8_1 form MMVQ would quantize it to: +// this registers a memo entry under MMVQ's key and returns its buffer (null +// when not worthwhile). Rows are padded to MATRIX_ROW_PADDING as MMVQ pads them. +static void ggml_cuda_q8_memo_for_producer(ggml_backend_cuda_context & ctx, const ggml_cgraph * cgraph, + int idx, void ** q8_out, int * ncols_q8) { + *q8_out = nullptr; + *ncols_q8 = 0; + const ggml_tensor * t = cgraph->nodes[idx]; + if (!ggml_cuda_luce_q8_memo_on() || !ctx.stream_context().concurrent_events.empty() || + t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || t->buffer == nullptr || + t->ne[0] % QK8_1 != 0) { + return; + } + int j = idx + 1; + while (j < cgraph->n_nodes && ggml_cuda_node_is_noop(cgraph->nodes[j])) ++j; + if (j >= cgraph->n_nodes) return; + const ggml_tensor * c = cgraph->nodes[j]; + if ((c->op != GGML_OP_MUL_MAT && c->op != GGML_OP_MUL_MAT_ID) || !ggml_is_quantized(c->src[0]->type)) return; + const ggml_tensor * src1 = c->src[1]; + while (src1 && src1->op == GGML_OP_RESHAPE && src1 != t) src1 = src1->src[0]; + if (src1 != t) return; + const int64_t padded = GGML_PAD(t->ne[0], MATRIX_ROW_PADDING); + const int64_t nrows = ggml_nrows(t); + ggml_backend_cuda_context::luce_q8_memo_entry ent; + ent.src1_node = (const void *) t; + ent.src1_data = t->data; + ent.src0_type = -1; + ent.ne[0] = t->ne[0]; ent.ne[1] = nrows; ent.ne[2] = 1; ent.ne[3] = 1; + ent.buf = std::make_unique>(ctx.pool(), (size_t) nrows * padded / QK8_1 * sizeof(block_q8_1)); + *q8_out = ent.buf->ptr; + *ncols_q8 = (int) padded; + ctx.luce_q8_memo.push_back(std::move(ent)); +} + +// RMS_NORM, MUL (one F32 weight row), a forward NORMAL/TAIL ROPE of that and +// an F16 copy of the ROPE result (the per-token K/V latent the flash lanes +// read) in one launch; the ROPE output is still written for its other readers. +static int ggml_cuda_try_rms_norm_rope_f16(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { + const int n = cgraph->n_nodes; + auto next = [&](int j) { + while (j < n && ggml_cuda_node_is_noop(cgraph->nodes[j])) ++j; + return j; + }; + ggml_tensor * norm = cgraph->nodes[i]; + if (i + 1 >= n) return -1; + ggml_tensor * mul = cgraph->nodes[i + 1]; + if (mul->op != GGML_OP_MUL || mul->src[0] != norm) return -1; + const int i_rope = next(i + 2); if (i_rope >= n) return -1; + ggml_tensor * rope = cgraph->nodes[i_rope]; + const int i_cpy = next(i_rope + 1); if (i_cpy >= n) return -1; + ggml_tensor * cpy = cgraph->nodes[i_cpy]; + if (rope->op != GGML_OP_ROPE || cpy->op != GGML_OP_CPY) return -1; + + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * w = mul->src[1]; + const int64_t ncols = x->ne[0]; + if (x->type != GGML_TYPE_F32 || x->nb[0] != sizeof(float) || x->ne[2] != 1 || x->ne[3] != 1 || + !ggml_is_contiguous(norm) || !ggml_is_contiguous(mul) || + w->type != GGML_TYPE_F32 || !ggml_is_contiguous(w) || w->ne[0] != ncols || ggml_nrows(w) != 1 || + !ggml_cuda_rms_norm_rope_f16_supported((int) ncols)) return -1; + // rope reads mul (directly or through one reshape) and copies to F16. + const ggml_tensor * rope_src = rope->src[0]; + if (rope_src != mul && !(rope_src->op == GGML_OP_RESHAPE && rope_src->src[0] == mul && + ggml_cuda_graph_use_count(cgraph, rope_src) == 1)) return -1; + const int mode = ((const int32_t *) rope->op_params)[2]; + if ((mode & ~GGML_ROPE_TYPE_TAIL) != GGML_ROPE_TYPE_NORMAL || rope->src[2] != nullptr || + rope->type != GGML_TYPE_F32 || rope_src->type != GGML_TYPE_F32 || !ggml_is_contiguous(rope) || + rope_src->ne[0] != ncols || ggml_nrows(rope_src) != ggml_nrows(x) || rope_src->ne[3] != 1 || + rope->src[1]->type != GGML_TYPE_I32) return -1; + const ggml_tensor * cpy_src = cpy->src[0]; + if (cpy_src != rope && !(cpy_src->op == GGML_OP_RESHAPE && cpy_src->src[0] == rope)) return -1; + if (cpy->type != GGML_TYPE_F16 || !ggml_is_contiguous(cpy) || ggml_nelements(cpy) != ggml_nelements(rope)) return -1; + // The CPY node writes into its src[1] (ggml_cpy's destination); ggml_cast's + // is the node itself. + if (cpy->src[1] && cpy->src[1]->data != cpy->data) return -1; + if (!ggml_cuda_elidable(cgraph, norm, 1) || !ggml_cuda_elidable(cgraph, mul, 1)) return -1; + + ggml_cuda_rms_norm_rope_f16(ctx, norm, mul, rope, cpy); + return i_cpy; +} + +// A prefill band's HC mix: RMS_NORM read only by a MUL_MAT against the 24 +// F16 mix rows, in one pass over the residual (ggml_cuda_ds4_rms_norm_mix) +// instead of a norm, an F16 copy and a GEMM of 24 rows. Only for bands of 64 +// tokens or more: decode and verify keep their batch-invariant products. +static int ggml_cuda_try_ds4_rms_norm_mix(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { + ggml_tensor * norm = cgraph->nodes[i]; + int j = i + 1; + while (j < cgraph->n_nodes && ggml_cuda_node_is_noop(cgraph->nodes[j])) ++j; + if (j >= cgraph->n_nodes) return -1; + ggml_tensor * mm = cgraph->nodes[j]; + const ggml_tensor * x = norm->src[0]; + if (mm->op != GGML_OP_MUL_MAT || mm->src[1] != norm || mm->type != GGML_TYPE_F32 || + !ggml_is_contiguous(mm)) return -1; + const ggml_tensor * w = mm->src[0]; + if (x->type != GGML_TYPE_F32 || x->nb[0] != sizeof(float) || x->ne[2] != 1 || x->ne[3] != 1 || + x->ne[1] < 64 || !ggml_is_contiguous(norm) || + w->type != GGML_TYPE_F16 || w->nb[0] != sizeof(half) || w->ne[0] != x->ne[0] || + w->ne[2] != 1 || w->ne[3] != 1 || mm->ne[1] != x->ne[1] || + !ggml_cuda_ds4_rms_norm_mix_supported(x->ne[0], w->ne[1]) || + !ggml_cuda_elidable(cgraph, norm, 1)) return -1; + float eps = 0.0f; + memcpy(&eps, norm->op_params, sizeof(float)); + ggml_cuda_ds4_rms_norm_mix(ctx, (const float *) x->data, (int64_t) (x->nb[1] / sizeof(float)), + (const half *) w->data, (int64_t) (w->nb[1] / sizeof(half)), + (float *) mm->data, (int64_t) (mm->nb[1] / sizeof(float)), + x->ne[0], w->ne[1], x->ne[1], eps); + return j; +} + +// A CONT of a strided F32 view read only by a quantized matmul that takes +// MMVQ: MMVQ would quantize the contiguous copy, so quantize the view itself +// into its memo entry (the same q8_1 bytes: rows of whole, unpadded blocks in +// the same flattened order) and run the matmul on it, without the copy +// launch. The copy's output is never materialized, so it must have no other +// reader. +static int ggml_cuda_try_cont_into_mmvq(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { + ggml_tensor * c = cgraph->nodes[i]; + if (c->op != GGML_OP_CONT) return -1; + const ggml_tensor * v = c->src[0]; + if (c->type != GGML_TYPE_F32 || v->type != GGML_TYPE_F32 || !ggml_is_contiguous(c) || + v->nb[0] != sizeof(float) || ggml_is_contiguous(v)) return -1; + if (!ggml_cuda_luce_q8_memo_on() || !ctx.stream_context().concurrent_events.empty() || + c->buffer == nullptr) return -1; + int j = i + 1; + while (j < cgraph->n_nodes && ggml_cuda_node_is_noop(cgraph->nodes[j])) ++j; + if (j >= cgraph->n_nodes) return -1; + ggml_tensor * mm = cgraph->nodes[j]; + if (mm->op != GGML_OP_MUL_MAT || mm->type != GGML_TYPE_F32) return -1; + const ggml_tensor * src0 = mm->src[0]; + const ggml_tensor * src1 = mm->src[1]; + // src1 is c or a chain of single-reader reshapes of it. + const ggml_tensor * walk = src1; + while (walk != c && walk->op == GGML_OP_RESHAPE && walk->src[0] && + ggml_cuda_graph_use_count(cgraph, walk) == 1) { + walk = walk->src[0]; + } + if (walk != c || !ggml_is_contiguous(src1) || !ggml_cuda_elidable(cgraph, c, 1)) return -1; + if (!ggml_cuda_mul_mat_takes_mmvq(ctx, mm)) return -1; + // Whole unpadded q8_1 rows on both sides. + if (v->ne[0] % MATRIX_ROW_PADDING != 0 || src1->ne[0] % MATRIX_ROW_PADDING != 0 || + src1->ne[0] % v->ne[0] != 0) return -1; + + ggml_backend_cuda_context::luce_q8_memo_entry ent; + ent.src1_node = (const void *) c; + ent.src1_data = c->data; + ent.src0_type = -1; + ent.ne[0] = src1->ne[0]; ent.ne[1] = ggml_nrows(src1); ent.ne[2] = 1; ent.ne[3] = 1; + ent.buf = std::make_unique>( + ctx.pool(), (size_t) ggml_nelements(c) / QK8_1 * sizeof(block_q8_1)); + char * q8 = ent.buf->ptr; + ctx.luce_q8_memo.push_back(std::move(ent)); + quantize_row_q8_1_cuda((const float *) v->data, nullptr, q8, src0->type, v->ne[0], + v->nb[1] / sizeof(float), v->nb[2] / sizeof(float), v->nb[3] / sizeof(float), + v->ne[0], v->ne[1], v->ne[2], v->ne[3], ctx.stream()); + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, mm, nullptr); + return j; +} + +// A forward NORMAL/TAIL ROPE whose output [D, H, T] is read next by an MMVQ +// matmul through a grouped view [D * hpg, T, H / hpg] (the DS4 grouped output +// projection): the rope writes the q8_1 form MMVQ would quantize that view to, +// under MMVQ's memo key, so the matmul skips its quantize launch. +static bool ggml_cuda_try_rope_q8(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { + const ggml_tensor * rope = cgraph->nodes[i]; + const ggml_tensor * x = rope->src[0]; + const int mode = ((const int32_t *) rope->op_params)[2]; + if ((mode & ~GGML_ROPE_TYPE_TAIL) != GGML_ROPE_TYPE_NORMAL || rope->src[2] != nullptr || + rope->type != GGML_TYPE_F32 || x->type != GGML_TYPE_F32 || !ggml_is_contiguous(rope) || + x->nb[0] != sizeof(float) || rope->ne[3] != 1 || rope->ne[0] % QK8_1 != 0 || + rope->src[1]->type != GGML_TYPE_I32 || rope->buffer == nullptr) return false; + if (!ggml_cuda_luce_q8_memo_on() || !ctx.stream_context().concurrent_events.empty()) return false; + int j = i + 1; + while (j < cgraph->n_nodes && ggml_cuda_node_is_noop(cgraph->nodes[j])) ++j; + if (j >= cgraph->n_nodes) return false; + const ggml_tensor * mm = cgraph->nodes[j]; + if (mm->op != GGML_OP_MUL_MAT || mm->type != GGML_TYPE_F32) return false; + const ggml_tensor * src0 = mm->src[0]; + const ggml_tensor * src1 = mm->src[1]; + const int64_t D = rope->ne[0], H = rope->ne[1], T = rope->ne[2]; + const int64_t G = src1->ne[0]; + if (src1->type != GGML_TYPE_F32 || src1->data != rope->data || ggml_is_contiguous(src1) || + G % D != 0 || G % MATRIX_ROW_PADDING != 0 || H % (G / D) != 0 || + src1->ne[1] != T || src1->ne[2] != H / (G / D) || src1->ne[3] != 1 || + src1->nb[0] != sizeof(float) || src1->nb[1] != (size_t) (D * H) * sizeof(float) || + src1->nb[2] != (size_t) G * sizeof(float)) return false; + if (!ggml_cuda_mul_mat_takes_mmvq(ctx, mm)) return false; + + ggml_backend_cuda_context::luce_q8_memo_entry ent; + ent.src1_node = (const void *) src1; + ent.src1_data = src1->data; + ent.src0_type = (int) src0->type; + ent.ne[0] = src1->ne[0]; ent.ne[1] = src1->ne[1]; ent.ne[2] = src1->ne[2]; ent.ne[3] = src1->ne[3]; + ent.buf = std::make_unique>( + ctx.pool(), (size_t) ggml_nelements(src1) / QK8_1 * sizeof(block_q8_1)); + void * q8 = ent.buf->ptr; + ctx.luce_q8_memo.push_back(std::move(ent)); + ggml_cuda_rope_q8(ctx, rope, q8, (int) (G / D)); + return true; +} + +static int ggml_cuda_try_ds4_hc_boundary(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { + const int n = cgraph->n_nodes; + // The next node at or after j that launches something. + auto next = [&](int j) { + while (j < n && ggml_cuda_node_is_noop(cgraph->nodes[j])) ++j; + return j; + }; + const auto is_hc = [](const ggml_tensor * t, int mode) { + return t->op == GGML_OP_DS4_HC && ggml_get_op_params_i32(t, 0) == mode; + }; + + ggml_tensor * n1 = cgraph->nodes[i]; + if (n1->op != GGML_OP_RMS_NORM) return -1; + const int i_mm = next(i + 1); if (i_mm >= n) return -1; + ggml_tensor * mm = cgraph->nodes[i_mm]; + const int i_pre = next(i_mm + 1); if (i_pre >= n) return -1; + ggml_tensor * pre = cgraph->nodes[i_pre]; + const int i_post = next(i_pre + 1); if (i_post >= n) return -1; + ggml_tensor * post = cgraph->nodes[i_post]; + const int i_coll = next(i_post + 1); if (i_coll >= n) return -1; + ggml_tensor * coll = cgraph->nodes[i_coll]; + const int i_n2 = next(i_coll + 1); if (i_n2 >= n) return -1; + ggml_tensor * n2 = cgraph->nodes[i_n2]; + const int i_out = next(i_n2 + 1); if (i_out >= n) return -1; + ggml_tensor * out = cgraph->nodes[i_out]; + + if (mm->op != GGML_OP_MUL_MAT || !is_hc(pre, 0) || !is_hc(post, 1) || !is_hc(coll, 6) || + n2->op != GGML_OP_RMS_NORM || out->op != GGML_OP_MUL) { + return -1; + } + const int n_embd = ggml_get_op_params_i32(pre, 1); + const int n_hc = ggml_get_op_params_i32(pre, 2); + const int64_t T = post->ne[1]; + const int64_t ncols_hc = (int64_t) n_embd * n_hc; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + + // R: the streams n1 normalizes, pre mixes and post carries. + const ggml_tensor * R = n1->src[0]; + const ggml_tensor * fn = mm->src[0]; + if (R->type != GGML_TYPE_F32 || R->nb[0] != sizeof(float) || R->ne[0] != ncols_hc || + R->ne[1] != T || R->ne[2] != 1 || R->ne[3] != 1) return -1; + if (n1->ne[1] != T || !ggml_is_contiguous(n1)) return -1; + const ggml_tensor * pre_R = pre->src[2]; + const ggml_tensor * post_R = post->src[0]; + const int64_t R_stride = ggml_cuda_f32_row_stride(R); + for (const ggml_tensor * r : {pre_R, post_R}) { + if (r->data != R->data || r->nb[0] != sizeof(float) || ggml_cuda_f32_row_stride(r) != R_stride) return -1; + } + // mm: the F16 fn product MMVF would run with half accumulation. + if (fn->type != GGML_TYPE_F16 || mm->src[1] != n1 || mm->type != GGML_TYPE_F32 || + fn->ne[0] != ncols_hc || fn->ne[1] != mix_dim || fn->ne[2] != 1 || fn->ne[3] != 1 || + mm->ne[1] != T || T > MMVF_MAX_BATCH_SIZE || !ggml_is_contiguous(mm) || + ggml_backend_buft_is_cuda_split(fn->buffer->buft) || ggml_mul_mat_is_grouped_src(mm)) return -1; + const int cc = ggml_cuda_info().devices[ctx.device].cc; + if (!ggml_cuda_should_use_mmvf(fn->type, cc, fn->ne, fn->nb, T) || + !fast_fp16_available(cc) || ggml_prec(mm->op_params[0]) != GGML_PREC_DEFAULT) return -1; + const int mmv_block = ggml_cuda_mmvf_block_size(ctx.device, ncols_hc, mix_dim, T, /*f16_half_acc=*/true); + if (!ggml_cuda_ds4_hc_boundary_supported(n_hc, n_embd, mmv_block)) return -1; + // pre reads mm (directly or through one reshape). + const ggml_tensor * mix = pre->src[0]; + if (mix != mm) { + if (mix->op != GGML_OP_RESHAPE || mix->src[0] != mm || ggml_cuda_graph_use_count(cgraph, mix) != 1) return -1; + } + const ggml_tensor * base = pre->src[1]; + if (base->type != GGML_TYPE_F32 || !ggml_is_contiguous(base) || base->ne[0] < mix_dim) return -1; + if (pre->ne[0] != n_embd + mix_dim || (T > 1 && pre->ne[1] != T) || !ggml_is_contiguous(pre) || T > 64) return -1; + const int64_t pre_stride = (int64_t) (n_embd + mix_dim); + // post: block_out and pre's split. + const ggml_tensor * block_out = post->src[1]; + const ggml_tensor * split = post->src[2]; + if (block_out->type != GGML_TYPE_F32 || block_out->nb[0] != sizeof(float) || block_out->ne[0] != n_embd) return -1; + if (ggml_get_op_params_i32(post, 1) != n_embd || ggml_get_op_params_i32(post, 2) != n_hc || + !ggml_is_contiguous(post)) return -1; + const char * pre_split = (const char *) pre->data + (size_t) n_embd * sizeof(float); + if (split->data != pre_split || (T > 1 && split->nb[1] != pre->nb[1])) return -1; + // coll: the collapse of post with pre's own pre coefficients. + const ggml_tensor * coll_pre = coll->src[1]; + if (coll->src[0] != post || coll_pre->data != pre_split || coll_pre->ne[0] != n_hc || + (T > 1 && coll_pre->nb[1] != pre->nb[1]) || ggml_get_op_params_i32(coll, 1) != n_embd || + ggml_get_op_params_i32(coll, 2) != n_hc || !ggml_is_contiguous(coll)) return -1; + // n2, out: rms_norm(coll) * w with w one F32 row. + const ggml_tensor * w = out->src[1]; + if (n2->src[0] != coll || out->src[0] != n2 || !ggml_is_contiguous(out) || out->ne[0] != n_embd || + w->type != GGML_TYPE_F32 || !ggml_is_contiguous(w) || w->ne[0] != n_embd || ggml_nrows(w) != 1) return -1; + // Elided intermediates have no reader outside the group. + if (!ggml_cuda_elidable(cgraph, n1, 1) || !ggml_cuda_elidable(cgraph, mm, 1) || + !ggml_cuda_elidable(cgraph, coll, 1) || !ggml_cuda_elidable(cgraph, n2, 1)) return -1; + // pre is written unless its only reader is the split view, read only by + // post and by the view of its pre coefficients that only coll reads. + const bool write_pre = !( + ggml_cuda_graph_use_count(cgraph, pre) == 1 && (pre->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && + split->op == GGML_OP_VIEW && split->src[0] == pre && (split->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && + ggml_cuda_graph_use_count(cgraph, split) == 2 && + coll_pre->op == GGML_OP_VIEW && coll_pre->src[0] == split && + (coll_pre->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && ggml_cuda_graph_use_count(cgraph, coll_pre) == 1); + + float eps_hc = 0.0f, eps_out = 0.0f; + memcpy(&eps_hc, n1->op_params, sizeof(float)); + memcpy(&eps_out, n2->op_params, sizeof(float)); + // out feeds the next sub-block's quantized projections: hand MMVQ its + // q8_1 form, written by the same launch, under MMVQ's memo key. + void * q8_out = nullptr; + if (ggml_cuda_luce_q8_memo_on() && ctx.stream_context().concurrent_events.empty() && + n_embd % MATRIX_ROW_PADDING == 0 && out->buffer != nullptr) { + ggml_backend_cuda_context::luce_q8_memo_entry ent; + ent.src1_node = (const void *) out; + ent.src1_data = out->data; + ent.src0_type = -1; + ent.ne[0] = n_embd; ent.ne[1] = T; ent.ne[2] = 1; ent.ne[3] = 1; + const size_t q8_bytes = (size_t) T * n_embd * sizeof(block_q8_1) / QK8_1; + ent.buf = std::make_unique>(ctx.pool(), q8_bytes); + q8_out = ent.buf->ptr; + ctx.luce_q8_memo.push_back(std::move(ent)); + } + ggml_cuda_ds4_hc_boundary( + ctx, (int) T, n_hc, mmv_block, + (const float *) R->data, R_stride, + (const half *) fn->data, (int64_t) (fn->nb[1] / sizeof(half)), + (const float *) base->data, + (const float *) block_out->data, ggml_cuda_f32_row_stride(block_out), + (const float *) w->data, + (float *) pre->data, pre_stride, + (float *) post->data, ggml_cuda_f32_row_stride(post), + (float *) out->data, ggml_cuda_f32_row_stride(out), + n_embd, eps_hc, eps_out, ggml_get_op_params_i32(pre, 3), + ggml_get_op_params_f32(pre, 4), ggml_get_op_params_f32(pre, 5), ggml_get_op_params_f32(pre, 6), + write_pre, q8_out); + return i_out; +} + static thread_local size_t g_copy_batch_run_count = 0; size_t ggml_backend_cuda_get_copy_batch_run_count(void) { @@ -5564,6 +6063,51 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud // start of fusion operations static bool disable_fusion = (getenv("GGML_CUDA_DISABLE_FUSION") != nullptr); if (!disable_fusion) { + if (node->op == GGML_OP_GLU && ggml_get_glu_op(node) == GGML_GLU_OP_SWIGLU_DS4 && + node->src[1] != nullptr && node->type == GGML_TYPE_F32) { + // A down projection reads it next: emit its q8_1 form too. + void * q8_out = nullptr; + int ncols_q8 = 0; + ggml_cuda_q8_memo_for_producer(*cuda_ctx, cgraph, i, &q8_out, &ncols_q8); + if (q8_out) { + ggml_cuda_op_swiglu_ds4(*cuda_ctx, node, q8_out, ncols_q8); + continue; + } + } + if (node->op == GGML_OP_ROPE && ggml_cuda_try_rope_q8(*cuda_ctx, cgraph, i)) { + continue; + } + if (node->op == GGML_OP_CONT) { + const int last = ggml_cuda_try_cont_into_mmvq(*cuda_ctx, cgraph, i); + if (last > i) { + i = last; + continue; + } + } + if (node->op == GGML_OP_DS4_HC && ggml_get_op_params_i32(node, 0) == 4) { + // The router's select then weights over the same logits. + int j = i + 1; + while (j < cgraph->n_nodes && ggml_cuda_node_is_noop(cgraph->nodes[j])) ++j; + ggml_tensor * wnode = j < cgraph->n_nodes ? cgraph->nodes[j] : nullptr; + if (wnode && wnode->op == GGML_OP_DS4_HC && ggml_get_op_params_i32(wnode, 0) == 5 && + wnode->src[0] == node->src[0] && wnode->src[1] == node && + ggml_is_contiguous(node) && ggml_is_contiguous(wnode) && + node->ne[0] == wnode->ne[0] && node->ne[0] <= 32) { + ggml_cuda_op_ds4_router_fused(*cuda_ctx, node, wnode); + i = j; + continue; + } + } + if (node->op == GGML_OP_RMS_NORM) { + int last = ggml_cuda_try_ds4_hc_boundary(*cuda_ctx, cgraph, i); + if (last < 0) last = ggml_cuda_try_rms_norm_rope_f16(*cuda_ctx, cgraph, i); + if (last < 0) last = ggml_cuda_try_ds4_rms_norm_mix(*cuda_ctx, cgraph, i); + if (last > i) { + i = last; + continue; + } + } + const bool qwen4exp_graph = ggml_cuda_qwen4exp_graph(); const bool qwen4exp_rdna35 = ggml_cuda_qwen4exp_enabled(); if (qwen4exp_graph) { @@ -5779,6 +6323,37 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud bool fused_mul_mat_vec = false; int fused_node_count = 0; + // An RDNA4 Q8_0 gate/up pair whose weights are views (the cluster + // shared expert's row slices): the views launch nothing, so the + // pair and its GLU are the next three nodes that launch something. + if (node->op == GGML_OP_MUL_MAT) { + int j = i + 1; + while (j < cgraph->n_nodes && ggml_cuda_node_is_noop(cgraph->nodes[j])) ++j; + int k = j + 1; + while (k < cgraph->n_nodes && ggml_cuda_node_is_noop(cgraph->nodes[k])) ++k; + if (k < cgraph->n_nodes && j > i + 1 && cgraph->nodes[j]->op == GGML_OP_MUL_MAT && + cgraph->nodes[k]->op == GGML_OP_GLU && + cgraph->nodes[k]->src[0] == node && cgraph->nodes[k]->src[1] == cgraph->nodes[j]) { + ggml_tensor * up_n = cgraph->nodes[j]; + ggml_tensor * glu_n = cgraph->nodes[k]; + const int idxs[3] = { i, j, k }; + const enum ggml_op ops3[3] = { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }; + const int outs[1] = { k }; + if (ggml_can_fuse_subgraph_ext(cgraph, idxs, 3, ops3, outs, 1) && + up_n->src[1]->ne[1] <= ggml_cuda_mmvq_max_ncols() && + !ggml_backend_buft_is_cuda_split(up_n->src[0]->buffer->buft) && + !ggml_backend_buft_is_cuda_split(node->src[0]->buffer->buft) && + ggml_cuda_mmvq_rdna4_glu_pair(node, up_n, glu_n)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = node->src[0]; + ggml_cuda_set_fusion_glu_params(fusion_data, glu_n); + ggml_cuda_mul_mat_vec_q(*cuda_ctx, up_n->src[0], up_n->src[1], nullptr, glu_n, &fusion_data); + i = k; + continue; + } + } + } + for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; @@ -5994,7 +6569,10 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL}, {})) { - ggml_cuda_op_rms_norm_fused(*cuda_ctx, node, cgraph->nodes[i+1]); + void * q8_out = nullptr; + int ncols_q8 = 0; + ggml_cuda_q8_memo_for_producer(*cuda_ctx, cgraph, i + 1, &q8_out, &ncols_q8); + ggml_cuda_op_rms_norm_fused(*cuda_ctx, node, cgraph->nodes[i+1], q8_out, ncols_q8); i++; continue; } @@ -6412,6 +6990,17 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + // The key is the first node's address, and a rebuilt graph (a new + // scheduler uid) can land in the same metadata arena: that is a new graph + // at an old key, not a graph whose properties keep flipping, so its churn + // history (and an eager verdict from it) does not carry over. + const bool new_generation = cgraph->uid != 0 && graph->uid != 0 && cgraph->uid != graph->uid; + if (new_generation && graph->disable_due_to_churn) { + graph->disable_due_to_churn = false; + graph->churn = 0; + graph->replay_streak = 0; + graph->warmup_complete = false; + } if (!ggml_cuda_graphs_disabled_override && !ggml_cuda_graphs_disabled_for_device(cuda_ctx->device) && graph->is_enabled()) { @@ -6453,7 +7042,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, // Properties changed - reset warmup, execute directly until stable again graph->warmup_complete = false; graph->replay_streak = 0; - if (++graph->churn >= ggml_cuda_graph::kMaxChurn) { + if (!new_generation && ++graph->churn >= ggml_cuda_graph::kMaxChurn) { graph->disable_due_to_churn = true; // This key evaluates eagerly from now on: release the // executable and captured graph it would otherwise @@ -7306,6 +7895,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return false; } } + if (b->type == GGML_TYPE_Q8_1) { + // Pre-quantized activation rows: MMVQ only (see the F32 -> Q8_1 copy). + return op->op == GGML_OP_MUL_MAT_ID && ggml_is_quantized(a->type) && + ggml_is_contiguous(b) && b->ne[0] % MATRIX_ROW_PADDING == 0 && + op->ne[2] <= MMVQ_MAX_MOE_BATCH_SIZE; + } if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) { return false; } @@ -7417,6 +8012,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g { ggml_type src0_type = op->src[0]->type; ggml_type src1_type = op->src[1]->type; + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q8_1) { + return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && + op->src[0]->ne[0] % MATRIX_ROW_PADDING == 0; + } if ((src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_BF16 || src0_type == GGML_TYPE_F16) && (src1_type == GGML_TYPE_F32 || src1_type == GGML_TYPE_BF16 || src1_type == GGML_TYPE_F16) ) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mix-wmma-moe.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mix-wmma-moe.cu index fdd3947ab..4b08b101c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mix-wmma-moe.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mix-wmma-moe.cu @@ -1,13 +1,17 @@ // F16 WMMA routed-expert GEMM for the ROCmFP2/ROCmFP3 MIX expert types -// (qtype 106 / 105) on RDNA3 (gfx11), used for large mul_mat_id batches +// (qtype 106 / 105, and plain ROCmFP2 107) on RDNA3 (gfx11) and RDNA4 (gfx12), +// used for large mul_mat_id batches // (prefill). The MMQ path re-decodes every weight tile for each 64-route tile // and runs the codebook unpack inside its int8 inner loop; here each K step // decodes a 128x32 weight tile once into LDS as F16 (scale * codebook level, // the same float product the dequantize kernels round to half) and eight // waves accumulate it against 64 gathered F16 routes on the matrix cores. // -// On by default on RDNA3.5 (gfx115x, the arch it is validated on; -// LUCE_MIX_WMMA_PREFILL=0 disables); other GPUs keep MMQ. +// On by default on RDNA3.5 (gfx115x) and RDNA4 (gfx120x), the archs it is +// validated on (LUCE_MIX_WMMA_PREFILL=0 disables); other GPUs keep MMQ. The +// two archs share the kernel and differ only in the WMMA fragment layout: +// gfx11 lanes hold all 16 K values of their row and C rows 2*e + hi, gfx12 +// lanes hold K values 8*hi .. 8*hi+7 and C rows 8*hi + e. // // The structure (dequantize a weight tile into padded LDS, then accumulate it // on the WMMA units against gathered routed activations) follows Piotr @@ -27,7 +31,7 @@ #include #include -#if defined(GGML_USE_HIP) && defined(__HIP_DEVICE_COMPILE__) && defined(RDNA3) +#if defined(GGML_USE_HIP) && defined(__HIP_DEVICE_COMPILE__) && (defined(RDNA3) || defined(RDNA4)) #define MIX_WMMA_DEVICE 1 #else #define MIX_WMMA_DEVICE 0 @@ -113,15 +117,31 @@ __global__ void mix_wmma_gather(const float * __restrict__ src1, const int32_t * } #if MIX_WMMA_DEVICE -using v16h = __attribute__((__vector_size__(16 * sizeof(_Float16)))) _Float16; using v8f = __attribute__((__vector_size__(8 * sizeof(float)))) float; +#if defined(RDNA4) +using mix_wmma_frag_t = __attribute__((__vector_size__(8 * sizeof(_Float16)))) _Float16; -__device__ __forceinline__ v16h mix_wmma_frag(const half * p) { - union { v16h f; uint4 u[2]; } cvt; +// gfx12: the lane's 8 K values (8 * hi ..) of the row at p (the K step). +__device__ __forceinline__ mix_wmma_frag_t mix_wmma_frag(const half * p, int hi) { + union { mix_wmma_frag_t f; uint4 u; } cvt; + cvt.u = *reinterpret_cast(p + 8 * hi); + return cvt.f; +} +#define MIX_WMMA_MMA __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12 +#define MIX_WMMA_C_ROW(e, hi) (8 * (hi) + (e)) +#else +using mix_wmma_frag_t = __attribute__((__vector_size__(16 * sizeof(_Float16)))) _Float16; + +// gfx11: all 16 K values of the row at p (the K step). +__device__ __forceinline__ mix_wmma_frag_t mix_wmma_frag(const half * p, int /*hi*/) { + union { mix_wmma_frag_t f; uint4 u[2]; } cvt; cvt.u[0] = *reinterpret_cast(p); cvt.u[1] = *reinterpret_cast(p + 8); return cvt.f; } +#define MIX_WMMA_MMA __builtin_amdgcn_wmma_f32_16x16x16_f16_w32 +#define MIX_WMMA_C_ROW(e, hi) (2 * (e) + (hi)) +#endif #endif #if MIX_WMMA_DEVICE @@ -328,18 +348,18 @@ __launch_bounds__(kThreads) __global__ void mix_wmma_moe_kernel( } #pragma unroll for (int ks = 0; ks < kBK2; ks += 16) { - v16h a[2], b[kNF]; + mix_wmma_frag_t a[2], b[kNF]; #pragma unroll - for (int i = 0; i < 2; ++i) a[i] = mix_wmma_frag(&s_w[(wm + i * 16 + sub) * kLds2 + ks]); + for (int i = 0; i < 2; ++i) a[i] = mix_wmma_frag(&s_w[(wm + i * 16 + sub) * kLds2 + ks], hl); #pragma unroll for (int j = 0; j < kNF; ++j) { - if (frag_live[j]) b[j] = mix_wmma_frag(&s_x[(wn + j * 16 + sub) * kLds2 + ks]); + if (frag_live[j]) b[j] = mix_wmma_frag(&s_x[(wn + j * 16 + sub) * kLds2 + ks], hl); } #pragma unroll for (int i = 0; i < 2; ++i) #pragma unroll for (int j = 0; j < kNF; ++j) - if (frag_live[j]) acc[i][j] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a[i], b[j], acc[i][j]); + if (frag_live[j]) acc[i][j] = MIX_WMMA_MMA(a[i], b[j], acc[i][j]); } __syncthreads(); } @@ -358,14 +378,14 @@ __launch_bounds__(kThreads) __global__ void mix_wmma_moe_kernel( for (int i = 0; i < 2; ++i) #pragma unroll for (int v = 0; v < 8; ++v) { - const int c = wm + i * 16 + 2 * v + hl; + const int c = wm + i * 16 + MIX_WMMA_C_ROW(v, hl); out_row[c] = ggml_cuda_op_swiglu_ds4_single(gate_row[c], acc[i][j][v], glu_limit); } } else { #pragma unroll for (int i = 0; i < 2; ++i) #pragma unroll - for (int v = 0; v < 8; ++v) out_row[wm + i * 16 + 2 * v + hl] = acc[i][j][v]; + for (int v = 0; v < 8; ++v) out_row[wm + i * 16 + MIX_WMMA_C_ROW(v, hl)] = acc[i][j][v]; } } #else @@ -401,8 +421,9 @@ static bool mix_wmma_prefill_enabled() { static std::atomic g_mix_wmma_launches{0}; bool ggml_cuda_mix_wmma_moe_available(int device) { - return mix_wmma_prefill_enabled() && device >= 0 && device < ggml_cuda_info().device_count && - GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[device].cc); + if (!mix_wmma_prefill_enabled() || device < 0 || device >= ggml_cuda_info().device_count) return false; + const int cc = ggml_cuda_info().devices[device].cc; + return GGML_CUDA_CC_IS_RDNA3_5(cc) || GGML_CUDA_CC_IS_RDNA4(cc); } uint64_t ggml_cuda_mix_wmma_moe_launch_count() { @@ -420,7 +441,7 @@ bool ggml_cuda_mix_wmma_moe_enabled(const ggml_tensor * src0, const ggml_tensor const long n = v && *v ? std::strtol(v, &end, 10) : 0; return (end && *end == '\0' && n >= 1 && n <= 1 << 20) ? (int) n : 64; }(); - if (!enabled || !GGML_CUDA_CC_IS_RDNA3_5(cc)) return false; + if (!enabled || !(GGML_CUDA_CC_IS_RDNA3_5(cc) || GGML_CUDA_CC_IS_RDNA4(cc))) return false; // Plain ROCmFP2 (qtype 107) is opt-in: LUCE_ROCMFP2_WMMA_PREFILL=1. static const bool plain_fp2 = [] { #ifdef ROCMFP2_AFFINE diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu index f6c276097..c06d104dc 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu @@ -5,6 +5,7 @@ #include "mmvq.cuh" #include "convert.cuh" +#include #include // On HIP the K loops below are unrolled (the "K-loop unroll" pragmas) so the @@ -432,25 +433,9 @@ static void mul_mat_vec_f_switch_fusion( } -template -void launch_mul_mat_vec_f_cuda( - const T * x, const float * y, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, - const int64_t ncols, const int64_t nrows, - const int64_t stride_row, const int64_t stride_col_y, const int64_t stride_col_dst, - const int64_t nchannels_x, const int64_t nchannels_y, const int64_t nchannels_dst, - const int64_t stride_channel_x, const int64_t stride_channel_y, const int64_t stride_channel_dst, const int64_t nsamples_x, - const int64_t nsamples_dst, const int64_t stride_sample_x, const int64_t stride_sample_y, const int64_t stride_sample_dst, - const int64_t nsamples_or_ntokens, const int64_t ids_stride, cudaStream_t stream) { - GGML_ASSERT(ncols % 2 == 0); - GGML_ASSERT(stride_row % 2 == 0); - GGML_ASSERT(stride_col_y % 2 == 0); - GGML_ASSERT(ids || nchannels_dst % nchannels_x == 0); - GGML_ASSERT( nsamples_dst % nsamples_x == 0); - const uint3 nchannels_y_fd = ids ? init_fastdiv_values(nchannels_y) : make_uint3(0, 0, 0); - const uint3 channel_ratio_fd = ids ? make_uint3(0, 0, 0) : init_fastdiv_values(nchannels_dst / nchannels_x); - const uint3 sample_ratio_fd = init_fastdiv_values(nsamples_dst / nsamples_x); - - const int device = ggml_cuda_get_device(); +// The block size the MMVF launch picks for a shape. It sets the partial-sum +// order, so a fused kernel that repeats MMVF's arithmetic must use it too. +int ggml_cuda_mmvf_block_size(int device, int64_t ncols, int64_t nrows, int64_t ncols_dst, bool f16_half_acc) { const int warp_size = ggml_cuda_info().devices[device].warp_size; int64_t block_size_best = warp_size; @@ -466,8 +451,7 @@ void launch_mul_mat_vec_f_cuda( block_size_best = block_size; } } - if constexpr (std::is_same_v && std::is_same_v && - ncols_dst == 5) { + if (f16_half_acc && ncols_dst == 5) { // RDNA 3.5 has enough independent rows at the drafter's width-5 F16 // projections to favor more, smaller Wave32 workgroups. This is a // measured shape table (gfx1151, K=4096: 256 rows -> 128 threads, @@ -483,6 +467,33 @@ void launch_mul_mat_vec_f_cuda( } } } + return (int) block_size_best; +} + +template +void launch_mul_mat_vec_f_cuda( + const T * x, const float * y, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, + const int64_t ncols, const int64_t nrows, + const int64_t stride_row, const int64_t stride_col_y, const int64_t stride_col_dst, + const int64_t nchannels_x, const int64_t nchannels_y, const int64_t nchannels_dst, + const int64_t stride_channel_x, const int64_t stride_channel_y, const int64_t stride_channel_dst, const int64_t nsamples_x, + const int64_t nsamples_dst, const int64_t stride_sample_x, const int64_t stride_sample_y, const int64_t stride_sample_dst, + const int64_t nsamples_or_ntokens, const int64_t ids_stride, cudaStream_t stream) { + GGML_ASSERT(ncols % 2 == 0); + GGML_ASSERT(stride_row % 2 == 0); + GGML_ASSERT(stride_col_y % 2 == 0); + GGML_ASSERT(ids || nchannels_dst % nchannels_x == 0); + GGML_ASSERT( nsamples_dst % nsamples_x == 0); + const uint3 nchannels_y_fd = ids ? init_fastdiv_values(nchannels_y) : make_uint3(0, 0, 0); + const uint3 channel_ratio_fd = ids ? make_uint3(0, 0, 0) : init_fastdiv_values(nchannels_dst / nchannels_x); + const uint3 sample_ratio_fd = init_fastdiv_values(nsamples_dst / nsamples_x); + + const int device = ggml_cuda_get_device(); + const int warp_size = ggml_cuda_info().devices[device].warp_size; + + const int64_t block_size_best = ggml_cuda_mmvf_block_size( + device, ncols, nrows, ncols_dst, + std::is_same_v && std::is_same_v); const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; @@ -666,6 +677,101 @@ static void mul_mat_vec_f_cuda( stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); } +// Skinny-K single-column F16 matvec (K = 64/128/256 over many rows, such as +// the DSpark Markov head). The generic kernel gives every row its own block, +// so a short row is a single trip and the launch is bound by block +// scheduling. Here every lane loads 16 bytes, K/8 lanes share a row and one +// wave covers several rows. Sums accumulate in F32, so values can differ from +// the generic F16-accumulating path in the last bits; it runs on RDNA4 only, +// the GPU it was tuned on, and every other device keeps the generic kernel. +template +static __global__ void __launch_bounds__(256) mul_mat_vec_f16_skinny( + const half * __restrict__ x, const float * __restrict__ y, const float * __restrict__ x_bias, + float * __restrict__ dst, const int nrows, const int64_t stride_row) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int lanes_per_row = K / 8; + static_assert(lanes_per_row <= warp_size, "one row must fit in a wave"); + constexpr int rows_per_warp = warp_size / lanes_per_row; + + const int lane = threadIdx.x % warp_size; + const int warp = threadIdx.x / warp_size; + const int sub = lane % lanes_per_row; + const int rsub = lane / lanes_per_row; + const int rows_per_block = (blockDim.x / warp_size) * rows_per_warp; + + float yr[8]; +#pragma unroll + for (int j = 0; j < 8; ++j) { + yr[j] = y[sub*8 + j]; + } + + // row0 is uniform across the block, so every lane runs the same number of + // trips and the shuffle reduction below is always fully populated. + for (int row0 = blockIdx.x*rows_per_block; row0 < nrows; row0 += gridDim.x*rows_per_block) { + const int row = row0 + warp*rows_per_warp + rsub; + float sum = 0.0f; + if (row < nrows) { + const uint4 q = *((const uint4 *) (x + row*stride_row) + sub); + const half2 * h = (const half2 *) &q; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 f = __half22float2(h[j]); + sum += f.x*yr[2*j] + f.y*yr[2*j + 1]; + } + } + sum = warp_reduce_sum(sum); + if (row < nrows && sub == 0) { + dst[row] = x_bias ? sum + x_bias[row] : sum; + } + } +} + +// Returns true when the skinny kernel handled the product. +static bool mul_mat_vec_f16_skinny_try(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, + ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion, cudaStream_t stream) { +#if defined(GGML_USE_HIP) + const int device = ggml_cuda_get_device(); + if (src0->type != GGML_TYPE_F16 || ids || !GGML_CUDA_CC_IS_RDNA4(ggml_cuda_info().devices[device].cc)) { + return false; + } + if (fusion && (fusion->gate || fusion->gate_bias)) { + return false; + } + const int64_t K = src0->ne[0]; + const int64_t nrows = src0->ne[1]; + if ((K != 64 && K != 128 && K != 256) || nrows < 2048 || nrows > INT32_MAX) { + return false; + } + if (src1->ne[1] != 1 || src1->ne[2] != 1 || src1->ne[3] != 1 || src0->ne[2] != 1 || src0->ne[3] != 1 || + dst->ne[1] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) { + return false; + } + const int64_t stride_row = src0->nb[1] / ggml_type_size(src0->type); + if (src0->nb[1] % 16 != 0 || ((uintptr_t) src0->data) % 16 != 0) { + return false; + } + const float * x_bias = fusion && fusion->x_bias ? (const float *) fusion->x_bias->data : nullptr; + + const int warp_size = ggml_cuda_info().devices[device].warp_size; + const int64_t rows_per_block = (256 / warp_size) * (warp_size / (K / 8)); + const int64_t nblocks = std::min((nrows + rows_per_block - 1) / rows_per_block, 4096); + const half * x = (const half *) src0->data; + const float * y = (const float *) src1->data; + float * d = (float *) dst->data; + switch (K) { + case 64: mul_mat_vec_f16_skinny< 64><<>>(x, y, x_bias, d, (int) nrows, stride_row); break; + case 128: mul_mat_vec_f16_skinny<128><<>>(x, y, x_bias, d, (int) nrows, stride_row); break; + case 256: mul_mat_vec_f16_skinny<256><<>>(x, y, x_bias, d, (int) nrows, stride_row); break; + default: return false; + } + CUDA_CHECK(cudaGetLastError()); + return true; +#else + GGML_UNUSED_VARS(src0, src1, ids, dst, fusion, stream); + return false; +#endif // defined(GGML_USE_HIP) +} + void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion) { GGML_ASSERT( src1->type == GGML_TYPE_F32); @@ -686,6 +792,10 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); GGML_ASSERT( nb0 == ts_dst); + if (mul_mat_vec_f16_skinny_try(src0, src1, ids, dst, fusion, ctx.stream())) { + return; + } + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const enum ggml_prec prec = fast_fp16_available(cc) ? ggml_prec(dst->op_params[0]) : GGML_PREC_F32; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cuh index a50f7c021..9f68ce4d9 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cuh @@ -12,3 +12,6 @@ void ggml_cuda_op_mul_mat_vec_f( const int64_t src1_padded_row_size, cudaStream_t stream); bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0_ne, const size_t * src0_nb, int64_t ne11); + +// The block size MMVF launches with for this shape (see mmvf.cu). +int ggml_cuda_mmvf_block_size(int device, int64_t ncols, int64_t nrows, int64_t ncols_dst, bool f16_half_acc); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu index beb82d810..d23cb4075 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -624,6 +624,9 @@ static bool mmid_grouped_type_ok(ggml_type type) { } } +// The memo entry kind of a cached expert grouping (q8_1 entries use a type or -1). +#define MMID_PREP_MEMO_KIND (-2) + static bool mmid_grouped_device_ok() { static const int selected = []() { const char * e = std::getenv("LUCE_MMID_GROUPED_DEVICE"); @@ -1862,7 +1865,8 @@ static __global__ void mmid_group_prep( // (n_expert_used x n_tokens) expert-matrix reads. template + int tokens_per_group = MMID_GROUPED_DEFAULT_TPG, + bool fp2_prefetch = false> __launch_bounds__(tokens_per_group*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q_moe_grouped( const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ meta, @@ -1938,39 +1942,46 @@ static __global__ void mul_mat_vec_q_moe_grouped( float tmp[c_rows_per_block] = {0.0f}; float tmp_gate[c_rows_per_block] = {0.0f}; - for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { - const int kby = kbx * (qk/QK8_1); - const int kqs = vdr * (threadIdx.x % (qi/vdr)); + // ROCmFP2 on gfx1151: the per-expert kernel's prefetching loop, the same + // terms in the same kbx order as the loop below (bit-identical). + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2 && !fp2_packed32 && fp2_prefetch) { + mul_mat_vec_q_moe_fp2_prefetch( + vx, vgate, use_gate, y, kbx_offset, stride_row_x, blocks_per_row_x, tmp, tmp_gate); + } else { + for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (threadIdx.x % (qi/vdr)); #pragma unroll - for (int i = 0; i < c_rows_per_block; ++i) { - if constexpr (fp2_packed32) { - tmp[i] += vec_dot_rocmfpx_fp2_q8_1_packed32( - vx, &y[kby], - kbx_offset + i*stride_row_x + kbx, kqs); - } else if constexpr (fp3_packed24) { - tmp[i] += vec_dot_rocmfpx_fp3_q8_1_packed24( - vx, &y[kby], - kbx_offset + i*stride_row_x + kbx, kqs); - } else { - tmp[i] += vec_dot_q_cuda( - vx, &y[kby], - kbx_offset + i*stride_row_x + kbx, kqs); - } - if constexpr (has_fusion) { - if (use_gate) { - if constexpr (fp2_packed32) { - tmp_gate[i] += vec_dot_rocmfpx_fp2_q8_1_packed32( - vgate, &y[kby], - kbx_offset + i*stride_row_x + kbx, kqs); - } else if constexpr (fp3_packed24) { - tmp_gate[i] += vec_dot_rocmfpx_fp3_q8_1_packed24( - vgate, &y[kby], - kbx_offset + i*stride_row_x + kbx, kqs); - } else { - tmp_gate[i] += vec_dot_q_cuda( - vgate, &y[kby], - kbx_offset + i*stride_row_x + kbx, kqs); + for (int i = 0; i < c_rows_per_block; ++i) { + if constexpr (fp2_packed32) { + tmp[i] += vec_dot_rocmfpx_fp2_q8_1_packed32( + vx, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else if constexpr (fp3_packed24) { + tmp[i] += vec_dot_rocmfpx_fp3_q8_1_packed24( + vx, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else { + tmp[i] += vec_dot_q_cuda( + vx, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } + if constexpr (has_fusion) { + if (use_gate) { + if constexpr (fp2_packed32) { + tmp_gate[i] += vec_dot_rocmfpx_fp2_q8_1_packed32( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else if constexpr (fp3_packed24) { + tmp_gate[i] += vec_dot_rocmfpx_fp3_q8_1_packed24( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else { + tmp_gate[i] += vec_dot_q_cuda( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } } } } @@ -2037,7 +2048,8 @@ static __global__ void mul_mat_vec_q_moe_grouped( template + int tokens_per_group = MMID_GROUPED_DEFAULT_TPG, + bool fp2_prefetch = false> static void mul_mat_vec_q_moe_grouped_launch( const void * vx, const void * vy, const int32_t * meta, const ggml_cuda_mm_fusion_args_device & fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, @@ -2055,7 +2067,7 @@ static void mul_mat_vec_q_moe_grouped_launch( if (has_fusion) { mul_mat_vec_q_moe_grouped< type, rows_per_block, true, fp3_packed24, fp2_packed32, - tokens_per_group> + tokens_per_group, fp2_prefetch> <<>>( vx, vy, meta, fusion, dst, (uint32_t) np, ncols_x, nchannels_y, nrows_x, stride_row_x, stride_col_y, @@ -2064,7 +2076,7 @@ static void mul_mat_vec_q_moe_grouped_launch( } else { mul_mat_vec_q_moe_grouped< type, rows_per_block, false, fp3_packed24, fp2_packed32, - tokens_per_group> + tokens_per_group, fp2_prefetch> <<>>( vx, vy, meta, fusion, dst, (uint32_t) np, ncols_x, nchannels_y, nrows_x, stride_row_x, stride_col_y, @@ -2073,6 +2085,18 @@ static void mul_mat_vec_q_moe_grouped_launch( } } +// The prefetching ROCmFP2 loop (mul_mat_vec_q_moe_fp2_prefetch), for the +// per-expert and the grouped expert kernels: LUCE_CUDA_MMVQ_MOE_FP2_PREFETCH=0/1 +// wins; unset, gfx1151 only. +static bool mmvq_moe_fp2_prefetch_on(const int cc) { + static const int setting = []() { + const char * e = std::getenv("LUCE_CUDA_MMVQ_MOE_FP2_PREFETCH"); + if (e == nullptr) return -1; + return (e[0] == '1' && e[1] == '\0') ? 1 : 0; + }(); + return setting >= 0 ? setting == 1 : is_gfx1151(cc); +} + static bool mul_mat_vec_q_grouped_dispatch( const ggml_type type, const void * vx, const void * vy, const int32_t * meta, const ggml_cuda_mm_fusion_args_device & fusion, float * dst, @@ -2124,6 +2148,13 @@ static bool mul_mat_vec_q_grouped_dispatch( nrows_x, stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst, max_groups, warp_size, stream); + } else if (mmvq_moe_fp2_prefetch_on(ggml_cuda_info().devices[ggml_cuda_get_device()].cc)) { + mul_mat_vec_q_moe_grouped_launch< + GGML_TYPE_Q2_0_ROCMFP2, false, false, MMID_GROUPED_DEFAULT_TPG, true>( + vx, vy, meta, fusion, dst, ncols_x, nchannels_y_fd, + nrows_x, stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + max_groups, warp_size, stream); } else { mul_mat_vec_q_moe_grouped_launch< GGML_TYPE_Q2_0_ROCMFP2>( @@ -2441,16 +2472,7 @@ static void mul_mat_vec_q_moe_launch( std::getenv("LUCE_CUDA_MMVQ_MOE_FP2_PACKED32"); return e && e[0] == '1' && e[1] == '\0'; }(); - // Explicit LUCE_CUDA_MMVQ_MOE_FP2_PREFETCH wins; unset defaults to the - // prefetching ROCmFP2 loop on gfx1151 only. - static const int fp2_prefetch_setting = []() { - const char * e = - std::getenv("LUCE_CUDA_MMVQ_MOE_FP2_PREFETCH"); - if (e == nullptr) return -1; - return (e[0] == '1' && e[1] == '\0') ? 1 : 0; - }(); - const bool fp2_prefetch = - fp2_prefetch_setting >= 0 ? fp2_prefetch_setting == 1 : gfx1151; + const bool fp2_prefetch = mmvq_moe_fp2_prefetch_on(ggml_cuda_info().devices[ggml_cuda_get_device()].cc); GGML_UNUSED(fp2_prefetch); #define GGML_MOE_LAUNCH_RPB(RPB) \ @@ -2590,14 +2612,18 @@ static void mul_mat_vec_q_moe_launch( // kernel at every width, so decode and verify stay consistent. #define MMVQ_Q8_RDNA4_MAX_COLS 8 -template +// glu: vx is the up matrix, vx_gate the gate matrix of the same shape; each +// product keeps the single kernel's partials and reduction, and dst gets +// ggml_cuda_op_swiglu_ds4_single(gate, up, glu_limit), the GLU kernel's value. +template __launch_bounds__(32*c_vw*c_row_groups, 1) static __global__ void mul_mat_vec_q8_0_rdna4( const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_y, const int stride_col_dst, const uint3 channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const uint3 sample_ratio, - const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst) { + const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, + const int kb_off, const void * __restrict__ vx_gate, const float glu_limit) { constexpr int nvirt = 8; // waves of the reference block constexpr int nq = nvirt / c_vw; // virtual waves per physical wave static_assert(nvirt % c_vw == 0, "c_vw must divide eight"); @@ -2619,28 +2645,41 @@ static __global__ void mul_mat_vec_q8_0_rdna4( const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_dst*stride_sample_y + channel_dst*stride_channel_y; const block_q8_0 * x = ((const block_q8_0 *) vx) + sample_x*stride_sample_x + channel_x*stride_channel_x; + const block_q8_0 * xg = glu ? ((const block_q8_0 *) vx_gate) + sample_x*stride_sample_x + channel_x*stride_channel_x : x; + const block_q8_0 * xr[c_rows]; + const block_q8_0 * xgr[c_rows]; bool row_ok[c_rows]; #pragma unroll for (int r = 0; r < c_rows; ++r) { const int row = row_base + r; row_ok[r] = row < nrows_x; - xr[r] = x + (int64_t) (row_ok[r] ? row : 0)*stride_row_x; + xr[r] = x + (int64_t) (row_ok[r] ? row : 0)*stride_row_x; + xgr[r] = xg + (int64_t) (row_ok[r] ? row : 0)*stride_row_x; } float part[c_rows][ncols_dst][nq]; + float partg[glu ? c_rows : 1][ncols_dst][nq]; #pragma unroll for (int r = 0; r < c_rows; ++r) #pragma unroll for (int j = 0; j < ncols_dst; ++j) #pragma unroll - for (int q = 0; q < nq; ++q) part[r][j][q] = 0.0f; + for (int q = 0; q < nq; ++q) { + part[r][j][q] = 0.0f; + if constexpr (glu) partg[r][j][q] = 0.0f; + } - for (int k0 = 0; k0 < nb; k0 += 64) { + // kb_off > 0: the rows are a column range starting at block kb_off of + // wider rows (ggml_backend_cuda_mul_mat_whole_row_lanes). Each block takes + // the lane and K step it takes in the whole rows, so every partial adds + // the same terms in the same order as the whole rows times a src1 that is + // zero outside the range: those products only add exact zeros. + for (int k0 = (kb_off/64)*64; k0 < kb_off + nb; k0 += 64) { #pragma unroll for (int q = 0; q < nq; ++q) { - const int kbx = k0 + 8*(g + c_vw*q) + boff; - if (kbx < nb) { + const int kbx = k0 + 8*(g + c_vw*q) + boff - kb_off; + if (kbx >= 0 && kbx < nb) { int u0[ncols_dst]; int u1[ncols_dst]; float dy[ncols_dst]; @@ -2664,12 +2703,26 @@ static __global__ void mul_mat_vec_q8_0_rdna4( sumi = ggml_cuda_dp4a(v1, u1[j], sumi); part[r][j][q] += dx*dy[j] * ((float) sumi); } + if constexpr (glu) { + const block_q8_0 * bg = xgr[r] + kbx; + const int g0 = get_int_b2(bg->qs, 2*quarter + 0); + const int g1 = get_int_b2(bg->qs, 2*quarter + 1); + const float dg = bg->d; +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + int sumg = 0; + sumg = ggml_cuda_dp4a(g0, u0[j], sumg); + sumg = ggml_cuda_dp4a(g1, u1[j], sumg); + partg[r][j][q] += dg*dy[j] * ((float) sumg); + } + } } } } } float sum[c_rows][ncols_dst]; + float sumg[glu ? c_rows : 1][ncols_dst]; if constexpr (c_vw == 1) { #pragma unroll for (int r = 0; r < c_rows; ++r) @@ -2681,15 +2734,27 @@ static __global__ void mul_mat_vec_q8_0_rdna4( t += part[r][j][q]; } sum[r][j] = t; + if constexpr (glu) { + float tg = partg[r][j][0]; +#pragma unroll + for (int q = 1; q < nvirt; ++q) { + tg += partg[r][j][q]; + } + sumg[r][j] = tg; + } } } else { __shared__ float sh[c_row_groups][c_vw][c_rows][ncols_dst][nq][32]; + __shared__ float shg[glu ? c_row_groups : 1][glu ? c_vw : 1][glu ? c_rows : 1][ncols_dst][glu ? nq : 1][32]; #pragma unroll for (int r = 0; r < c_rows; ++r) #pragma unroll for (int j = 0; j < ncols_dst; ++j) #pragma unroll - for (int q = 0; q < nq; ++q) sh[rg][g][r][j][q][lane] = part[r][j][q]; + for (int q = 0; q < nq; ++q) { + sh[rg][g][r][j][q][lane] = part[r][j][q]; + if constexpr (glu) shg[rg][g][r][j][q][lane] = partg[r][j][q]; + } __syncthreads(); if (g != 0) { return; @@ -2704,6 +2769,14 @@ static __global__ void mul_mat_vec_q8_0_rdna4( t += sh[rg][v % c_vw][r][j][v / c_vw][lane]; } sum[r][j] = t; + if constexpr (glu) { + float tg = shg[rg][0][r][j][0][lane]; +#pragma unroll + for (int v = 1; v < nvirt; ++v) { + tg += shg[rg][v % c_vw][r][j][v / c_vw][lane]; + } + sumg[r][j] = tg; + } } } @@ -2713,7 +2786,12 @@ static __global__ void mul_mat_vec_q8_0_rdna4( #pragma unroll for (int j = 0; j < ncols_dst; ++j) { const float t = warp_reduce_sum<32>(sum[r][j]); - if (lane == 0 && row_ok[r]) { + if constexpr (glu) { + const float tg = warp_reduce_sum<32>(sumg[r][j]); + if (lane == 0 && row_ok[r]) { + dst[j*stride_col_dst + row_base + r] = ggml_cuda_op_swiglu_ds4_single(tg, t, glu_limit); + } + } else if (lane == 0 && row_ok[r]) { dst[j*stride_col_dst + row_base + r] = t; } } @@ -2735,14 +2813,38 @@ static void mul_mat_vec_q8_0_rdna4_launch_cfg( const int stride_col_dst, const int nchannels_dst, const uint3 channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const int nsamples_dst, const uint3 sample_ratio, const int stride_sample_x, - const int stride_sample_y, const int stride_sample_dst, cudaStream_t stream) { + const int stride_sample_y, const int stride_sample_dst, const int kb_off, cudaStream_t stream) { constexpr int rows_per_block = c_rows*c_row_groups; const dim3 block_nums((nrows_x + rows_per_block - 1)/rows_per_block, nchannels_dst, nsamples_dst); const dim3 block_dims(32, c_vw*c_row_groups, 1); mul_mat_vec_q8_0_rdna4<<>>( vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, - sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst); + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, kb_off, nullptr, 0.0f); +} + +// The gate/up pair with SwiGLU-DS4 (ggml_cuda_mmvq_rdna4_glu_pair): one block +// shape, two rows per wave over four waves, whose cross-wave partials of both +// products fit the LDS at eight columns. Every shape is bit-exact, so the pair +// gives the value of the two table-shaped products and the GLU kernel. +template +static void mul_mat_vec_q8_0_rdna4_launch_glu( + const void * vx, const void * vx_gate, const void * vy, float * dst, const float glu_limit, + const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_y, + const int stride_col_dst, const int nchannels_dst, const uint3 channel_ratio, + const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const int nsamples_dst, const uint3 sample_ratio, const int stride_sample_x, + const int stride_sample_y, const int stride_sample_dst, cudaStream_t stream) { + constexpr int c_rows = 2, c_vw = 4, c_row_groups = 1; + static_assert(2u*c_row_groups*c_vw*c_rows*ncols_dst*(8/c_vw)*32*sizeof(float) <= 65536, + "both products' partials must fit the LDS"); + constexpr int rows_per_block = c_rows*c_row_groups; + const dim3 block_nums((nrows_x + rows_per_block - 1)/rows_per_block, nchannels_dst, nsamples_dst); + const dim3 block_dims(32, c_vw*c_row_groups, 1); + mul_mat_vec_q8_0_rdna4<<>>( + vx, vy, dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, 0, vx_gate, glu_limit); } template @@ -2752,12 +2854,12 @@ static void mul_mat_vec_q8_0_rdna4_launch_nc( const int stride_col_dst, const int nchannels_dst, const uint3 channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const int nsamples_dst, const uint3 sample_ratio, const int stride_sample_x, - const int stride_sample_y, const int stride_sample_dst, cudaStream_t stream) { + const int stride_sample_y, const int stride_sample_dst, const int kb_off, cudaStream_t stream) { #define Q8_RDNA4_LAUNCH(R, VW, G) \ mul_mat_vec_q8_0_rdna4_launch_cfg(vx, vy, dst, ncols_x, nrows_x, \ stride_row_x, stride_col_y, stride_col_dst, nchannels_dst, channel_ratio, \ stride_channel_x, stride_channel_y, stride_channel_dst, nsamples_dst, sample_ratio, \ - stride_sample_x, stride_sample_y, stride_sample_dst, stream) + stride_sample_x, stride_sample_y, stride_sample_dst, kb_off, stream) // Block shapes from the gfx1201 sweep of the DS4.1 dense projections // (server/test/bench/bench_ds41_q8_mmvq.cpp). Few rows (attn_kv, 512) or one // K step per row (attn_q_b, K=1280): four waves share a row pair's eight @@ -2773,17 +2875,41 @@ static void mul_mat_vec_q8_0_rdna4_launch_nc( #undef Q8_RDNA4_LAUNCH } +// vx_gate set: the gate/up pair (vx is the up matrix), 1..8 columns, no column range. +// One column is a single-token decode, fused by the generic gate/GLU fusion: it +// takes the same kernel as a verify's columns, so the two stay batch-invariant. static void mul_mat_vec_q8_0_rdna4_launch( const void * vx, const void * vy, float * dst, const int ncols_dst, const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_y, const int stride_col_dst, const int nchannels_dst, const uint3 channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const int nsamples_dst, const uint3 sample_ratio, const int stride_sample_x, - const int stride_sample_y, const int stride_sample_dst, cudaStream_t stream) { + const int stride_sample_y, const int stride_sample_dst, const int kb_off, + const void * vx_gate, const float glu_limit, cudaStream_t stream) { + if (vx_gate) { + GGML_ASSERT(kb_off == 0); +#define Q8_RDNA4_GLU_NC(NC) case NC: mul_mat_vec_q8_0_rdna4_launch_glu(vx, vx_gate, vy, dst, glu_limit, \ + ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, nchannels_dst, channel_ratio, \ + stride_channel_x, stride_channel_y, stride_channel_dst, nsamples_dst, sample_ratio, \ + stride_sample_x, stride_sample_y, stride_sample_dst, stream); break + switch (ncols_dst) { + Q8_RDNA4_GLU_NC(1); + Q8_RDNA4_GLU_NC(2); + Q8_RDNA4_GLU_NC(3); + Q8_RDNA4_GLU_NC(4); + Q8_RDNA4_GLU_NC(5); + Q8_RDNA4_GLU_NC(6); + Q8_RDNA4_GLU_NC(7); + Q8_RDNA4_GLU_NC(8); + default: GGML_ABORT("unreachable q8_0 rdna4 glu width"); + } +#undef Q8_RDNA4_GLU_NC + return; + } #define Q8_RDNA4_NC(NC) case NC: mul_mat_vec_q8_0_rdna4_launch_nc(vx, vy, dst, ncols_x, nrows_x, \ stride_row_x, stride_col_y, stride_col_dst, nchannels_dst, channel_ratio, stride_channel_x, \ stride_channel_y, stride_channel_dst, nsamples_dst, sample_ratio, stride_sample_x, \ - stride_sample_y, stride_sample_dst, stream); break + stride_sample_y, stride_sample_dst, kb_off, stream); break switch (ncols_dst) { Q8_RDNA4_NC(1); Q8_RDNA4_NC(2); @@ -2799,6 +2925,44 @@ static void mul_mat_vec_q8_0_rdna4_launch( } #endif // defined(GGML_USE_HIP) +// A Q8_0 gate/up pair (two MUL_MATs on the same activations) whose SwiGLU-DS4 +// reads exactly their outputs, at 2..8 columns of batch-invariant products on +// RDNA4: ggml_cuda_mul_mat_vec_q with the gate as fusion runs it as one launch +// of mul_mat_vec_q8_0_rdna4, bit-identical to the two products and the GLU +// kernel (the caller rules out split buffers). +bool ggml_cuda_mmvq_rdna4_glu_pair(const ggml_tensor * gate, const ggml_tensor * up, const ggml_tensor * glu) { +#if defined(GGML_USE_HIP) + if (!mmvq_q8_0_rdna4_enabled() || !ggml_cuda_mmvq_batch_invariant()) return false; + const int device = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[device].cc; + if (!GGML_CUDA_CC_IS_RDNA4(cc) || ggml_cuda_info().devices[device].warp_size != 32 || + get_device_table_id(cc) != MMVQ_PARAMETERS_RDNA4 || + calc_nwarps(GGML_TYPE_Q8_0, 1, MMVQ_PARAMETERS_RDNA4) != 8) { + return false; + } + if (gate->op != GGML_OP_MUL_MAT || up->op != GGML_OP_MUL_MAT) return false; + const ggml_tensor * g0 = gate->src[0]; + const ggml_tensor * u0 = up->src[0]; + const ggml_tensor * x = up->src[1]; + if (g0->type != GGML_TYPE_Q8_0 || u0->type != GGML_TYPE_Q8_0 || gate->src[1] != x || + !ggml_are_same_shape(g0, u0) || !ggml_are_same_stride(g0, u0) || + x->type != GGML_TYPE_F32 || up->type != GGML_TYPE_F32 || gate->type != GGML_TYPE_F32 || + x->ne[2] != 1 || x->ne[3] != 1 || g0->ne[2] != 1 || g0->ne[3] != 1) { + return false; + } + if (x->ne[1] < 2 || x->ne[1] > MMVQ_Q8_RDNA4_MAX_COLS) return false; + if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU_DS4 || + glu->src[0] != gate || glu->src[1] != up || ggml_get_op_params_i32(glu, 1) != 0 || + glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu)) { + return false; + } + return true; +#else + GGML_UNUSED_VARS(gate, up, glu); + return false; +#endif // defined(GGML_USE_HIP) +} + // True when the single-column launch for this type, device and K is the generic kernel with one wave per row: // each lane walks K in the same order and one warp reduction sums the row, which is also what the multi-token // MoE kernel does for every (token, row). Batch-invariant MUL_MAT_ID can then keep the MoE kernel, which serves @@ -2831,7 +2995,8 @@ static void mul_mat_vec_q_switch_ncols_dst( const int nchannels_x, const int nchannels_y, const int nchannels_dst, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, - const int ids_stride, cudaStream_t stream, const bool ids_tokenwise_samples = false) { + const int ids_stride, cudaStream_t stream, const bool ids_tokenwise_samples = false, + const int kb_off = 0) { GGML_ASSERT(ncols_x % ggml_blck_size(type) == 0); // Tokens as samples (see the batch-invariant MUL_MAT_ID case below) are a single-column launch. @@ -2974,7 +3139,12 @@ static void mul_mat_vec_q_switch_ncols_dst( // arithmetic, packed rows, see mul_mat_vec_q8_0_rdna4. #if defined(GGML_USE_HIP) if constexpr (type == GGML_TYPE_Q8_0) { - if (!has_ids && !has_fusion && (ncols_dst == 1 || width_invariant) && + // A gate/up pair with SwiGLU-DS4 and nothing else fused (see + // ggml_cuda_mmvq_rdna4_glu_pair) runs as one launch of this kernel. + const bool rdna4_glu = has_fusion && fusion.gate != nullptr && fusion.x_bias == nullptr && + fusion.gate_bias == nullptr && fusion.glu_op == GGML_GLU_OP_SWIGLU_DS4 && + fusion.gate_value_scale == 1.0f && fusion.x_value_scale == 1.0f && kb_off == 0; + if (!has_ids && (!has_fusion || rdna4_glu) && (ncols_dst == 1 || width_invariant) && ncols_dst <= MMVQ_Q8_RDNA4_MAX_COLS && GGML_CUDA_CC_IS_RDNA4(cc) && warp_size == 32 && table_id == MMVQ_PARAMETERS_RDNA4 && calc_nwarps(type, 1, table_id) == 8 && mmvq_q8_0_rdna4_enabled()) { @@ -2982,11 +3152,15 @@ static void mul_mat_vec_q_switch_ncols_dst( vx, vy, dst, ncols_dst, ncols_x, nrows_x, stride_row_x, stride_col_y, stride_col_dst, nchannels_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, nsamples_dst, sample_ratio_fd, - stride_sample_x, stride_sample_y, stride_sample_dst, stream); + stride_sample_x, stride_sample_y, stride_sample_dst, kb_off, + rdna4_glu ? fusion.gate : nullptr, fusion.glu_param0, stream); return; } } #endif // defined(GGML_USE_HIP) + // The generic kernels fuse one column only; several columns are the RDNA4 + // Q8_0 gate/up pair above (ggml_cuda_mmvq_rdna4_glu_pair). + GGML_ASSERT(!has_fusion || has_ids || ncols_dst == 1); if constexpr (type == GGML_TYPE_Q8_0) { if (width_invariant) { #define GGML_MMVQ_INVARIANT_LAUNCH(NC) \ @@ -3329,7 +3503,7 @@ static void mul_mat_vec_q_switch_type( const int nchannels_x, const int nchannels_y, const int nchannels_dst, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, - const int ids_stride, cudaStream_t stream) { + const int ids_stride, cudaStream_t stream, const int kb_off = 0) { switch (type_x) { case GGML_TYPE_Q2_0: mul_mat_vec_q_switch_ncols_dst @@ -3365,7 +3539,8 @@ static void mul_mat_vec_q_switch_type( mul_mat_vec_q_switch_ncols_dst (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, - nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream, + /*ids_tokenwise_samples=*/false, kb_off); break; case GGML_TYPE_MXFP4: mul_mat_vec_q_switch_ncols_dst @@ -3505,10 +3680,18 @@ static void mul_mat_vec_q_switch_type( } } +bool ggml_cuda_luce_q8_memo_on() { + static const bool on = []() { + const char * e = getenv("LUCE_Q8_MEMO"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + return on; +} + void ggml_cuda_mul_mat_vec_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion) { - GGML_ASSERT( src1->type == GGML_TYPE_F32); + GGML_ASSERT( src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_Q8_1); GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. @@ -3536,7 +3719,11 @@ void ggml_cuda_mul_mat_vec_q( ggml_cuda_mm_fusion_args_device fusion_local{}; if (fusion) { - GGML_ASSERT( ids || dst->ne[1] == 1); + // Several columns only for the RDNA4 Q8_0 gate/up pair (a gate, no + // biases: ggml_cuda_mmvq_rdna4_glu_pair), which switch_ncols_dst runs + // on mul_mat_vec_q8_0_rdna4. + GGML_ASSERT( ids || dst->ne[1] == 1 || + (src0->type == GGML_TYPE_Q8_0 && fusion->gate && !fusion->x_bias && !fusion->gate_bias)); if (fusion->x_bias) { GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32); GGML_ASSERT(fusion->x_bias->ne[0] == dst->ne[0]); @@ -3576,13 +3763,15 @@ void ggml_cuda_mul_mat_vec_q( // matmuls of one token is a straight win on every backend measured, and // requiring an env var to get the tuned path invites running the slow one // by accident. LUCE_Q8_MEMO=0 opts out. - static const bool luce_q8_memo_on = []() { - const char * e = getenv("LUCE_Q8_MEMO"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); + static const bool luce_q8_memo_on = ggml_cuda_luce_q8_memo_on(); const size_t q8_bytes = ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1; ggml_cuda_pool_alloc src1_q8_1(ctx.pool()); char * src1_q8_d = nullptr; + if (src1->type == GGML_TYPE_Q8_1) { + // Already the q8_1 rows read below (quantized on another device). + GGML_ASSERT(ggml_is_contiguous(src1) && ne10 == ne10_padded); + src1_q8_d = (char *) src1->data; + } // The src1->q8_1 quantization depends only on src0->type and src1's dims/strides, // not on ids (ids only affects the matmul kernel's channel/dst strides below), so // the memo is valid for MUL_MAT_ID too: gate/up and the shared expert re-quantize @@ -3595,13 +3784,33 @@ void ggml_cuda_mul_mat_vec_q( // entry filled on one stream could be read on another while the // quantize kernel is still in flight. Skip the memo whenever concurrent // streams are active for this evaluation. - const bool use_q8_memo = luce_q8_memo_on && src1->buffer != nullptr && + const bool use_q8_memo = luce_q8_memo_on && src1->buffer != nullptr && src1->type == GGML_TYPE_F32 && ctx.stream_context().concurrent_events.empty(); + // quantize_row_q8_1_cuda ignores src0's type and writes rows in their + // flattened order, so a contiguous src1 is keyed by the node it reshapes + // (a reshape only reinterprets the same bytes), its address, its row + // length and its row count: the routed experts' [n_embd, 1, T] input and + // the shared expert's [n_embd, T] activations share one quantization. + const ggml_tensor * q8_key_node = src1; + int64_t q8_key_ne[4] = {ne10, ne11, ne12, ne13}; + int q8_key_type = (int) src0->type; + if (ggml_is_contiguous(src1)) { + // src[0], not view_src: view_src is the allocation root and would + // skip an in-place node between the reshape and that root. + while (q8_key_node->op == GGML_OP_RESHAPE && q8_key_node->src[0]) { + q8_key_node = q8_key_node->src[0]; + } + q8_key_ne[1] = ne11*ne12*ne13; + q8_key_ne[2] = 1; + q8_key_ne[3] = 1; + q8_key_type = -1; + } if (use_q8_memo) { for (const auto & e : ctx.luce_q8_memo) { - if (e.src1_node == (const void *) src1 && e.src1_data == (const void *) src1_d && - e.src0_type == (int) src0->type && - e.ne[0] == ne10 && e.ne[1] == ne11 && e.ne[2] == ne12 && e.ne[3] == ne13) { + if (e.src1_node == (const void *) q8_key_node && e.src1_data == (const void *) src1_d && + e.src0_type == q8_key_type && + e.ne[0] == q8_key_ne[0] && e.ne[1] == q8_key_ne[1] && + e.ne[2] == q8_key_ne[2] && e.ne[3] == q8_key_ne[3]) { src1_q8_d = e.buf->ptr; break; } @@ -3611,10 +3820,10 @@ void ggml_cuda_mul_mat_vec_q( char * q8_dst; if (use_q8_memo) { ggml_backend_cuda_context::luce_q8_memo_entry ent; - ent.src1_node = (const void *) src1; + ent.src1_node = (const void *) q8_key_node; ent.src1_data = (const void *) src1_d; - ent.src0_type = (int) src0->type; - ent.ne[0] = ne10; ent.ne[1] = ne11; ent.ne[2] = ne12; ent.ne[3] = ne13; + ent.src0_type = q8_key_type; + ent.ne[0] = q8_key_ne[0]; ent.ne[1] = q8_key_ne[1]; ent.ne[2] = q8_key_ne[2]; ent.ne[3] = q8_key_ne[3]; ent.buf = std::make_unique>(ctx.pool(), q8_bytes); q8_dst = ent.buf->ptr; ctx.luce_q8_memo.push_back(std::move(ent)); @@ -3671,7 +3880,8 @@ void ggml_cuda_mul_mat_vec_q( // Batches above MMID_GROUPED_MAX_PAIRS fall through to the legacy // per-expert kernel instead of aborting the request. const int np = (int) (nchannels_dst*ncols_dst); - ggml_cuda_pool_alloc mmid_meta(ctx.pool(), MMID_META_INTS); + ggml_cuda_pool_alloc mmid_meta(ctx.pool()); + int32_t * meta_d = nullptr; float * gate_w = nullptr; int gate_w_stride = 0; float gate_tau = 0.0f; @@ -3681,31 +3891,63 @@ void ggml_cuda_mul_mat_vec_q( gate_w_stride = (int) (gx->weights->nb[1]/sizeof(float)); gate_tau = gx->tau; } - // Adaptive-k writes -1 drop sentinels, so it must gate a scratch copy: - // sibling weights can still take the legacy kernel, which interprets - // ids as unsigned. Fixed top-k only reads ids and avoids both the - // allocation and the tiny device-to-device copy on every expert op. - ggml_cuda_pool_alloc ids_gated(ctx.pool()); - int32_t * prep_ids = const_cast(ids_d); - int prep_ids_stride = (int) ids_stride; - if (gate_w != nullptr) { - prep_ids = ids_gated.alloc((size_t) np); - prep_ids_stride = (int) nchannels_dst; - CUDA_CHECK(cudaMemcpy2DAsync(prep_ids, nchannels_dst*sizeof(int32_t), - ids_d, ids_stride*sizeof(int32_t), - nchannels_dst*sizeof(int32_t), ncols_dst, - cudaMemcpyDeviceToDevice, stream)); - } - // q4/top4 has just 16 pairs. Launch the smallest whole-warp block - // instead of 256 threads; __syncthreads() still covers every pair. - const int prep_warp = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; - const int prep_threads = ((np + prep_warp - 1)/prep_warp)*prep_warp; - mmid_group_prep<<<1, prep_threads, 0, stream>>>( - prep_ids, mmid_meta.ptr, (int) nchannels_dst, (int) ncols_dst, prep_ids_stride, - gate_w, gate_w_stride, gate_tau); - CUDA_CHECK(cudaGetLastError()); + // The gate, up and down products of an MoE layer read the same ids, so + // with fixed top-k (nothing gates them) their expert grouping is the + // same: it is computed once per evaluation and kept beside the + // memoized q8_1 activations (same lifetime, same LIFO release), keyed + // by the ids node like those are by their source node. Only when no + // local pool block is outstanding, which the release order needs. + const bool prep_memo = gate_w == nullptr && src1_q8_1.ptr == nullptr && + ids->buffer != nullptr && ctx.stream_context().concurrent_events.empty(); + if (prep_memo) { + for (const auto & e : ctx.luce_q8_memo) { + if (e.src0_type == MMID_PREP_MEMO_KIND && e.src1_node == (const void *) ids && + e.src1_data == (const void *) ids_d && e.ne[0] == np && e.ne[1] == ncols_dst && + e.ne[2] == nchannels_dst && e.ne[3] == ids_stride) { + meta_d = (int32_t *) e.buf->ptr; + break; + } + } + } + if (meta_d == nullptr) { + if (prep_memo) { + ggml_backend_cuda_context::luce_q8_memo_entry ent; + ent.src1_node = (const void *) ids; + ent.src1_data = (const void *) ids_d; + ent.src0_type = MMID_PREP_MEMO_KIND; + ent.ne[0] = np; ent.ne[1] = ncols_dst; ent.ne[2] = nchannels_dst; ent.ne[3] = ids_stride; + ent.buf = std::make_unique>(ctx.pool(), MMID_META_INTS*sizeof(int32_t)); + meta_d = (int32_t *) ent.buf->ptr; + ctx.luce_q8_memo.push_back(std::move(ent)); + } else { + meta_d = mmid_meta.alloc(MMID_META_INTS); + } + // Adaptive-k writes -1 drop sentinels, so it must gate a scratch copy: + // sibling weights can still take the legacy kernel, which interprets + // ids as unsigned. Fixed top-k only reads ids and avoids both the + // allocation and the tiny device-to-device copy on every expert op. + ggml_cuda_pool_alloc ids_gated(ctx.pool()); + int32_t * prep_ids = const_cast(ids_d); + int prep_ids_stride = (int) ids_stride; + if (gate_w != nullptr) { + prep_ids = ids_gated.alloc((size_t) np); + prep_ids_stride = (int) nchannels_dst; + CUDA_CHECK(cudaMemcpy2DAsync(prep_ids, nchannels_dst*sizeof(int32_t), + ids_d, ids_stride*sizeof(int32_t), + nchannels_dst*sizeof(int32_t), ncols_dst, + cudaMemcpyDeviceToDevice, stream)); + } + // q4/top4 has just 16 pairs. Launch the smallest whole-warp block + // instead of 256 threads; __syncthreads() still covers every pair. + const int prep_warp = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + const int prep_threads = ((np + prep_warp - 1)/prep_warp)*prep_warp; + mmid_group_prep<<<1, prep_threads, 0, stream>>>( + prep_ids, meta_d, (int) nchannels_dst, (int) ncols_dst, prep_ids_stride, + gate_w, gate_w_stride, gate_tau); + CUDA_CHECK(cudaGetLastError()); + } if (mul_mat_vec_q_grouped_dispatch( - src0->type, src0->data, src1_q8_d, mmid_meta.ptr, fusion_local, dst_d, + src0->type, src0->data, src1_q8_d, meta_d, fusion_local, dst_d, (int) ne00, (int) ne01, (int) nchannels_y, (int) s01, (int) stride_col_y, (int) stride_col_dst, (int) s02, (int) stride_channel_y, (int) stride_channel_dst, @@ -3757,11 +3999,19 @@ void ggml_cuda_mul_mat_vec_q( } } + // A flagged column range of Q8_0 rows: its first block in the whole rows + // (ggml_backend_cuda_mul_mat_whole_row_lanes). + int kb_off = 0; + if (!ids && src0->type == GGML_TYPE_Q8_0 && dst->op == GGML_OP_MUL_MAT && + dst->op_params[1] == GGML_CUDA_WHOLE_ROW_LANES && + src0->view_src != nullptr && src0->nb[1] == src0->view_src->nb[1]) { + kb_off = (int) ((src0->view_offs % src0->nb[1]) / sizeof(block_q8_0)); + } mul_mat_vec_q_switch_type( src0->data, src0->type, src1_q8_d, ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, - ne03, ne3, s03, s13, s3, ids_stride, stream); + ne03, ne3, s03, s13, s3, ids_stride, stream, kb_off); } void ggml_cuda_op_mul_mat_vec_q( diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cuh index b8b3d721e..ce116473a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cuh @@ -19,6 +19,8 @@ bool ggml_cuda_mmvq_mmid_grouped_enabled( // Set by ggml_backend_cuda_set_mmvq_batch_invariant (calling thread). bool ggml_cuda_mmvq_batch_invariant(); +bool ggml_cuda_mmvq_rdna4_glu_pair(const ggml_tensor * gate, const ggml_tensor * up, const ggml_tensor * glu); + void ggml_cuda_mul_mat_vec_q(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion = nullptr); @@ -27,3 +29,8 @@ void ggml_cuda_op_mul_mat_vec_q( const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, const int64_t src1_padded_row_size, cudaStream_t stream); + +// Whether MMVQ memoizes q8_1 activations within an evaluation (LUCE_Q8_MEMO). +// A producer that writes an activation's q8_1 form itself may register it in +// ctx.luce_q8_memo under the key MMVQ uses (see ggml_cuda_mul_mat_vec_q). +bool ggml_cuda_luce_q8_memo_on(); 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 60e5553b7..32c4f15f8 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 @@ -373,6 +373,56 @@ static __global__ void moe_combine_vec4_kernel( (size_t)t * output_nb1) = sum; } +// GGML_MOE_FUSED_COMBINE_MASKED: laguna_moe_combine_kernel with the owner's +// masked route weight formed inline: w = weights[e, t] * valid[ids[e, t], t], +// the product the separate mul(weights, get_rows(valid, ids)) rounds. +static __global__ void moe_combine_masked_kernel( + const char * __restrict__ experts, + const char * __restrict__ weights, + const char * __restrict__ valid, + const char * __restrict__ ids, + char * __restrict__ output, + const int n_embd, + const int n_used, + const int n_tokens, + const size_t experts_nb1, + const size_t experts_nb2, + const size_t weights_nb1, + const size_t valid_nb1, + const size_t valid_nb2, + const size_t ids_nb1, + const size_t output_nb1) { + // The replaced mul, mul and repeat_back round every product and sum on + // their own. The product must not fuse into the following add: __fmul_rn + // keeps it a separate rounding on CUDA, where nvcc ignores the clang + // pragma and builds with fast math; on HIP the pragma keeps the plain add + // out of an FMA (HIP's __fadd_rn is a header function that would fuse). +#pragma clang fp contract(off) + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int total = n_embd * n_tokens; + if (idx >= total) return; + + const int h = idx % n_embd; + const int t = idx / n_embd; + float sum = 0.0f; + for (int e = 0; e < n_used; ++e) { + const int32_t id = *(const int32_t *)(ids + (size_t)e * sizeof(int32_t) + (size_t)t * ids_nb1); + // A route masked to -1 (another rank's expert) has no validity row. + if (id < 0) { + continue; + } + const float valid_w = *(const float *)(valid + (size_t)id * valid_nb1 + (size_t)t * valid_nb2); + const float w = *(const float *)(weights + (size_t)e * sizeof(float) + (size_t)t * weights_nb1) * valid_w; + if (w == 0.0f) { + continue; + } + const float v = *(const float *)(experts + (size_t)h * sizeof(float) + + (size_t)e * experts_nb1 + (size_t)t * experts_nb2); + sum = sum + __fmul_rn(v, w); + } + *(float *)(output + (size_t)h * sizeof(float) + (size_t)t * output_nb1) = sum; +} + static void launch_moe_combine( cudaStream_t stream, const char * experts, @@ -1031,6 +1081,26 @@ void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) ggml_cuda_op_ds4_moe_owner_split(ctx, dst); return; } + if (mode == GGML_MOE_FUSED_COMBINE_MASKED) { + const ggml_tensor * experts = dst->src[0]; // [n_embd, n_used, n_tokens] + const ggml_tensor * weights = dst->src[1]; // [n_used, n_tokens] + const ggml_tensor * valid = dst->src[2]; // [1, n_expert, n_tokens] + const ggml_tensor * ids = dst->src[3]; // [n_used, n_tokens] + GGML_ASSERT(experts->nb[0] == sizeof(float) && dst->nb[0] == sizeof(float)); + const int n_embd = (int) experts->ne[0]; + const int n_used = (int) experts->ne[1]; + const int n_tokens = (int) experts->ne[2]; + const int total = n_embd * n_tokens; + constexpr int block = 256; + moe_combine_masked_kernel<<<(total + block - 1) / block, block, 0, ctx.stream()>>>( + (const char *) experts->data, (const char *) weights->data, + (const char *) valid->data, (const char *) ids->data, (char *) dst->data, + n_embd, n_used, n_tokens, + experts->nb[1], experts->nb[2], weights->nb[1], + valid->nb[1], valid->nb[2], ids->nb[1], dst->nb[1]); + CUDA_CHECK(cudaGetLastError()); + return; + } if (mode == GGML_MOE_FUSED_COMBINE) { const ggml_tensor * experts = dst->src[0]; // [n_embd, n_used, n_tokens] const ggml_tensor * weights = dst->src[1]; // [n_used, n_tokens] diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu index 701be1507..cc7045179 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cu @@ -135,7 +135,7 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr } } -template +template static __global__ void rms_norm_f32(const float * x, float * dst, const int ncols, @@ -158,7 +158,9 @@ static __global__ void rms_norm_f32(const float * x, const uint3 add_ncols_packed = make_uint3(0, 0, 0), const uint3 add_nrows_packed = make_uint3(0, 0, 0), const uint3 add_nchannels_packed = make_uint3(0, 0, 0), - const uint3 add_nsamples_packed = make_uint3(0, 0, 0)) { + const uint3 add_nsamples_packed = make_uint3(0, 0, 0), + block_q8_1 * q8 = nullptr, + const int ncols_q8 = 0) { const int nrows = gridDim.x; const int nchannels = gridDim.y; @@ -200,6 +202,32 @@ static __global__ void rms_norm_f32(const float * x, const float mean = tmp / ncols; const float scale = rsqrtf(mean + eps); + if constexpr (do_q8) { + // The output's q8_1 form as quantize_q8_1 writes it for MMVQ (rows + // padded with zero blocks to ncols_q8): a warp holds one 32-wide + // block, reduced in the same butterfly order. + static_assert(do_multiply && !do_add, "q8 output is fused with rms_norm * w only"); + block_q8_1 * yrow = q8 + (int64_t) ((sample*nchannels + channel)*nrows + row) * (ncols_q8 / QK8_1); + for (int col = tid; col < ncols_q8; col += block_size) { + float v = 0.0f; + if (col < ncols) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + v = scale * x[col] * mul[mul_col]; + dst[col] = v; + } + float amax = fabsf(v); + float sum = v; + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + const float d = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(v / d); + yrow[col / QK8_1].qs[col % QK8_1] = q; + if (col % QK8_1 == 0) { + yrow[col / QK8_1].ds = make_half2(d, sum); + } + } + return; + } for (int col = tid; col < ncols; col += block_size) { if constexpr (do_multiply && do_add) { const int mul_col = fastmodulo(col, mul_ncols_packed); @@ -484,7 +512,9 @@ static void rms_norm_mul_f32_cuda(const float * x, const uint32_t add_nchannels, const uint32_t add_nsamples, const float eps, - cudaStream_t stream) { + cudaStream_t stream, + block_q8_1 * q8 = nullptr, + const int ncols_q8 = 0) { const dim3 blocks_num(nrows, nchannels, nsamples); if (mul == nullptr) { rms_norm_f32_cuda(x, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream); @@ -495,16 +525,31 @@ static void rms_norm_mul_f32_cuda(const float * x, const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); + const uint3 z = make_uint3(0, 0, 0); if (ncols < 1024) { const dim3 block_dims(256, 1, 1); - rms_norm_f32<256, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( - x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, - mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed); + if (q8) { + rms_norm_f32<256, true, false, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, + nullptr, 0, 0, 0, z, z, z, z, q8, ncols_q8); + } else { + rms_norm_f32<256, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed); + } } else { const dim3 block_dims(1024, 1, 1); - rms_norm_f32<1024, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( - x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, - mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed); + if (q8) { + rms_norm_f32<1024, true, false, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, + nullptr, 0, 0, 0, z, z, z, z, q8, ncols_q8); + } else { + rms_norm_f32<1024, true><< WARP_SIZE ? 32 * sizeof(float): 0, stream>>>( + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed); + } } } else { const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols); @@ -624,7 +669,8 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } -void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) { +void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor, + void * q8_out, int ncols_q8) { const ggml_tensor * rms_norm_src = (ggml_tensor *) dst->src[0]; float eps = 0.0f; @@ -681,7 +727,7 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * mul_ncols, mul_nrows, mul_nchannels, mul_nsamples, /*add_s00*/ 0, 0, 0, 0, 0, 0, 0, - eps, stream); + eps, stream, (block_q8_1 *) q8_out, ncols_q8); } // dflash: ADD (residual) + RMS_NORM + MUL in one launch. `add_tensor` is the diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh index 69ac62333..3b5c3b698 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/norm.cuh @@ -13,7 +13,10 @@ void ggml_hip_vision_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); #endif void ggml_cuda_op_rms_norm_scale(ggml_backend_cuda_context & ctx, ggml_tensor * norm, ggml_tensor * scale); -void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor); +// q8_out: optional, the output's q8_1 form as MMVQ quantizes it, rows padded +// to ncols_q8 (a multiple of QK8_1) with zero blocks. +void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor, + void * q8_out = nullptr, int ncols_q8 = 0); void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu index cfcd97613..3bd2d006b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu @@ -160,6 +160,208 @@ static __global__ void rope_norm(const T * x, store_coaelsced(x0 * cos_theta - x1 * sin_theta, x0 * sin_theta + x1 * cos_theta); } +// A forward NORMAL (optionally TAIL) rope without frequency factors, as the +// fused rope kernels below run it: its launch parameters, read from the ROPE +// node the way ggml_cuda_op_rope reads them, and one rotated pair with +// rope_norm's arithmetic. +struct rope_norm_fwd { + int n_dims; + int rot_offset; + float freq_scale; + float ext_factor; + float attn_factor; + float theta_scale; + rope_corr_dims corr_dims; +}; + +static rope_norm_fwd rope_norm_fwd_of(const ggml_tensor * rope, const int ncols) { + const int32_t * op = (const int32_t *) rope->op_params; + float freq_base, beta_fast, beta_slow; + rope_norm_fwd r; + r.n_dims = op[1]; + const int n_ctx_orig = op[4]; + memcpy(&freq_base, op + 5, sizeof(float)); + memcpy(&r.freq_scale, op + 6, sizeof(float)); + memcpy(&r.ext_factor, op + 7, sizeof(float)); + memcpy(&r.attn_factor, op + 8, sizeof(float)); + memcpy(&beta_fast, op + 9, sizeof(float)); + memcpy(&beta_slow, op + 10, sizeof(float)); + r.rot_offset = (op[2] & GGML_ROPE_TYPE_TAIL) ? ncols - r.n_dims : 0; + ggml_rope_yarn_corr_dims(r.n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, r.corr_dims.v); + r.theta_scale = powf(freq_base, -2.0f / r.n_dims); + return r; +} + +static __device__ __forceinline__ void rope_norm_fwd_pair( + const rope_norm_fwd & r, const int32_t p, const int i0, const float x0, const float x1, + float & o0, float & o1) { + if (i0 < r.rot_offset || i0 >= r.rot_offset + r.n_dims) { + o0 = x0; + o1 = x1; + return; + } + const int j0 = i0 - r.rot_offset; + const double theta_base = rope_theta_fp64(p, r.theta_scale, j0/2); + float cos_theta; + float sin_theta; + rope_yarn(theta_base/1.0f, r.freq_scale, r.corr_dims, j0, r.ext_factor, r.attn_factor, cos_theta, sin_theta); + o0 = x0 * cos_theta - x1 * sin_theta; + o1 = x0 * sin_theta + x1 * cos_theta; +} + +// RMS_NORM * w, then a NORMAL (optionally TAIL) forward ROPE, then the F32 -> +// F16 copy of the result, one row per block: rms_norm_f32's +// arithmetic, rope_norm's per pair (no frequency factors) and the cast's +// rounding, so the F32 rope output and its F16 copy are bit-identical to the +// three launches they replace. Rows are [ncols] with positions pos[row / ne01]. +template +static __global__ void rms_norm_rope_f16_kernel( + const float * __restrict__ x, const int64_t x_stride, + const float * __restrict__ w, + float * __restrict__ rope_dst, half * __restrict__ dst_f16, + const int ncols, const float eps, const int ne01, + const int32_t * __restrict__ pos, const rope_norm_fwd r) { + const int row = (int) blockIdx.x; + const int tid = (int) threadIdx.x; + x += (int64_t) row * x_stride; + rope_dst += (int64_t) row * ncols; + dst_f16 += (int64_t) row * ncols; + + extern __shared__ float s_normed[]; + __shared__ float s_sum[32]; + + float tmp = 0.0f; + for (int col = tid; col < ncols; col += block_size) { + const float xi = x[col]; + tmp += xi * xi; + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + for (int col = tid; col < ncols; col += block_size) { + s_normed[col] = scale * x[col] * w[col]; + } + __syncthreads(); + + const int i2 = row / ne01; + for (int p = tid; p < ncols / 2; p += block_size) { + const int i0 = 2 * p; + float o0; + float o1; + rope_norm_fwd_pair(r, pos[i2], i0, s_normed[i0 + 0], s_normed[i0 + 1], o0, o1); + rope_dst[i0 + 0] = o0; + rope_dst[i0 + 1] = o1; + dst_f16[i0 + 0] = __float2half(o0); + dst_f16[i0 + 1] = __float2half(o1); + } +} + +// A forward NORMAL/TAIL rope (rope_norm's arithmetic, no frequency factors) +// that also writes its output's q8_1 form in the layout MMVQ quantizes a +// grouped view [G = D * heads_per_group, T, n_groups] of it to: element +// (d, head, t) -> row (head / hpg) * T + t, column (head % hpg) * D + d. A +// half-warp holds one 32-wide block (two values per thread); amax and the +// sum follow quantize_q8_1's 32-lane butterfly exactly. +static __global__ void rope_norm_q8_kernel( + const float * __restrict__ x, float * __restrict__ dst, const int D, const int H, const int T, + const int64_t s01, const int64_t s02, + const int32_t * __restrict__ pos, const rope_norm_fwd r, + block_q8_1 * __restrict__ q8, const int hpg) { + const int row = (int) blockIdx.x; // head + H * token + const int p = (int) (blockIdx.y * blockDim.x + threadIdx.x); + const int i0 = 2 * p; + if (i0 >= D) return; // whole half-warps: D % 32 == 0 + const int head = row % H; + const int t = row / H; + const float x0 = x[(int64_t) head * s01 + (int64_t) t * s02 + i0 + 0]; + const float x1 = x[(int64_t) head * s01 + (int64_t) t * s02 + i0 + 1]; + float o0; + float o1; + rope_norm_fwd_pair(r, pos[t], i0, x0, x1, o0, o1); + dst[(int64_t) row * D + i0 + 0] = o0; + dst[(int64_t) row * D + i0 + 1] = o1; + + // quantize_q8_1: lane v of the block holds element v; here thread v/2. + float amax = fmaxf(fabsf(o0), fabsf(o1)); + float v0 = o0; + float v1 = o1; +#pragma unroll + for (int off = 8; off > 0; off >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, off, 16)); + v0 += __shfl_xor_sync(0xffffffff, v0, off, 16); + v1 += __shfl_xor_sync(0xffffffff, v1, off, 16); + } + const float sum = v0 + v1; + const float d = amax / 127.0f; + const int8_t q0 = amax == 0.0f ? 0 : roundf(o0 / d); + const int8_t q1 = amax == 0.0f ? 0 : roundf(o1 / d); + const int G = D * hpg; + const int g = head / hpg; + const int64_t col = (int64_t) (head % hpg) * D + i0; + const int64_t qi = ((int64_t) g * T + t) * G + col; + block_q8_1 * yb = q8 + qi / QK8_1; + yb->qs[qi % QK8_1 + 0] = q0; + yb->qs[qi % QK8_1 + 1] = q1; + if (qi % QK8_1 == 0) { + yb->ds = make_half2(d, sum); + } +} + +void ggml_cuda_rope_q8(ggml_backend_cuda_context & ctx, const ggml_tensor * rope, void * q8, int hpg) { + const ggml_tensor * x = rope->src[0]; + const int D = (int) x->ne[0]; + const int H = (int) x->ne[1]; + const int T = (int) x->ne[2]; + const dim3 grid(H * T, (D / 2 + 255) / 256, 1); + rope_norm_q8_kernel<<>>( + (const float *) x->data, (float *) rope->data, D, H, T, + x->nb[1] / sizeof(float), x->nb[2] / sizeof(float), + (const int32_t *) rope->src[1]->data, rope_norm_fwd_of(rope, D), + (block_q8_1 *) q8, hpg); + CUDA_CHECK(cudaGetLastError()); +} + +bool ggml_cuda_rms_norm_rope_f16_supported(int ncols) { +#if defined(GGML_USE_HIP) + return ncols % 2 == 0 && (size_t) ncols * sizeof(float) <= 32u * 1024u; +#else + // Bit-exact on HIP only: nvcc builds with fast math and contracts the + // rotation's products differently here than in rope_norm, so CUDA keeps + // the three launches. + GGML_UNUSED(ncols); + return false; +#endif +} + +void ggml_cuda_rms_norm_rope_f16(ggml_backend_cuda_context & ctx, const ggml_tensor * norm, + const ggml_tensor * mul, const ggml_tensor * rope, const ggml_tensor * cpy) { + const ggml_tensor * x = norm->src[0]; + const ggml_tensor * w = mul->src[0] == norm ? mul->src[1] : mul->src[0]; + float eps = 0.0f; + memcpy(&eps, norm->op_params, sizeof(float)); + const int ncols = (int) x->ne[0]; + const int nrows = (int) ggml_nrows(x); + const ggml_tensor * rope_src = rope->src[0]; + const int ne01 = (int) rope_src->ne[1]; + + const rope_norm_fwd r = rope_norm_fwd_of(rope, ncols); + + const size_t shmem = (size_t) ncols * sizeof(float); + const int64_t x_stride = x->nb[1] / sizeof(float); + if (ncols < 1024) { + rms_norm_rope_f16_kernel<256><<>>( + (const float *) x->data, x_stride, (const float *) w->data, + (float *) rope->data, (half *) cpy->data, ncols, eps, ne01, + (const int32_t *) rope->src[1]->data, r); + } else { + rms_norm_rope_f16_kernel<1024><<>>( + (const float *) x->data, x_stride, (const float *) w->data, + (float *) rope->data, (half *) cpy->data, ncols, eps, ne01, + (const int32_t *) rope->src[1]->data, r); + } + CUDA_CHECK(cudaGetLastError()); +} + template static __global__ void rope_neox(const T * x, D * dst, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cuh index 304ef60ae..7dca1360c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cuh @@ -8,5 +8,15 @@ void ggml_cuda_op_rope_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * set_rows); +// RMS_NORM * w -> forward NORMAL/TAIL ROPE -> F16 copy in one launch (see +// rope.cu); writes the ROPE node's F32 output and the copy's F16 output. +bool ggml_cuda_rms_norm_rope_f16_supported(int ncols); +void ggml_cuda_rms_norm_rope_f16(ggml_backend_cuda_context & ctx, const ggml_tensor * norm, + const ggml_tensor * mul, const ggml_tensor * rope, const ggml_tensor * cpy); + +// A forward NORMAL/TAIL rope of a contiguous [D, H, T] F32 tensor that also +// writes the q8_1 form MMVQ quantizes its grouped view [D * hpg, T, H / hpg] +// to (see rope.cu). +void ggml_cuda_rope_q8(ggml_backend_cuda_context & ctx, const ggml_tensor * rope, void * q8, int hpg); // ROPE (multi-section, F32) -> PERMUTE -> CONT in one kernel: the rotation is written straight into cont's layout. void ggml_cuda_op_rope_permuted(ggml_backend_cuda_context & ctx, ggml_tensor * rope, ggml_tensor * cont); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu index 1ed928d28..8df3d3f42 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu @@ -55,6 +55,16 @@ struct soft_max_params { #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wpass-failed" #endif // __clang__ +// The sink-column mode reproduces scale, an additive mask and the sink concat +// as separate launches round them: the scaled score is rounded before the +// mask is added (no FMA). +static __device__ __forceinline__ float soft_max_sink_col_value(float x, float scale, float m) { + // __fmul_rn keeps the product its own rounding on CUDA, where nvcc ignores + // the clang pragma and builds with fast math; the pragma covers HIP. +#pragma clang fp contract(off) + return __fmul_rn(x, scale) + m; +} + template static __global__ void soft_max_f32( const float * x, const T * mask, const float * sinks, float * dst, const soft_max_params p) { @@ -100,7 +110,9 @@ static __global__ void soft_max_f32( const float val = (p.sink_col && col == ncols_x) ? sinks[rowx] - : x[col]*p.scale + (mask ? slope*t2f32(mask[col]) : 0.0f); + : p.sink_col + ? soft_max_sink_col_value(x[col], p.scale, mask ? slope*t2f32(mask[col]) : 0.0f) + : x[col]*p.scale + (mask ? slope*t2f32(mask[col]) : 0.0f); // The global-memory cache is dst, which has no slot for the virtual sink. if (use_shared || col < ncols_x) { @@ -456,7 +468,8 @@ void ggml_cuda_op_soft_max(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { memcpy(&sink_col, (const int32_t *) dst->op_params + 2, sizeof(int32_t)); params.sink_col = sink_col; if (sink_col) { - GGML_ASSERT(src1 == nullptr && src2 != nullptr); + GGML_ASSERT(src2 != nullptr); + GGML_ASSERT(max_bias == 0.0f); GGML_ASSERT(ggml_nelements(src2) == nrows_x); params.ncols = ne00 + 1; // one virtual sink column per row } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cu index ed7f3d02a..ce1773163 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cu @@ -357,13 +357,49 @@ static __global__ void swiglu_ds4_kernel(const T * gate, const T * up, T * dst, dst[i] = ggml_cuda_op_swiglu_ds4_single(gate_v, up_v, limit); } +// swiglu_ds4_kernel that also writes the result's q8_1 form as MMVQ's +// quantize_q8_1 would (rows of n padded with zero blocks to n_q8): a warp's +// 32 consecutive elements are one block (n % QK8_1 == 0). Threads past the +// k results write the padding blocks. +static __global__ void swiglu_ds4_q8_kernel(const float * gate, const float * up, float * dst, + const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, + float limit, block_q8_1 * q8, const int64_t n_q8) { + const int64_t i = int64_t(blockDim.x)*blockIdx.x + threadIdx.x; + const int64_t rows = k / n; + const int64_t pad = n_q8 - n; + float v = 0.0f; + int64_t qi; // element index in the padded q8 layout + if (i < k) { + const int64_t j0 = (i / n) * o0 + (i % n); + const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); + v = ggml_cuda_op_swiglu_ds4_single(gate[j0], up[j1], limit); + dst[i] = v; + qi = (i / n) * n_q8 + (i % n); + } else if (i - k < rows * pad) { + const int64_t p = i - k; + qi = (p / pad) * n_q8 + n + (p % pad); + } else { + return; // whole warps only: k and rows * pad are multiples of QK8_1 + } + float amax = fabsf(v); + float sum = v; + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + const float d = amax / 127.0f; + const int8_t q = amax == 0.0f ? 0 : roundf(v / d); + q8[qi / QK8_1].qs[qi % QK8_1] = q; + if (qi % QK8_1 == 0) { + q8[qi / QK8_1].ds = make_half2(d, sum); + } +} + template static void swiglu_ds4_cuda(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, const float limit, cudaStream_t stream) { const int64_t num_blocks = (k + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE; swiglu_ds4_kernel<<>>(gate, up, dst, k, n, o0, o1, limit); } -void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst, void * q8_out, int ncols_q8) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; void * src0_d = src0->data; @@ -392,6 +428,16 @@ void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const float limit = ggml_get_op_params_f32(dst, 2); + if (q8_out) { + GGML_ASSERT(nc % QK8_1 == 0 && ncols_q8 % QK8_1 == 0 && ncols_q8 >= nc); + const int64_t k = ggml_nelements(dst); + const int64_t total = k + ggml_nrows(dst) * (ncols_q8 - nc); + const int64_t num_blocks = (total + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE; + swiglu_ds4_q8_kernel<<>>( + (const float *) src0_d, (const float *) src1_d, (float *) dst_d, k, nc, + src0_o / sizeof(float), src1_o / sizeof(float), limit, (block_q8_1 *) q8_out, ncols_q8); + return; + } swiglu_ds4_cuda((float *) src0_d, (float *) src1_d, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), limit, stream); } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cuh index 1864ce2dc..39ca8dad9 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cuh @@ -81,7 +81,8 @@ void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +// q8_out: optional, the result's q8_1 form as MMVQ quantizes it (rows padded to ncols_q8). +void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst, void * q8_out = nullptr, int ncols_q8 = 0); void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 6052b6367..010c3a936 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -4426,6 +4426,25 @@ struct ggml_tensor * ggml_soft_max_ext_sink_col( return result; } +struct ggml_tensor * ggml_soft_max_ext_sink_col_mask( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + struct ggml_tensor * sinks, + float scale) { + struct ggml_tensor * result = ggml_soft_max_ext_sink_col(ctx, a, sinks, scale); + if (mask) { + GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); + GGML_ASSERT(mask->nb[0] == ggml_type_size(mask->type)); + GGML_ASSERT(mask->ne[0] == a->ne[0]); + GGML_ASSERT(mask->ne[1] >= a->ne[1]); + GGML_ASSERT(a->ne[2] % mask->ne[2] == 0); + GGML_ASSERT(a->ne[3] % mask->ne[3] == 0); + result->src[1] = mask; + } + return result; +} + void ggml_soft_max_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks) { @@ -5831,6 +5850,8 @@ void ggml_flash_attn_ext_set_ds4_inverse_rope( bool q_unrotated) { GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); GGML_ASSERT(kv_start >= 0); + // Lane rows (bit 2) hold src[6] for their ring rows; no fused RoPE then. + GGML_ASSERT((ggml_get_op_params_i32(a, 7) & 4) == 0); // Bit 0: inverse RoPE on attention output. Bit 1: Q still needs forward // RoPE. Both directions share the position and YaRN parameters below. // Preserve optional DS4 causal-layout metadata in the upper 16 bits so @@ -5862,6 +5883,31 @@ void ggml_flash_attn_ext_set_ds4_rope_positions( a->src[6] = positions; } +void ggml_flash_attn_ext_set_ds4_lane_rows( + struct ggml_tensor * a, + struct ggml_tensor * ring_rows, + struct ggml_tensor * lane_kv) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + // src[6] holds the ring rows: no fused RoPE, whose positions live there. + GGML_ASSERT((ggml_get_op_params_i32(a, 7) & 0xffff) == 0); + GGML_ASSERT(a->src[6] == NULL && a->src[9] == NULL); + // One latent K/V tensor (MLA), as for the segmented layout. + GGML_ASSERT(a->src[1] == a->src[2]); + const int64_t n_query = a->src[0]->ne[1]; + GGML_ASSERT(ring_rows && ring_rows->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(ring_rows)); + GGML_ASSERT(ring_rows->ne[0] == n_query && ring_rows->ne[1] == 1 && + ring_rows->ne[2] == 1 && ring_rows->ne[3] == 1); + GGML_ASSERT(lane_kv && lane_kv->type == a->src[1]->type); + GGML_ASSERT(ggml_is_contiguous(lane_kv)); + GGML_ASSERT(lane_kv->ne[0] == a->src[1]->ne[0] && lane_kv->ne[1] == n_query && + lane_kv->ne[2] == 1 && lane_kv->ne[3] == 1); + const uint32_t flags = (uint32_t) ggml_get_op_params_i32(a, 7); + ggml_set_op_params_i32(a, 7, (int32_t) (flags | 4u)); + a->src[6] = ring_rows; + a->src[9] = lane_kv; +} + bool ggml_flash_attn_ext_is_ds4(const struct ggml_tensor * a) { return a && a->op == GGML_OP_FLASH_ATTN_EXT && (ggml_get_op_params_i32(a, 6) != 0 || @@ -9126,6 +9172,33 @@ struct ggml_tensor * ggml_laguna_moe_combine( return result; } +struct ggml_tensor * ggml_moe_combine_masked( + struct ggml_context * ctx, + struct ggml_tensor * experts, + struct ggml_tensor * weights, + struct ggml_tensor * valid, + struct ggml_tensor * ids) { + GGML_ASSERT(experts->type == GGML_TYPE_F32 && weights->type == GGML_TYPE_F32); + GGML_ASSERT(valid->type == GGML_TYPE_F32 && ids->type == GGML_TYPE_I32); + GGML_ASSERT(experts->ne[1] == weights->ne[0] && experts->ne[2] == weights->ne[1]); + GGML_ASSERT(ids->ne[0] == weights->ne[0] && ids->ne[1] == weights->ne[1]); + GGML_ASSERT(valid->ne[0] == 1 && valid->ne[2] == weights->ne[1]); + GGML_ASSERT(experts->nb[0] == sizeof(float) && weights->nb[0] == sizeof(float)); + GGML_ASSERT(ids->nb[0] == sizeof(int32_t)); + + const int64_t ne[4] = { experts->ne[0], experts->ne[2], 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne); + + result->op = GGML_OP_MOE_FUSED; + result->src[0] = experts; + result->src[1] = weights; + result->src[2] = valid; + result->src[3] = ids; + + ggml_set_op_params_i32(result, 0, GGML_MOE_FUSED_COMBINE_MASKED); + return result; +} + struct ggml_tensor * ggml_ds4_moe_owner( struct ggml_context * ctx, struct ggml_tensor * input, @@ -9353,16 +9426,19 @@ struct ggml_tensor * ggml_ds4_deferred_peer_copy( return result; } -struct ggml_tensor * ggml_cluster_allreduce( +static struct ggml_tensor * ggml_cluster_allreduce_impl( struct ggml_context * ctx, struct ggml_tensor * a, ggml_cluster_allreduce_fn fn, - void * user) { + void * user, + bool inplace) { GGML_ASSERT(a->type == GGML_TYPE_F32); GGML_ASSERT(ggml_is_contiguous(a)); GGML_ASSERT(fn != NULL); - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + // In place, the backend sees dst->data == src->data and enqueues only + // the collective. + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); result->op = GGML_OP_MOE_FUSED; result->src[0] = a; @@ -9375,6 +9451,22 @@ struct ggml_tensor * ggml_cluster_allreduce( return result; } +struct ggml_tensor * ggml_cluster_allreduce( + struct ggml_context * ctx, + struct ggml_tensor * a, + ggml_cluster_allreduce_fn fn, + void * user) { + return ggml_cluster_allreduce_impl(ctx, a, fn, user, false); +} + +struct ggml_tensor * ggml_cluster_allreduce_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + ggml_cluster_allreduce_fn fn, + void * user) { + return ggml_cluster_allreduce_impl(ctx, a, fn, user, true); +} + struct ggml_tensor * ggml_ds4_hc_pre( struct ggml_context * ctx, struct ggml_tensor * mix, @@ -9595,6 +9687,39 @@ struct ggml_tensor * ggml_ds4_router_weights( return result; } +// GGML_OP_DS4_HC modes 7 and 8. +struct ggml_tensor * ggml_ds4_argmax_pair( + struct ggml_context * ctx, + struct ggml_tensor * logits, + int row_offset, + int slot, + int n_slots) { + GGML_ASSERT(logits->type == GGML_TYPE_F32 && ggml_is_contiguous(logits)); + GGML_ASSERT(n_slots >= 1 && n_slots <= 8 && slot >= 0 && slot < n_slots && row_offset >= 0); + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 8 * n_slots, logits->ne[1]); + result->op = GGML_OP_DS4_HC; + result->src[0] = logits; + ggml_set_op_params_i32(result, 0, 7); + ggml_set_op_params_i32(result, 1, row_offset); + ggml_set_op_params_i32(result, 2, slot); + ggml_set_op_params_i32(result, 3, n_slots); + return result; +} + +struct ggml_tensor * ggml_ds4_argmax_pick( + struct ggml_context * ctx, + struct ggml_tensor * pairs, + int n_slots) { + GGML_ASSERT(pairs->type == GGML_TYPE_F32 && ggml_is_contiguous(pairs)); + GGML_ASSERT(pairs->ne[0] == 8 * n_slots); + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, pairs->ne[1]); + result->op = GGML_OP_DS4_HC; + result->src[0] = pairs; + ggml_set_op_params_i32(result, 0, 8); + ggml_set_op_params_i32(result, 3, n_slots); + return result; +} + // GGML_OP_DS4_HC mode 6. struct ggml_tensor * ggml_ds4_hc_collapse( struct ggml_context * ctx, diff --git a/server/scripts/cluster/build_expert_placement.py b/server/scripts/cluster/build_expert_placement.py new file mode 100644 index 000000000..76b87da8f --- /dev/null +++ b/server/scripts/cluster/build_expert_placement.py @@ -0,0 +1,245 @@ +#!/usr/bin/env python3 +"""Build a cluster expert placement JSON from a routing hotness CSV. + +Input: the CSV written by MoeHybridRoutingStats::save_csv +(DFLASH_DS4_ROUTING_STATS_OUT): + + # hotness table: n_layer=43 n_expert=256 n_expert_used=6 + # format: one row per layer, columns are expert activation counts (expert 0..N-1) + 12,0,5,... <- layer 0 + ... + +Output: the JSON consumed by ClusterExpertPlacement::load_json +(server/src/cluster/cluster_expert_placement.cpp) and passed to +--cluster-expert-placement : + + {"n_ranks":N,"n_layer":L,"n_expert":E,"n_expert_used":K, + "replicate_hot":R,"owner":[[...E ints...] x L]} + +owner[layer][expert] is the owning rank or -1 (replicated on every rank; the +rank that evaluates expert e for token slot t is (e + t) % N). + +The algorithms mirror the C++ builders exactly (uniform: e % N; balanced: +hottest replicate_hot experts replicated, the rest sorted by count descending +and greedily assigned to the least-loaded rank below ceil(remaining/N) owned +experts, ties to the lowest rank) so a JSON produced here hashes identically +to a placement built in-process from the same CSV. + +stdlib only. +""" + +import argparse +import json +import re +import sys + + +def load_hotness_csv(path): + n_layer = n_expert = n_expert_used = None + rows = [] + header_re = re.compile(r"#\s*hotness table:\s*n_layer=(\d+)\s+n_expert=(\d+)\s+n_expert_used=(\d+)") + with open(path, "r", encoding="utf-8") as f: + for raw in f: + line = raw.strip() + if not line: + continue + if line.startswith("#"): + m = header_re.match(line) + if m: + n_layer, n_expert, n_expert_used = (int(m.group(i)) for i in (1, 2, 3)) + continue + rows.append([int(x) for x in line.split(",") if x.strip() != ""]) + if n_layer is None: + raise ValueError("missing '# hotness table:' header") + if len(rows) != n_layer: + raise ValueError("expected %d layer rows, found %d" % (n_layer, len(rows))) + for il, row in enumerate(rows): + if len(row) != n_expert: + raise ValueError("layer %d has %d columns, expected %d" % (il, len(row), n_expert)) + return n_layer, n_expert, n_expert_used, rows + + +def ranked_experts(counts): + # count descending, expert id ascending (stable, same as C++ ranked_experts) + return sorted(range(len(counts)), key=lambda e: (-counts[e], e)) + + +def build_uniform(n_ranks, n_layer, n_expert, rows, replicate_hot): + owner = [] + for il in range(n_layer): + row = [e % n_ranks for e in range(n_expert)] + if replicate_hot > 0: + for e in ranked_experts(rows[il])[:replicate_hot]: + row[e] = -1 + owner.append(row) + return owner + + +def build_balanced(n_ranks, n_layer, n_expert, rows, replicate_hot): + owner = [] + for il in range(n_layer): + counts = rows[il] + ranked = ranked_experts(counts) + row = [0] * n_expert + for e in ranked[:replicate_hot]: + row[e] = -1 + remaining = n_expert - replicate_hot + cap = -(-remaining // n_ranks) # ceil + load = [0] * n_ranks + owned = [0] * n_ranks + for e in ranked[replicate_hot:]: + best = -1 + for r in range(n_ranks): + if owned[r] >= cap: + continue + if best < 0 or load[r] < load[best]: + best = r + if best < 0: + raise RuntimeError("no rank below cap at layer %d" % il) + row[e] = best + load[best] += counts[e] + owned[best] += 1 + owner.append(row) + return owner + + +def build_all_on_rank(rank, n_ranks, n_layer, n_expert): + # Diagnostic: one rank owns every expert, the others hold empty shards. + # The cluster must then reproduce the single-node output exactly. + if rank < 0 or rank >= n_ranks: + raise ValueError("--all-on-rank must be in [0, ranks)") + return [[rank] * n_expert for _ in range(n_layer)] + + +def validate(owner, n_ranks, n_expert, allow_empty=False): + for il, row in enumerate(owner): + replicated = sum(1 for o in row if o == -1) + owned = [0] * n_ranks + for e, o in enumerate(row): + if o == -1: + continue + if o < 0 or o >= n_ranks: + raise ValueError("layer %d expert %d owner %d out of range" % (il, e, o)) + owned[o] += 1 + if n_expert - replicated >= n_ranks and min(owned) == 0: + if allow_empty: + print("warning: layer %d: rank %d owns no expert (diagnostic placement)" + % (il, owned.index(0)), file=sys.stderr) + else: + raise ValueError("layer %d: rank %d owns no expert" % (il, owned.index(0))) + + +def report(owner, rows, n_ranks, n_layer, n_expert): + total_owned = [0] * n_ranks + replicated_total = 0 + worst = 0.0 + ratio_sum = 0.0 + layers_counted = 0 + per_layer = [] + for il in range(n_layer): + row = owner[il] + counts = rows[il] + load = [0.0] * n_ranks + owned = [0] * n_ranks + total = float(sum(counts)) + for e, o in enumerate(row): + if o == -1: + replicated_total += 1 + for r in range(n_ranks): + load[r] += counts[e] / n_ranks + else: + owned[o] += 1 + total_owned[o] += 1 + load[o] += counts[e] + if total > 0: + mean = total / n_ranks + mx = max(load) + ratio = mx / mean if mean > 0 else 1.0 + worst = max(worst, ratio) + ratio_sum += ratio + layers_counted += 1 + per_layer.append((il, owned, mx, mean, ratio)) + else: + per_layer.append((il, owned, 0.0, 0.0, 1.0)) + + print("ranks=%d layers=%d experts=%d replicated_total=%d" % (n_ranks, n_layer, n_expert, replicated_total)) + print("owned experts per rank (all layers): %s" % total_owned) + print("layer owned_per_rank max_load mean_load max/mean") + for il, owned, mx, mean, ratio in per_layer: + print("%5d %-24s %11.0f %11.0f %.3f" % (il, owned, mx, mean, ratio)) + if layers_counted: + print("expected imbalance (max/mean): worst=%.3f avg=%.3f over %d layers" + % (worst, ratio_sum / layers_counted, layers_counted)) + else: + print("no routing counts in CSV (all zero): load report skipped") + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("csv", nargs="?", help="routing hotness CSV (MoeHybridRoutingStats::save_csv); " + "optional with --dims") + ap.add_argument("--dims", help="L,E,K (n_layer,n_expert,n_expert_used) when no CSV is given, " + "e.g. 43,256,6 for DeepSeek V4 Flash") + ap.add_argument("--all-on-rank", type=int, default=None, + help="diagnostic: rank R owns every expert, other ranks get empty shards") + ap.add_argument("--ranks", "-n", type=int, required=True, help="cluster size N (2..8)") + ap.add_argument("--mode", choices=("uniform", "balanced"), default="balanced") + ap.add_argument("--replicate-hot", type=int, default=0, help="replicate the k hottest experts per layer") + ap.add_argument("--out", "-o", required=True, help="output placement JSON") + ap.add_argument("--quiet", action="store_true") + args = ap.parse_args() + + if args.ranks < 1: + sys.exit("--ranks must be >= 1") + if args.replicate_hot < 0: + sys.exit("--replicate-hot must be >= 0") + + if args.csv: + n_layer, n_expert, n_expert_used, rows = load_hotness_csv(args.csv) + elif args.dims: + n_layer, n_expert, n_expert_used = (int(x) for x in args.dims.split(",")) + rows = [[0] * n_expert for _ in range(n_layer)] + else: + sys.exit("give a hotness CSV or --dims L,E,K") + if args.replicate_hot > n_expert - args.ranks: + sys.exit("--replicate-hot leaves fewer experts than ranks") + + if args.all_on_rank is not None: + owner = build_all_on_rank(args.all_on_rank, args.ranks, n_layer, n_expert) + args.replicate_hot = 0 + elif args.mode == "uniform": + owner = build_uniform(args.ranks, n_layer, n_expert, rows, args.replicate_hot) + else: + if not args.csv: + sys.exit("balanced placement needs a hotness CSV") + owner = build_balanced(args.ranks, n_layer, n_expert, rows, args.replicate_hot) + validate(owner, args.ranks, n_expert, allow_empty=args.all_on_rank is not None) + + doc = { + "n_ranks": args.ranks, + "n_layer": n_layer, + "n_expert": n_expert, + "n_expert_used": n_expert_used, + "replicate_hot": args.replicate_hot, + "owner": owner, + } + with open(args.out, "w", encoding="utf-8") as f: + # one owner row per line keeps the file diffable + f.write("{\n") + f.write(' "n_ranks": %d,\n' % doc["n_ranks"]) + f.write(' "n_layer": %d,\n' % doc["n_layer"]) + f.write(' "n_expert": %d,\n' % doc["n_expert"]) + f.write(' "n_expert_used": %d,\n' % doc["n_expert_used"]) + f.write(' "replicate_hot": %d,\n' % doc["replicate_hot"]) + f.write(' "owner": [\n') + for il, row in enumerate(owner): + f.write(" " + json.dumps(row, separators=(",", ":")) + (",\n" if il + 1 < n_layer else "\n")) + f.write(" ]\n}\n") + + if not args.quiet: + report(owner, rows, args.ranks, n_layer, n_expert) + print("wrote %s" % args.out) + + +if __name__ == "__main__": + main() diff --git a/server/src/cluster/cluster_comm.cpp b/server/src/cluster/cluster_comm.cpp new file mode 100644 index 000000000..ee7e846d3 --- /dev/null +++ b/server/src/cluster/cluster_comm.cpp @@ -0,0 +1,871 @@ +// cluster_comm.cpp - collective communication between cluster ranks +// (cluster_comm.h). +// +// Two implementations: +// * LoopbackComm: always compiled. size-1 communicator whose collectives +// are no-ops that still count into the stats, so cluster-enabled code +// paths run unchanged on a single node and in GPU-free unit tests. +// * RcclComm: compiled only under LUCE_CLUSTER_RCCL (set by +// -DLUCE_CLUSTER=ON on a HIP build). Wraps one ncclComm_t created +// with the non-blocking config so a watchdog can ncclCommAbort a hung +// collective; every device collective is enqueued on the caller's HIP +// stream and never synchronizes the host. +// +// Without the define the RCCL entry points return false / nullptr with +// err = "RCCL support not compiled in (build with -DLUCE_CLUSTER=ON on HIP)". + +#include "cluster/cluster_comm.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef LUCE_CLUSTER_RCCL +#include +#include +#endif + +namespace luce::cluster { + +uint64_t cluster_device_bytes_in_use(int device) { +#ifdef LUCE_CLUSTER_RCCL + if (device < 0) return 0; + int previous = 0; + if (hipGetDevice(&previous) != hipSuccess) return 0; + if (device != previous && hipSetDevice(device) != hipSuccess) return 0; + size_t free_bytes = 0; + size_t total_bytes = 0; + const hipError_t rc = hipMemGetInfo(&free_bytes, &total_bytes); + if (device != previous) hipSetDevice(previous); + if (rc != hipSuccess || total_bytes == 0 || free_bytes > total_bytes) return 0; + return (uint64_t) (total_bytes - free_bytes); +#else + (void) device; + return 0; +#endif +} + +namespace { + +const char * const kNoRcclMsg = + "RCCL support not compiled in (build with -DLUCE_CLUSTER=ON on HIP)"; + +void set_err(std::string * err, const std::string & msg) { + if (err) *err = msg; +} + +uint64_t now_us() { + using namespace std::chrono; + return (uint64_t)duration_cast(steady_clock::now().time_since_epoch()).count(); +} + +// ─── Loopback ─────────────────────────────────────────────────────────── + +class LoopbackComm final : public IClusterComm { +public: + LoopbackComm(int rank, int size) : rank_(rank), size_(size) {} + + int rank() const override { return rank_; } + int size() const override { return size_; } + + bool allreduce_sum_f32(void * dev_ptr, size_t n, DeviceStream stream, std::string * err) override { + (void)dev_ptr; (void)stream; (void)err; + std::lock_guard lk(mu_); + stats_.allreduce_calls++; + stats_.allreduce_bytes += n * sizeof(float); + return true; + } + + bool allreduce_sum_bf16_compressed(void * dev_f32, size_t n, void * scratch_bf16, + DeviceStream stream, std::string * err) override { + (void)dev_f32; (void)scratch_bf16; (void)stream; (void)err; + std::lock_guard lk(mu_); + stats_.allreduce_calls++; + stats_.allreduce_bytes += n * 2; + return true; + } + + bool broadcast_bytes(void * dev_buf, size_t bytes, int root, DeviceStream stream, + std::string * err) override { + (void)dev_buf; (void)stream; + if (root != rank_) { + set_err(err, "loopback broadcast root " + std::to_string(root) + " is not this rank"); + return false; + } + std::lock_guard lk(mu_); + stats_.broadcast_calls++; + stats_.broadcast_bytes += bytes; + return true; + } + + bool broadcast_host_i32(int32_t * host_buf, size_t count, int root, std::string * err) override { + (void)host_buf; + if (root != rank_) { + set_err(err, "loopback broadcast root " + std::to_string(root) + " is not this rank"); + return false; + } + std::lock_guard lk(mu_); + stats_.broadcast_calls++; + stats_.broadcast_bytes += count * sizeof(int32_t); + return true; + } + + bool barrier(std::string * err) override { + (void)err; + std::lock_guard lk(mu_); + stats_.barrier_calls++; + return true; + } + + bool wait_stream(DeviceStream stream, uint32_t deadline_ms, std::string * err) override { + (void)stream; (void)deadline_ms; (void)err; + return true; + } + + void abort() override { + std::lock_guard lk(mu_); + if (!aborted_) { + aborted_ = true; + stats_.aborts++; + } + } + + bool async_error(std::string * err) const override { + (void)err; + return false; + } + + ClusterCommStats stats() const override { + std::lock_guard lk(mu_); + return stats_; + } + + void reset_stats() override { + std::lock_guard lk(mu_); + stats_ = ClusterCommStats{}; + } + + const char * backend_name() const override { return "loopback"; } + +private: + const int rank_; + const int size_; + mutable std::mutex mu_; + ClusterCommStats stats_; + bool aborted_ = false; +}; + +} // namespace + +std::unique_ptr create_loopback_cluster_comm(int rank, int size) { + if (size < 1) size = 1; + if (rank < 0 || rank >= size) rank = 0; + return std::unique_ptr(new LoopbackComm(rank, size)); +} + +// ═══════════════════════════════════════════════════════════════════════ +#ifdef LUCE_CLUSTER_RCCL + +static_assert(sizeof(ncclUniqueId) == kRcclUniqueIdSize, + "ncclUniqueId size differs from kRcclUniqueIdSize (NCCL_UNIQUE_ID_BYTES)"); + +// Defined in cluster_comm_kernels.cu (LANGUAGE HIP). +bool cluster_kernel_f32_to_bf16(const float * in, uint16_t * out, size_t n, void * stream); +bool cluster_kernel_bf16_to_f32(const uint16_t * in, float * out, size_t n, void * stream); +bool cluster_kernel_f32_to_q8(const float * in, void * slot, size_t n, void * stream); +bool cluster_kernel_q8_sum(const void * slots, int ranks, size_t slot_bytes, size_t n, float * out, + void * stream); + +namespace { + +std::string hip_err(const char * what, hipError_t e) { + return std::string(what) + ": " + hipGetErrorString(e); +} + +std::string nccl_err(const char * what, ncclResult_t r) { + return std::string(what) + ": " + ncclGetErrorString(r); +} + +// Non-blocking communicators report enqueue success as ncclInProgress. +bool nccl_ok(ncclResult_t r) { + return r == ncclSuccess || r == ncclInProgress; +} + +class RcclComm final : public IClusterComm { +public: + RcclComm(int rank, int size, int device, uint32_t timeout_ms) + : rank_(rank), size_(size), device_(device), timeout_ms_(timeout_ms) {} + + ~RcclComm() override { + (void)hipSetDevice(device_); + if (comm_ != nullptr) { + if (!aborted_) { + // Non-blocking comms may report InProgress here; nothing to do + // about it at teardown. + (void)ncclCommDestroy(comm_); + } + comm_ = nullptr; + } + if (wait_event_ != nullptr) (void)hipEventDestroy(wait_event_); + if (internal_stream_ != nullptr) (void)hipStreamDestroy(internal_stream_); + if (staging_dev_ != nullptr) (void)hipFree(staging_dev_); + if (staging_host_ != nullptr) (void)hipHostFree(staging_host_); + } + + bool init(const RcclUniqueId & id, bool blocking, std::string * err) { + hipError_t he = hipSetDevice(device_); + if (he != hipSuccess) { + set_err(err, hip_err("hipSetDevice", he)); + return false; + } + he = hipStreamCreateWithFlags(&internal_stream_, hipStreamNonBlocking); + if (he != hipSuccess) { + set_err(err, hip_err("hipStreamCreateWithFlags", he)); + return false; + } + he = hipEventCreateWithFlags(&wait_event_, hipEventDisableTiming); + if (he != hipSuccess) { + set_err(err, hip_err("hipEventCreateWithFlags", he)); + return false; + } + he = hipMalloc(&staging_dev_, kStagingBytes); + if (he != hipSuccess) { + set_err(err, hip_err("hipMalloc(staging)", he)); + return false; + } + he = hipHostMalloc(&staging_host_, kStagingBytes, hipHostMallocDefault); + if (he != hipSuccess) { + set_err(err, hip_err("hipHostMalloc(staging)", he)); + return false; + } + + ncclUniqueId nid; + std::memcpy(&nid, id.data(), sizeof(nid)); + ncclConfig_t cfg = NCCL_CONFIG_INITIALIZER; + cfg.blocking = blocking ? 1 : 0; + + ncclResult_t r = ncclCommInitRankConfig(&comm_, size_, nid, rank_, &cfg); + if (!nccl_ok(r)) { + set_err(err, nccl_err("ncclCommInitRankConfig", r)); + comm_ = nullptr; + return false; + } + if (!blocking) { + // Poll until the communicator leaves InProgress or the deadline + // hits; on timeout abort so peers unblock too. + const uint64_t deadline = now_us() + (uint64_t)timeout_ms_ * 1000ull; + for (;;) { + ncclResult_t state = ncclSuccess; + r = ncclCommGetAsyncError(comm_, &state); + if (r != ncclSuccess) { + set_err(err, nccl_err("ncclCommGetAsyncError during init", r)); + abort(); + return false; + } + if (state == ncclSuccess) break; + if (state != ncclInProgress) { + set_err(err, nccl_err("RCCL communicator init failed", state)); + abort(); + return false; + } + if (now_us() > deadline) { + set_err(err, "RCCL communicator init timed out after " + + std::to_string(timeout_ms_) + " ms (rank " + + std::to_string(rank_) + "/" + std::to_string(size_) + ")"); + abort(); + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + } + ready_ = true; + return true; + } + + int rank() const override { return rank_; } + int size() const override { return size_; } + + bool allreduce_sum_f32(void * dev_ptr, size_t n, DeviceStream stream, std::string * err) override { + if (!check_ready(err)) return false; + if (n == 0) return true; + (void)hipSetDevice(device_); + const ncclResult_t r = ncclAllReduce(dev_ptr, dev_ptr, n, ncclFloat32, ncclSum, comm_, + (hipStream_t)stream); + if (!nccl_ok(r)) { + set_err(err, nccl_err("ncclAllReduce(f32)", r)); + return false; + } + std::lock_guard lk(mu_); + stats_.allreduce_calls++; + stats_.allreduce_bytes += n * sizeof(float); + return true; + } + + bool allreduce_sum_bf16_compressed(void * dev_f32, size_t n, void * scratch_bf16, + DeviceStream stream, std::string * err) override { + if (!check_ready(err)) return false; + if (n == 0) return true; + if (scratch_bf16 == nullptr) { + set_err(err, "allreduce_sum_bf16_compressed needs a device scratch buffer"); + return false; + } + (void)hipSetDevice(device_); + if (!cluster_kernel_f32_to_bf16((const float *)dev_f32, (uint16_t *)scratch_bf16, n, stream)) { + set_err(err, hip_err("f32->bf16 conversion launch", hipGetLastError())); + return false; + } + const ncclResult_t r = ncclAllReduce(scratch_bf16, scratch_bf16, n, ncclBfloat16, ncclSum, + comm_, (hipStream_t)stream); + if (!nccl_ok(r)) { + set_err(err, nccl_err("ncclAllReduce(bf16)", r)); + return false; + } + if (!cluster_kernel_bf16_to_f32((const uint16_t *)scratch_bf16, (float *)dev_f32, n, stream)) { + set_err(err, hip_err("bf16->f32 conversion launch", hipGetLastError())); + return false; + } + std::lock_guard lk(mu_); + stats_.allreduce_calls++; + stats_.allreduce_bytes += n * 2; + return true; + } + + bool allreduce_sum_q8_exchange(void * dev_f32, size_t n, void * scratch, + DeviceStream stream, std::string * err) override { + if (!check_ready(err)) return false; + if (n == 0) return true; + if (scratch == nullptr || n % kQ8ExchangeBlock != 0) { + set_err(err, "allreduce_sum_q8_exchange needs a device scratch and n % 32 == 0"); + return false; + } + (void)hipSetDevice(device_); + const size_t slot = q8_exchange_slot_bytes(n); + uint8_t * base = (uint8_t *)scratch; + if (!cluster_kernel_f32_to_q8((const float *)dev_f32, base + (size_t)rank_ * slot, n, stream)) { + set_err(err, hip_err("f32->q8 launch", hipGetLastError())); + return false; + } + const ncclResult_t r = ncclAllGather(base + (size_t)rank_ * slot, base, slot, ncclInt8, + comm_, (hipStream_t)stream); + if (!nccl_ok(r)) { + set_err(err, nccl_err("ncclAllGather(q8)", r)); + return false; + } + if (!cluster_kernel_q8_sum(base, size_, slot, n, (float *)dev_f32, stream)) { + set_err(err, hip_err("q8 sum launch", hipGetLastError())); + return false; + } + std::lock_guard lk(mu_); + stats_.allreduce_calls++; + stats_.allreduce_bytes += slot; + return true; + } + + bool broadcast_bytes(void * dev_buf, size_t bytes, int root, DeviceStream stream, + std::string * err) override { + if (!check_ready(err)) return false; + if (root < 0 || root >= size_) { + set_err(err, "broadcast root " + std::to_string(root) + " out of range"); + return false; + } + if (bytes == 0) return true; + (void)hipSetDevice(device_); + const ncclResult_t r = ncclBroadcast(dev_buf, dev_buf, bytes, ncclUint8, root, comm_, + (hipStream_t)stream); + if (!nccl_ok(r)) { + set_err(err, nccl_err("ncclBroadcast", r)); + return false; + } + std::lock_guard lk(mu_); + stats_.broadcast_calls++; + stats_.broadcast_bytes += bytes; + return true; + } + + bool broadcast_host_i32(int32_t * host_buf, size_t count, int root, std::string * err) override { + if (!check_ready(err)) return false; + if (root < 0 || root >= size_) { + set_err(err, "broadcast root " + std::to_string(root) + " out of range"); + return false; + } + (void)hipSetDevice(device_); + // Serialize users of the internal staging buffers. + std::lock_guard slk(staging_mu_); + const size_t max_per_chunk = kStagingBytes / sizeof(int32_t); + size_t done = 0; + while (done < count) { + const size_t n = std::min(max_per_chunk, count - done); + const size_t bytes = n * sizeof(int32_t); + if (rank_ == root) { + std::memcpy(staging_host_, host_buf + done, bytes); + hipError_t he = hipMemcpyAsync(staging_dev_, staging_host_, bytes, + hipMemcpyHostToDevice, internal_stream_); + if (he != hipSuccess) { + set_err(err, hip_err("hipMemcpyAsync(H2D staging)", he)); + return false; + } + } + const ncclResult_t r = ncclBroadcast(staging_dev_, staging_dev_, bytes, ncclUint8, root, + comm_, internal_stream_); + if (!nccl_ok(r)) { + set_err(err, nccl_err("ncclBroadcast(host_i32)", r)); + return false; + } + hipError_t he = hipMemcpyAsync(staging_host_, staging_dev_, bytes, + hipMemcpyDeviceToHost, internal_stream_); + if (he != hipSuccess) { + set_err(err, hip_err("hipMemcpyAsync(D2H staging)", he)); + return false; + } + if (!wait_stream_impl(internal_stream_, timeout_ms_, /*count_allreduce=*/false, err)) { + return false; + } + std::memcpy(host_buf + done, staging_host_, bytes); + done += n; + } + std::lock_guard lk(mu_); + stats_.broadcast_calls++; + stats_.broadcast_bytes += count * sizeof(int32_t); + return true; + } + + bool barrier(std::string * err) override { + if (!check_ready(err)) return false; + const uint64_t t0 = now_us(); + (void)hipSetDevice(device_); + std::lock_guard slk(staging_mu_); + // Tiny all-reduce on the internal stream + host sync = barrier. + const ncclResult_t r = ncclAllReduce(staging_dev_, staging_dev_, 1, ncclFloat32, ncclSum, + comm_, internal_stream_); + if (!nccl_ok(r)) { + set_err(err, nccl_err("ncclAllReduce(barrier)", r)); + return false; + } + if (!wait_stream_impl(internal_stream_, timeout_ms_, /*count_allreduce=*/false, err)) { + return false; + } + std::lock_guard lk(mu_); + stats_.barrier_calls++; + stats_.barrier_us += now_us() - t0; + return true; + } + + bool wait_stream(DeviceStream stream, uint32_t deadline_ms, std::string * err) override { + if (!check_ready(err)) return false; + return wait_stream_impl((hipStream_t)stream, deadline_ms, /*count_allreduce=*/true, err); + } + + void abort() override { + std::lock_guard lk(mu_); + if (aborted_) return; + aborted_ = true; + stats_.aborts++; + if (comm_ != nullptr) { + (void)ncclCommAbort(comm_); + comm_ = nullptr; + } + ready_ = false; + } + + bool async_error(std::string * err) const override { + std::lock_guard lk(mu_); + if (comm_ == nullptr) { + if (aborted_) set_err(err, "RCCL communicator aborted"); + return aborted_; + } + ncclResult_t state = ncclSuccess; + const ncclResult_t r = ncclCommGetAsyncError(comm_, &state); + if (r != ncclSuccess) { + set_err(err, nccl_err("ncclCommGetAsyncError", r)); + return true; + } + if (state != ncclSuccess && state != ncclInProgress) { + set_err(err, nccl_err("RCCL asynchronous error", state)); + return true; + } + return false; + } + + ClusterCommStats stats() const override { + std::lock_guard lk(mu_); + return stats_; + } + + void reset_stats() override { + std::lock_guard lk(mu_); + stats_ = ClusterCommStats{}; + } + + const char * backend_name() const override { return "rccl"; } + +private: + static constexpr size_t kStagingBytes = 64 * 1024; + + bool check_ready(std::string * err) const { + std::lock_guard lk(mu_); + if (aborted_) { + set_err(err, "RCCL communicator aborted"); + return false; + } + if (!ready_ || comm_ == nullptr) { + set_err(err, "RCCL communicator not initialized"); + return false; + } + return true; + } + + // Record an event on `stream` and poll it until completion or deadline. + // On timeout: abort() so peers blocked in the same collective fail too. + bool wait_stream_impl(hipStream_t stream, uint32_t deadline_ms, bool count_allreduce, + std::string * err) { + const uint64_t t0 = now_us(); + const uint64_t deadline = t0 + (uint64_t)deadline_ms * 1000ull; + hipEvent_t ev = nullptr; + { + std::lock_guard lk(event_mu_); + ev = wait_event_; + const hipError_t he = hipEventRecord(ev, stream); + if (he != hipSuccess) { + set_err(err, hip_err("hipEventRecord", he)); + return false; + } + // Poll under event_mu_ so a concurrent wait cannot re-record the + // single event underneath us. + uint32_t spins = 0; + for (;;) { + const hipError_t q = hipEventQuery(ev); + if (q == hipSuccess) break; + if (q != hipErrorNotReady) { + set_err(err, hip_err("hipEventQuery", q)); + return false; + } + std::string aerr; + if (async_error(&aerr)) { + set_err(err, "collective failed while waiting: " + aerr); + abort(); + return false; + } + if (now_us() > deadline) { + set_err(err, "wait_stream timed out after " + std::to_string(deadline_ms) + + " ms on rank " + std::to_string(rank_) + "; aborting communicator"); + abort(); + return false; + } + // Busy-poll briefly (sub-100 us collectives), then back off. + if (++spins < 2000) { + std::this_thread::yield(); + } else { + std::this_thread::sleep_for(std::chrono::microseconds(50)); + } + } + } + if (count_allreduce) { + std::lock_guard lk(mu_); + stats_.allreduce_wait_us += now_us() - t0; + } + return true; + } + + const int rank_; + const int size_; + const int device_; + const uint32_t timeout_ms_; + + ncclComm_t comm_ = nullptr; + hipStream_t internal_stream_ = nullptr; + hipEvent_t wait_event_ = nullptr; + void * staging_dev_ = nullptr; + void * staging_host_ = nullptr; + + mutable std::mutex mu_; // stats_, aborted_, ready_, comm_ lifetime + std::mutex staging_mu_; // staging buffers + internal stream users + std::mutex event_mu_; // wait_event_ + ClusterCommStats stats_; + bool aborted_ = false; + bool ready_ = false; +}; + +// ─── Self-test (device path) ──────────────────────────────────────────── + +struct SelftestBuffers { + int device = 0; + hipStream_t stream = nullptr; + hipEvent_t e0 = nullptr; + hipEvent_t e1 = nullptr; + float * dev = nullptr; + uint16_t * scratch = nullptr; + float * host_src = nullptr; + std::vector host_out; + + ~SelftestBuffers() { + (void)hipSetDevice(device); + if (e0) (void)hipEventDestroy(e0); + if (e1) (void)hipEventDestroy(e1); + if (dev) (void)hipFree(dev); + if (scratch) (void)hipFree(scratch); + if (host_src) (void)hipHostFree(host_src); + if (stream) (void)hipStreamDestroy(stream); + } + + bool init(int dev_id, size_t max_n, float fill, std::string * err) { + device = dev_id; + hipError_t he = hipSetDevice(device); + if (he != hipSuccess) { set_err(err, hip_err("hipSetDevice", he)); return false; } + he = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + if (he != hipSuccess) { set_err(err, hip_err("hipStreamCreate", he)); return false; } + he = hipEventCreate(&e0); + if (he != hipSuccess) { set_err(err, hip_err("hipEventCreate", he)); return false; } + he = hipEventCreate(&e1); + if (he != hipSuccess) { set_err(err, hip_err("hipEventCreate", he)); return false; } + he = hipMalloc((void **)&dev, max_n * sizeof(float)); + if (he != hipSuccess) { set_err(err, hip_err("hipMalloc(selftest)", he)); return false; } + he = hipMalloc((void **)&scratch, max_n * sizeof(uint16_t)); + if (he != hipSuccess) { set_err(err, hip_err("hipMalloc(selftest scratch)", he)); return false; } + he = hipHostMalloc((void **)&host_src, max_n * sizeof(float), hipHostMallocDefault); + if (he != hipSuccess) { set_err(err, hip_err("hipHostMalloc(selftest)", he)); return false; } + for (size_t i = 0; i < max_n; ++i) host_src[i] = fill; + host_out.resize(max_n); + return true; + } + + // Upload fresh inputs, run `op`, and measure its device time in µs. + template + bool timed(size_t n, Op && op, double & us, IClusterComm & comm, uint32_t timeout_ms, + std::string * err) { + hipError_t he = hipMemcpyAsync(dev, host_src, n * sizeof(float), hipMemcpyHostToDevice, stream); + if (he != hipSuccess) { set_err(err, hip_err("hipMemcpyAsync(H2D)", he)); return false; } + he = hipEventRecord(e0, stream); + if (he != hipSuccess) { set_err(err, hip_err("hipEventRecord", he)); return false; } + if (!op()) return false; + he = hipEventRecord(e1, stream); + if (he != hipSuccess) { set_err(err, hip_err("hipEventRecord", he)); return false; } + if (!comm.wait_stream(stream, timeout_ms, err)) return false; + float ms = 0.0f; + he = hipEventElapsedTime(&ms, e0, e1); + if (he != hipSuccess) { set_err(err, hip_err("hipEventElapsedTime", he)); return false; } + us = (double)ms * 1000.0; + return true; + } + + bool verify(size_t n, float expected, const char * what, std::string * err) { + const hipError_t he = hipMemcpy(host_out.data(), dev, n * sizeof(float), hipMemcpyDeviceToHost); + if (he != hipSuccess) { set_err(err, hip_err("hipMemcpy(D2H verify)", he)); return false; } + for (size_t i = 0; i < n; ++i) { + if (host_out[i] != expected) { + set_err(err, std::string(what) + ": element " + std::to_string(i) + " = " + + std::to_string(host_out[i]) + ", expected " + std::to_string(expected)); + return false; + } + } + return true; + } +}; + +} // namespace + +bool rccl_get_unique_id(RcclUniqueId & out, std::string * err) { + ncclUniqueId id; + const ncclResult_t r = ncclGetUniqueId(&id); + if (r != ncclSuccess) { + set_err(err, nccl_err("ncclGetUniqueId", r)); + return false; + } + std::memcpy(out.data(), &id, sizeof(id)); + return true; +} + +std::unique_ptr create_rccl_cluster_comm(const RcclCommInit & init, std::string * err) { + if (init.size < 1 || init.rank < 0 || init.rank >= init.size) { + set_err(err, "bad RCCL rank/size (" + std::to_string(init.rank) + "/" + + std::to_string(init.size) + ")"); + return nullptr; + } + std::unique_ptr comm(new RcclComm(init.rank, init.size, init.device, + init.timeout_ms == 0 ? 30000u : init.timeout_ms)); + if (!comm->init(init.unique_id, init.blocking, err)) { + return nullptr; + } + return std::unique_ptr(comm.release()); +} + +std::string rccl_version_string() { + int v = 0; + if (ncclGetVersion(&v) != ncclSuccess || v <= 0) return "n/a"; + // NCCL >= 2.9 encodes MAJOR*10000 + MINOR*100 + PATCH. + const int major = v / 10000; + const int minor = (v % 10000) / 100; + const int patch = v % 100; + return std::to_string(major) + "." + std::to_string(minor) + "." + std::to_string(patch); +} + +#else // !LUCE_CLUSTER_RCCL ────────────────────────────────────────── + +bool rccl_get_unique_id(RcclUniqueId & out, std::string * err) { + out.fill(0); + set_err(err, kNoRcclMsg); + return false; +} + +std::unique_ptr create_rccl_cluster_comm(const RcclCommInit & init, std::string * err) { + (void)init; + set_err(err, kNoRcclMsg); + return nullptr; +} + +std::string rccl_version_string() { + return "n/a"; +} + +#endif // LUCE_CLUSTER_RCCL + +// ─── Self-test ────────────────────────────────────────────────────────── + +namespace { + +struct LatencyStats { + std::vector samples; + + double percentile(double p) const { + if (samples.empty()) return 0.0; + std::vector s = samples; + std::sort(s.begin(), s.end()); + const double idx = p / 100.0 * (double)(s.size() - 1); + const size_t lo = (size_t)std::floor(idx); + const size_t hi = std::min(lo + 1, s.size() - 1); + const double frac = idx - (double)lo; + return s[lo] + (s[hi] - s[lo]) * frac; + } + + void print(const char * label, size_t n_floats, size_t bytes) const { + std::printf("[cluster-selftest] %-14s n=%zu bytes=%zu iters=%zu p50=%.1f us p90=%.1f us " + "p99=%.1f us max=%.1f us\n", + label, n_floats, bytes, samples.size(), percentile(50.0), percentile(90.0), + percentile(99.0), percentile(100.0)); + } +}; + +} // namespace + +bool run_cluster_selftest(IClusterComm & comm, int device, int iters, size_t small_n, size_t large_n, + std::string * err) { + if (iters < 1) iters = 1; + if (small_n == 0) small_n = 16384; + if (large_n == 0) large_n = small_n; + const int size = comm.size(); + const float fill = (float)(comm.rank() + 1); + const float expected = (float)((size * (size + 1)) / 2); + constexpr int kLargeIters = 43; // one per DeepSeek4 MoE layer + + std::printf("[cluster-selftest] rank %d/%d backend=%s rccl=%s device=%d small_n=%zu large_n=%zu " + "iters=%d\n", + comm.rank(), size, comm.backend_name(), rccl_version_string().c_str(), device, + small_n, large_n, iters); + +#ifdef LUCE_CLUSTER_RCCL + SelftestBuffers buf; + const size_t max_n = std::max(small_n, large_n); + if (!buf.init(device, max_n, fill, err)) return false; + const uint32_t timeout_ms = 30000; + + LatencyStats small_lat; + const int check_every = std::max(1, iters / 10); + for (int i = 0; i < iters; ++i) { + double us = 0.0; + auto op = [&]() { + return comm.allreduce_sum_f32(buf.dev, small_n, buf.stream, err); + }; + if (!buf.timed(small_n, op, us, comm, timeout_ms, err)) return false; + small_lat.samples.push_back(us); + if (i % check_every == 0 || i + 1 == iters) { + if (!buf.verify(small_n, expected, "small f32 all-reduce", err)) return false; + } + } + small_lat.print("f32 small", small_n, small_n * sizeof(float)); + + LatencyStats large_f32; + LatencyStats large_bf16; + for (int i = 0; i < kLargeIters; ++i) { + const bool use_bf16 = (i % 2) == 1; + double us = 0.0; + auto op = [&]() { + return use_bf16 + ? comm.allreduce_sum_bf16_compressed(buf.dev, large_n, buf.scratch, buf.stream, err) + : comm.allreduce_sum_f32(buf.dev, large_n, buf.stream, err); + }; + if (!buf.timed(large_n, op, us, comm, timeout_ms, err)) return false; + (use_bf16 ? large_bf16 : large_f32).samples.push_back(us); + // rank+1 sums are small integers: exactly representable in bf16. + if (!buf.verify(large_n, expected, use_bf16 ? "large bf16 all-reduce" : "large f32 all-reduce", + err)) { + return false; + } + } + large_f32.print("f32 large", large_n, large_n * sizeof(float)); + large_bf16.print("bf16 large", large_n, large_n * 2); + + // Host broadcast of a few decision-sized words from rank 0. + int32_t words[8]; + for (int i = 0; i < 8; ++i) words[i] = comm.rank() == 0 ? 1000 + i : -1; + if (!comm.broadcast_host_i32(words, 8, 0, err)) return false; + for (int i = 0; i < 8; ++i) { + if (words[i] != 1000 + i) { + set_err(err, "broadcast_host_i32 mismatch at " + std::to_string(i) + ": " + + std::to_string(words[i])); + return false; + } + } + if (!comm.barrier(err)) return false; + + std::string aerr; + if (comm.async_error(&aerr)) { + set_err(err, "async error after self-test: " + aerr); + return false; + } + const ClusterCommStats st = comm.stats(); + std::printf("[cluster-selftest] rank %d OK: allreduce_calls=%llu allreduce_bytes=%llu " + "allreduce_wait_us=%llu broadcast_calls=%llu barrier_us=%llu\n", + comm.rank(), (unsigned long long)st.allreduce_calls, + (unsigned long long)st.allreduce_bytes, (unsigned long long)st.allreduce_wait_us, + (unsigned long long)st.broadcast_calls, (unsigned long long)st.barrier_us); + return true; +#else + // Without RCCL only the loopback communicator exists; it never touches + // the buffers, so host memory stands in for device memory. + (void)device; + if (std::string(comm.backend_name()) != "loopback") { + set_err(err, kNoRcclMsg); + return false; + } + std::vector host(std::max(small_n, large_n), fill); + std::vector scratch(host.size()); + LatencyStats small_lat; + for (int i = 0; i < iters; ++i) { + const uint64_t t0 = now_us(); + if (!comm.allreduce_sum_f32(host.data(), small_n, nullptr, err)) return false; + if (!comm.wait_stream(nullptr, 30000, err)) return false; + small_lat.samples.push_back((double)(now_us() - t0)); + } + small_lat.print("f32 small", small_n, small_n * sizeof(float)); + for (int i = 0; i < kLargeIters; ++i) { + if (!comm.allreduce_sum_bf16_compressed(host.data(), large_n, scratch.data(), nullptr, err)) { + return false; + } + } + for (size_t i = 0; i < small_n; ++i) { + if (host[i] != expected) { + set_err(err, "loopback self-test value mismatch at " + std::to_string(i)); + return false; + } + } + if (!comm.barrier(err)) return false; + std::printf("[cluster-selftest] rank %d OK (loopback)\n", comm.rank()); + return true; +#endif +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_comm.h b/server/src/cluster/cluster_comm.h new file mode 100644 index 000000000..ff6f298ab --- /dev/null +++ b/server/src/cluster/cluster_comm.h @@ -0,0 +1,160 @@ +// cluster_comm.h - collective communication between cluster ranks. +// +// Thin interface over RCCL (ROCm's NCCL) so the model code never includes +// rccl.h directly and unit tests can substitute a loopback or mock. All +// device-side collectives are enqueued on the caller-supplied HIP stream (the +// ggml backend's own stream, see ggml_backend_cuda_get_stream) so they are +// stream-ordered with the kernels that produce and consume the tensors; no +// host synchronization is needed around a call. +// +// Bootstrap: rank 0 obtains a unique id (rccl_get_unique_id), ships it to the +// workers over the control channel (WelcomeMsg), and every rank then calls +// create_rccl_cluster_comm. The communicator is created with the +// non-blocking config so a hung collective can be aborted from a watchdog +// (wait_stream deadline -> abort()). +// +// gfx1151 has no GPUDirect RDMA; RCCL bounces through host memory +// (NCCL_NET_GDR_LEVEL=0). On Strix Halo that host memory is the same DRAM +// as GPU memory, so the bounce is a memcpy, not a PCIe hop. + +#pragma once + +#include +#include +#include +#include +#include + +namespace luce::cluster { + +inline constexpr size_t kRcclUniqueIdSize = 128; // == NCCL_UNIQUE_ID_BYTES +using RcclUniqueId = std::array; + +// Device memory in use on `device` right now (total - free), in bytes, or 0 +// when the backend cannot report it. Sampled at request boundaries to fill +// RequestReportMsg::peak_device_bytes, which is therefore a high-water mark +// over sampled request ends, not a continuous peak. +uint64_t cluster_device_bytes_in_use(int device); + +struct ClusterCommStats { + uint64_t allreduce_calls = 0; + uint64_t allreduce_bytes = 0; // payload bytes handed to RCCL (f32 or bf16) + uint64_t allreduce_wait_us = 0; // time spent in wait_stream() attributable to collectives + uint64_t broadcast_calls = 0; + uint64_t broadcast_bytes = 0; + uint64_t barrier_calls = 0; + uint64_t barrier_us = 0; + uint64_t aborts = 0; +}; + +// Opaque stream handle: hipStream_t on HIP builds, cudaStream_t on CUDA +// builds, ignored by the loopback implementation. +using DeviceStream = void *; + +// The 8-bit exchange (allreduce_sum_q8_exchange): blocks of 32 values with one +// f32 scale. A rank's slot holds its n int8 values, then its n / 32 scales. +constexpr size_t kQ8ExchangeBlock = 32; +constexpr size_t q8_exchange_slot_bytes(size_t n) { + return n + (n / kQ8ExchangeBlock) * sizeof(float); +} + +class IClusterComm { +public: + virtual ~IClusterComm() = default; + + virtual int rank() const = 0; + virtual int size() const = 0; + + // In-place sum of n float32 values at dev_ptr across all ranks. Enqueued + // on `stream`; returns as soon as the collective is enqueued (or false on + // enqueue failure). Result is bit-identical on every rank. + virtual bool allreduce_sum_f32(void * dev_ptr, size_t n, DeviceStream stream, + std::string * err) = 0; + + // Bandwidth-saving variant for large prefill partials: converts dev_f32 + // to bf16 in `scratch_bf16` (device, >= n * 2 bytes), all-reduces the + // bf16 buffer with f32 accumulation semantics as in + // ggml_backend_cuda_comm_allreduce_tensor, and writes the f32 sum back to + // dev_f32. Same stream-ordering contract as allreduce_sum_f32. + virtual bool allreduce_sum_bf16_compressed(void * dev_f32, size_t n, void * scratch_bf16, + DeviceStream stream, std::string * err) = 0; + + // About 0.56x the bytes of the bf16 path, for large prefill partials: + // every rank quantizes its f32 partial to 8-bit blocks in its slot of + // `scratch` (device, >= size() * q8_exchange_slot_bytes(n) bytes), + // all-gathers the slots, and writes the sum of all ranks' dequantized + // values, in rank order, back to dev_f32: the same bits on every rank. + // n must be a multiple of kQ8ExchangeBlock. False when unsupported. + virtual bool allreduce_sum_q8_exchange(void * dev_f32, size_t n, void * scratch, + DeviceStream stream, std::string * err) { + (void)dev_f32; (void)n; (void)scratch; (void)stream; + if (err) *err = "this communicator has no 8-bit exchange"; + return false; + } + + // Broadcast `bytes` bytes at `buf` (device memory) from `root`. + virtual bool broadcast_bytes(void * dev_buf, size_t bytes, int root, DeviceStream stream, + std::string * err) = 0; + + // Host-side convenience: broadcast small host buffers (decisions, draft + // tokens) via a pinned staging buffer and a dedicated internal stream; + // blocks until the data is available on the host on every rank. Used + // only when the control channel is not preferred for a message. + virtual bool broadcast_host_i32(int32_t * host_buf, size_t count, int root, + std::string * err) = 0; + + // Cluster-wide barrier (host blocking). + virtual bool barrier(std::string * err) = 0; + + // Block the host until all work enqueued on `stream` so far has + // completed, or until deadline_ms elapses. On timeout returns false, + // records an abort and calls abort() so peers unblock with an error + // instead of hanging. + virtual bool wait_stream(DeviceStream stream, uint32_t deadline_ms, std::string * err) = 0; + + // Tear down the communicator abruptly (ncclCommAbort). Idempotent. + virtual void abort() = 0; + + // True if a collective failed asynchronously (ncclCommGetAsyncError). + virtual bool async_error(std::string * err) const = 0; + + virtual ClusterCommStats stats() const = 0; + virtual void reset_stats() = 0; + + virtual const char * backend_name() const = 0; // "rccl" | "loopback" | "mock" +}; + +struct RcclCommInit { + int rank = -1; + int size = 0; + RcclUniqueId unique_id{}; + int device = 0; // local HIP device ordinal + uint32_t timeout_ms = 30000; // default deadline for wait_stream + bool blocking = false; // false -> ncclCommInitRankConfig non-blocking +}; + +// Rank 0 only: obtain the unique id to distribute in WelcomeMsg. +// Compiled only when LUCE_CLUSTER_RCCL is defined; otherwise returns +// false with err = "RCCL support not compiled in". +bool rccl_get_unique_id(RcclUniqueId & out, std::string * err); + +// Create the RCCL-backed communicator for this rank (all ranks must call it +// with the same unique id). nullptr + err on failure. +std::unique_ptr create_rccl_cluster_comm(const RcclCommInit & init, std::string * err); + +// size-1 communicator: every collective is a no-op that succeeds. Used for +// single-node runs of cluster-enabled code paths and for tests. +std::unique_ptr create_loopback_cluster_comm(int rank = 0, int size = 1); + +// Human-readable RCCL version string ("2.27.7" etc.) or "n/a". +std::string rccl_version_string(); + +// Runs the startup self-test used by --cluster-selftest: `iters` all-reduces +// of `small_n` floats plus 43 all-reduces of `large_n` floats on device +// buffers filled with rank-dependent values, checks the sums against the +// host-computed expectation, and reports latency percentiles to stdout. +// Returns false on any mismatch or timeout. +bool run_cluster_selftest(IClusterComm & comm, int device, int iters, size_t small_n, size_t large_n, + std::string * err); + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_comm_kernels.cu b/server/src/cluster/cluster_comm_kernels.cu new file mode 100644 index 000000000..f7eae31fb --- /dev/null +++ b/server/src/cluster/cluster_comm_kernels.cu @@ -0,0 +1,110 @@ +// cluster_comm_kernels.cu - device kernels for the RCCL cluster communicator. +// +// Compiled as HIP (LANGUAGE HIP in server/CMakeLists.txt) only when +// LUCE_CLUSTER=ON. Provides the f32 <-> bf16 conversions used by +// allreduce_sum_bf16_compressed (bf16 is written by hand, round-to-nearest- +// even on the raw bits, so this file needs neither hip_bf16.h nor ggml) and +// the 8-bit blocks of allreduce_sum_q8_exchange. + +#include + +#include +#include + +namespace luce::cluster { + +namespace { + +constexpr int kBlock = 256; + +__device__ __forceinline__ uint16_t f32_to_bf16_bits(float f) { + uint32_t u = __float_as_uint(f); + if ((u & 0x7f800000u) == 0x7f800000u) { + // Inf/NaN: truncate, keep NaN quiet. + uint16_t r = (uint16_t)(u >> 16); + if (u & 0x007fffffu) r |= 0x0040u; + return r; + } + const uint32_t lsb = (u >> 16) & 1u; + u += 0x7fffu + lsb; + return (uint16_t)(u >> 16); +} + +__global__ void k_f32_to_bf16(const float * __restrict__ in, uint16_t * __restrict__ out, size_t n) { + const size_t i = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; + if (i < n) out[i] = f32_to_bf16_bits(in[i]); +} + +__global__ void k_bf16_to_f32(const uint16_t * __restrict__ in, float * __restrict__ out, size_t n) { + const size_t i = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; + if (i < n) out[i] = __uint_as_float((uint32_t)in[i] << 16); +} + +// One thread per block of 32 values: the scale is the block's absolute +// maximum over 127, the values round to nearest. +__global__ void k_f32_to_q8(const float * __restrict__ in, int8_t * __restrict__ q, + float * __restrict__ d, size_t nb) { + const size_t b = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; + if (b >= nb) return; + const float * x = in + b * 32; + float amax = 0.0f; + for (int i = 0; i < 32; ++i) amax = fmaxf(amax, fabsf(x[i])); + const float scale = amax / 127.0f; + const float inv = scale > 0.0f ? 1.0f / scale : 0.0f; + int8_t * qb = q + b * 32; + for (int i = 0; i < 32; ++i) qb[i] = (int8_t)__float2int_rn(x[i] * inv); + d[b] = scale; +} + +// out[i] = the ranks' dequantized values summed in rank order. +__global__ void k_q8_sum(const uint8_t * __restrict__ slots, int ranks, size_t slot_bytes, + size_t n, float * __restrict__ out) { + const size_t i = (size_t)blockIdx.x * (size_t)blockDim.x + (size_t)threadIdx.x; + if (i >= n) return; + float acc = 0.0f; + for (int r = 0; r < ranks; ++r) { + const uint8_t * slot = slots + (size_t)r * slot_bytes; + const float scale = ((const float *)(slot + n))[i / 32]; + acc += (float)((const int8_t *)slot)[i] * scale; + } + out[i] = acc; +} + +unsigned grid_for(size_t n) { + const size_t blocks = (n + (size_t)kBlock - 1) / (size_t)kBlock; + return blocks > 0x7fffffffull ? 0x7fffffffu : (unsigned)blocks; +} + +} // namespace + +// Declared in cluster_comm.cpp under LUCE_CLUSTER_RCCL. +bool cluster_kernel_f32_to_bf16(const float * in, uint16_t * out, size_t n, void * stream) { + if (n == 0) return true; + k_f32_to_bf16<<>>(in, out, n); + return hipGetLastError() == hipSuccess; +} + +bool cluster_kernel_bf16_to_f32(const uint16_t * in, float * out, size_t n, void * stream) { + if (n == 0) return true; + k_bf16_to_f32<<>>(in, out, n); + return hipGetLastError() == hipSuccess; +} + +// `slot`: n int8 values, then n / 32 f32 scales (q8_exchange_slot_bytes). +bool cluster_kernel_f32_to_q8(const float * in, void * slot, size_t n, void * stream) { + if (n == 0) return true; + const size_t nb = n / 32; + k_f32_to_q8<<>>( + in, (int8_t *)slot, (float *)((uint8_t *)slot + n), nb); + return hipGetLastError() == hipSuccess; +} + +bool cluster_kernel_q8_sum(const void * slots, int ranks, size_t slot_bytes, size_t n, float * out, + void * stream) { + if (n == 0) return true; + k_q8_sum<<>>( + (const uint8_t *)slots, ranks, slot_bytes, n, out); + return hipGetLastError() == hipSuccess; +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_config.cpp b/server/src/cluster/cluster_config.cpp new file mode 100644 index 000000000..a7fa79148 --- /dev/null +++ b/server/src/cluster/cluster_config.cpp @@ -0,0 +1,278 @@ +// cluster_config.cpp - validation, parsing helpers and RCCL environment +// export for the caller-requested cluster configuration (cluster_config.h). +// +// Pure C++/POSIX: no ggml, HIP or RCCL dependency so the feature gate, the +// CLI parser and the GPU-free unit tests can link this unconditionally. + +#include "cluster/cluster_config.h" + +#include "common/platform_env.h" + +#include +#include +#include +#include + +namespace luce::cluster { + +namespace { + +// Strict decimal integer parse: whole string must be consumed. +bool parse_int_strict(const std::string & text, int & out) { + if (text.empty()) return false; + errno = 0; + char * end = nullptr; + const long v = std::strtol(text.c_str(), &end, 10); + if (errno == ERANGE || end == text.c_str() || *end != '\0') return false; + if (v < -2147483647L - 1L || v > 2147483647L) return false; + out = (int)v; + return true; +} + +void set_err(std::string * err, const std::string & msg) { + if (err) *err = msg; +} + +bool ends_with(const std::string & s, const char * suffix) { + const size_t n = std::strlen(suffix); + return s.size() >= n && s.compare(s.size() - n, n, suffix) == 0; +} + +// True when the variable is set to a non-empty value other than "0". +bool env_flag_set(const char * name) { + const char * v = std::getenv(name); + if (!v || v[0] == '\0') return false; + return std::strcmp(v, "0") != 0; +} + +// setenv only when absent; returns 1 when a variable was written. +int export_if_absent(const char * name, const std::string & value) { + if (std::getenv(name) != nullptr) return 0; + if (value.empty()) return 0; + return common::set_environment_variable(name, value.c_str(), false) == 0 ? 1 : 0; +} + +} // namespace + +// ─── ClusterConfig ────────────────────────────────────────────────────── + +std::string ClusterConfig::validate() const { + if (!enabled()) { + return {}; + } + if (size < kClusterMinSize || size > kClusterMaxSize) { + return "--cluster-size must be in [" + std::to_string(kClusterMinSize) + ", " + + std::to_string(kClusterMaxSize) + "], got " + std::to_string(size); + } + if (rank < 0 || rank >= size) { + return "--cluster-rank must be in [0, " + std::to_string(size) + "), got " + + std::to_string(rank); + } + if (head_host.empty()) { + return "--cluster-head is required in cluster mode"; + } + if (head_port <= 0 || head_port > 65535) { + return "--cluster-head port must be in [1, 65535], got " + std::to_string(head_port); + } + if (gid_index < -1) { + return "--cluster-gid-index must be non-negative, got " + std::to_string(gid_index); + } + if (replicate_hot < 0) { + return "--cluster-replicate-hot must be non-negative, got " + std::to_string(replicate_hot); + } + if (placement_source == PlacementSource::File && placement_file.empty()) { + return "--cluster-expert-placement requires a file path"; + } + if (placement_source != PlacementSource::File && !placement_file.empty()) { + return "internal: placement file set without a File placement source"; + } + if (timeout_ms == 0) { + return "--cluster-timeout-ms must be positive"; + } + return {}; +} + +// ─── Parsing helpers ──────────────────────────────────────────────────── + +bool parse_host_port(const std::string & text, std::string & host, int & port, std::string * err) { + if (text.empty()) { + set_err(err, "expected host:port, got an empty string"); + return false; + } + + std::string host_part; + std::string port_part; + bool has_port = false; + if (text[0] == '[') { + // Bracketed IPv6 literal: [addr]:port or [addr] + const size_t close = text.find(']'); + if (close == std::string::npos) { + set_err(err, "unterminated '[' in host:port '" + text + "'"); + return false; + } + host_part = text.substr(1, close - 1); + if (close + 1 < text.size()) { + if (text[close + 1] != ':') { + set_err(err, "expected ':' after ']' in '" + text + "'"); + return false; + } + port_part = text.substr(close + 2); + has_port = true; + } + } else { + const size_t colon = text.rfind(':'); + if (colon != std::string::npos && text.find(':') != colon) { + // More than one ':' without brackets: a bare IPv6 address. Only + // accepted when unambiguous (no port). + host_part = text; + } else if (colon != std::string::npos) { + host_part = text.substr(0, colon); + port_part = text.substr(colon + 1); + has_port = true; + } else { + host_part = text; + } + } + + if (host_part.empty()) { + set_err(err, "empty host in '" + text + "'"); + return false; + } + if (has_port && port_part.empty()) { + set_err(err, "missing port after ':' in '" + text + "'"); + return false; + } + + int p = kClusterDefaultPort; + if (has_port) { + if (!parse_int_strict(port_part, p) || p <= 0 || p > 65535) { + set_err(err, "bad port '" + port_part + "' in '" + text + "' (expected 1..65535)"); + return false; + } + } + + host = host_part; + port = p; + return true; +} + +bool parse_shared_expert_mode(const std::string & text, SharedExpertMode & out, std::string * err) { + if (text == "replicate") { + out = SharedExpertMode::Replicate; + return true; + } + if (text == "shard") { + out = SharedExpertMode::Shard; + return true; + } + set_err(err, "--cluster-shared-expert expects replicate or shard, got '" + text + "'"); + return false; +} + +bool parse_allreduce_dtype(const std::string & text, AllreduceDType & out, std::string * err) { + if (text == "f32") { + out = AllreduceDType::F32; + return true; + } + if (text == "bf16") { + out = AllreduceDType::BF16; + return true; + } + if (text == "auto") { + out = AllreduceDType::Auto; + return true; + } + set_err(err, "--cluster-allreduce-dtype expects f32, bf16, or auto, got '" + text + "'"); + return false; +} + +bool parse_placement_source(const std::string & text, PlacementSource & out, + std::string & file, std::string * err) { + if (text == "uniform") { + out = PlacementSource::Uniform; + file.clear(); + return true; + } + if (text == "balanced") { + out = PlacementSource::Balanced; + file.clear(); + return true; + } + if (ends_with(text, ".json") && text.size() > 5) { + out = PlacementSource::File; + file = text; + return true; + } + set_err(err, "--cluster-expert-placement expects uniform, balanced, or a path ending in " + ".json, got '" + text + "'"); + return false; +} + +// ─── Names ────────────────────────────────────────────────────────────── + +const char * shared_expert_mode_name(SharedExpertMode mode) { + switch (mode) { + case SharedExpertMode::Replicate: return "replicate"; + case SharedExpertMode::Shard: return "shard"; + } + return "unknown"; +} + +const char * allreduce_dtype_name(AllreduceDType dtype) { + switch (dtype) { + case AllreduceDType::F32: return "f32"; + case AllreduceDType::BF16: return "bf16"; + case AllreduceDType::Auto: return "auto"; + } + return "unknown"; +} + +const char * placement_source_name(PlacementSource source) { + switch (source) { + case PlacementSource::Uniform: return "uniform"; + case PlacementSource::Balanced: return "balanced"; + case PlacementSource::File: return "file"; + } + return "unknown"; +} + +// ─── Environment ──────────────────────────────────────────────────────── + +int export_rccl_environment(const ClusterConfig & cfg) { + int n = 0; + + // gfx1151 has no GPUDirect RDMA: RCCL must bounce through host memory. + n += export_if_absent("NCCL_NET_GDR_LEVEL", "0"); + + // Operator-supplied transport hints. + n += export_if_absent("NCCL_SOCKET_IFNAME", cfg.ifname); + n += export_if_absent("NCCL_IB_HCA", cfg.ib_hca); + if (cfg.gid_index >= 0) { + n += export_if_absent("NCCL_IB_GID_INDEX", std::to_string(cfg.gid_index)); + } + + // Hardening defaults measured on the 4-node Strix Halo RoCE v2 cluster. + n += export_if_absent("NCCL_IB_DISABLE", "0"); + n += export_if_absent("NCCL_ASYNC_ERROR_HANDLING", "1"); + n += export_if_absent("NCCL_IB_QPS_PER_CONNECTION", "2"); + n += export_if_absent("NCCL_IB_TIMEOUT", "22"); + n += export_if_absent("NCCL_IB_RETRY_CNT", "7"); + // On the R9700, RCCL runs reliably only with the Simple protocol. + n += export_if_absent("NCCL_PROTO", "Simple"); + + return n; +} + +bool cluster_env_no_ingraph_allreduce() { + return env_flag_set("LUCE_CLUSTER_NO_INGRAPH_ALLREDUCE"); +} + +bool cluster_env_no_graph_capture() { + return env_flag_set("LUCE_CLUSTER_NO_GRAPH_CAPTURE"); +} + +bool cluster_env_trace() { + return env_flag_set("LUCE_CLUSTER_TRACE"); +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_config.h b/server/src/cluster/cluster_config.h new file mode 100644 index 000000000..4e28676c6 --- /dev/null +++ b/server/src/cluster/cluster_config.h @@ -0,0 +1,118 @@ +// cluster_config.h - caller-requested multi-node cluster configuration. +// +// lucebox-halo-cluster: SPMD expert-parallel execution of one DeepSeek4 model +// across N machines (rank 0 = head that serves HTTP, ranks 1..N-1 = workers), +// joined per MoE layer by one RCCL all-reduce over RoCE v2. +// +// This header carries only what the operator asked for on the command line. +// It has no ggml/HIP/RCCL dependency so the feature gate, the CLI parser and +// the GPU-free unit tests can include it unconditionally. Runtime facts +// (communicator handle, resolved placement, peer identities) live in the +// cluster runtime objects, not here. +// +// Policy (variables.md): what the operator chooses ships as CLI flags. The +// measured tuning defaults of cluster mode are environment variables that +// cluster_launch_env() (server/launch_profiles.h) installs, each with a 0 +// kill switch; debug tracing stays in LUCE_CLUSTER_TRACE. + +#pragma once + +#include +#include + +namespace luce::cluster { + +inline constexpr int kClusterMinSize = 2; +inline constexpr int kClusterMaxSize = 8; +inline constexpr int kClusterDefaultPort = 9400; +inline constexpr uint32_t kClusterDefaultTimeoutMs = 30000; + +// How the shared expert (present on every DeepSeek4 layer) is evaluated. +enum class SharedExpertMode { + Replicate, // every rank computes it; all-reduce covers routed experts only + Shard, // n_ff/N slice per rank, summed inside the same all-reduce +}; + +// Wire precision of the per-layer partial-sum reduction. +enum class AllreduceDType { + F32, + BF16, + Auto, // F32 below kAutoBf16ThresholdBytes, BF16 above (prefill chunks) +}; + +inline constexpr uint64_t kAutoBf16ThresholdBytes = 256ull * 1024ull; + +// Expert-to-rank ownership source (--cluster-expert-placement). +enum class PlacementSource { + Uniform, // expert e of every layer owned by rank e % N + Balanced, // greedy LPT on routing hotness (needs LUCE_DS4_HOTNESS_CSV) + File, // explicit JSON produced by scripts/cluster/build_expert_placement.py +}; + +struct ClusterConfig { + // --cluster-rank / --cluster-size. size == 0 means "not a cluster run". + int rank = -1; + int size = 0; + + // --cluster-head : rank 0 binds here, workers connect here. + std::string head_host; + int head_port = kClusterDefaultPort; + + // RCCL transport hints (--cluster-ifname / --cluster-ib-hca / + // --cluster-gid-index). Exported to NCCL_SOCKET_IFNAME, NCCL_IB_HCA and + // NCCL_IB_GID_INDEX unless the environment already sets them. + std::string ifname; + std::string ib_hca; + int gid_index = -1; + + // Expert placement. + PlacementSource placement_source = PlacementSource::Uniform; + std::string placement_file; // when placement_source == File + int replicate_hot = 0; // --cluster-replicate-hot + SharedExpertMode shared_expert = SharedExpertMode::Shard; + + // Reduction and robustness knobs. + AllreduceDType allreduce_dtype = AllreduceDType::Auto; + uint32_t timeout_ms = kClusterDefaultTimeoutMs; // --cluster-timeout-ms + + // --cluster-selftest: bootstrap, run the collective self-test, exit. + bool selftest = false; + + bool enabled() const { return size > 0; } + bool is_head() const { return enabled() && rank == 0; } + bool is_worker() const { return enabled() && rank > 0; } + + // Structural validation independent of model/backend (rank range, size + // range, head endpoint present, replicate_hot >= 0, file present when + // placement_source == File). Returns an empty string when valid. + std::string validate() const; +}; + +// Parsing helpers shared by server_main.cpp and the unit tests. Each returns +// false and fills *err on malformed input; none touches the process +// environment. +bool parse_host_port(const std::string & text, std::string & host, int & port, std::string * err); +bool parse_shared_expert_mode(const std::string & text, SharedExpertMode & out, std::string * err); +bool parse_allreduce_dtype(const std::string & text, AllreduceDType & out, std::string * err); +// "uniform" | "balanced" | +bool parse_placement_source(const std::string & text, PlacementSource & out, + std::string & file, std::string * err); + +const char * shared_expert_mode_name(SharedExpertMode mode); +const char * allreduce_dtype_name(AllreduceDType dtype); +const char * placement_source_name(PlacementSource source); + +// Exports the RCCL/NCCL environment for this node. Existing variables win; +// NCCL_NET_GDR_LEVEL=0 is always set because gfx1151 has no GPUDirect RDMA. +// Also sets the RCCL hardening defaults measured on the 4-node Strix Halo +// cluster (NCCL_IB_DISABLE=0, NCCL_ASYNC_ERROR_HANDLING=1, +// NCCL_IB_QPS_PER_CONNECTION=2, NCCL_IB_TIMEOUT=22, NCCL_IB_RETRY_CNT=7) and +// NCCL_PROTO=Simple unless already present. Returns the number of variables it set. +int export_rccl_environment(const ClusterConfig & cfg); + +// Kill switches / debug (environment, documented in variables.md). +bool cluster_env_no_ingraph_allreduce(); // LUCE_CLUSTER_NO_INGRAPH_ALLREDUCE +bool cluster_env_no_graph_capture(); // LUCE_CLUSTER_NO_GRAPH_CAPTURE +bool cluster_env_trace(); // LUCE_CLUSTER_TRACE + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_control.cpp b/server/src/cluster/cluster_control.cpp new file mode 100644 index 000000000..424b0f364 --- /dev/null +++ b/server/src/cluster/cluster_control.cpp @@ -0,0 +1,1069 @@ +// cluster_control.cpp - POSIX TCP control channel between head and workers +// (cluster_control.h). +// +// Design: one socket per worker. After the Hello/Welcome handshake every peer +// owns a service thread that (a) polls the socket and moves complete frames +// into a queue, filtering heartbeats and refreshing liveness on any inbound +// frame, and (b) sends a HeartbeatMsg every heartbeat_period_ms. recv() pops +// from the queue under a deadline, so a dead peer becomes an error, never a +// hang, even when nobody is currently receiving from it. +// +// Windows: the whole implementation is compiled out and every entry point +// returns false with "cluster control channel is POSIX-only" so the Windows +// CI leg still links the protocol tests. + +#include "cluster/cluster_control.h" + +#include +#include + +#if !defined(_WIN32) + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace luce::cluster { + +namespace { + +void set_err(std::string * err, const std::string & msg) { + if (err) *err = msg; +} + +std::string errno_string(const std::string & what) { + return what + ": " + std::strerror(errno); +} + +uint64_t deadline_from_now_us(uint32_t ms) { + return monotonic_us() + (uint64_t)ms * 1000ull; +} + +// Remaining milliseconds until an absolute monotonic deadline (>= 0). +int remaining_ms(uint64_t deadline_us) { + const uint64_t now = monotonic_us(); + if (now >= deadline_us) return 0; + const uint64_t rem = (deadline_us - now + 999) / 1000; + return rem > 0x7fffffffull ? 0x7fffffff : (int)rem; +} + +// Poll one fd for `events` until the deadline. Returns 1 ready, 0 timeout, +// -1 error (errno set; EINTR retried). +int poll_until(int fd, short events, uint64_t deadline_us) { + for (;;) { + struct pollfd p; + p.fd = fd; + p.events = events; + p.revents = 0; + const int r = ::poll(&p, 1, remaining_ms(deadline_us)); + if (r < 0) { + if (errno == EINTR) continue; + return -1; + } + if (r == 0) return 0; + if (p.revents & (POLLERR | POLLNVAL)) { + errno = ECONNRESET; + return -1; + } + // POLLHUP with readable data still delivers the tail; a read returning + // 0 reports EOF afterwards. + return 1; + } +} + +bool set_nonblocking(int fd, bool on) { + const int flags = ::fcntl(fd, F_GETFL, 0); + if (flags < 0) return false; + const int next = on ? (flags | O_NONBLOCK) : (flags & ~O_NONBLOCK); + return ::fcntl(fd, F_SETFL, next) == 0; +} + +void apply_stream_options(int fd) { + int one = 1; + (void)::setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one)); + (void)::setsockopt(fd, SOL_SOCKET, SO_KEEPALIVE, &one, sizeof(one)); +#ifdef TCP_KEEPIDLE + int idle = 5; + (void)::setsockopt(fd, IPPROTO_TCP, TCP_KEEPIDLE, &idle, sizeof(idle)); +#endif +#ifdef TCP_KEEPINTVL + int intvl = 2; + (void)::setsockopt(fd, IPPROTO_TCP, TCP_KEEPINTVL, &intvl, sizeof(intvl)); +#endif +#ifdef TCP_KEEPCNT + int cnt = 5; + (void)::setsockopt(fd, IPPROTO_TCP, TCP_KEEPCNT, &cnt, sizeof(cnt)); +#endif +} + +bool write_exact_deadline_impl(int fd, const void * data, size_t n, uint64_t deadline_us, + std::string * err) { + const uint8_t * p = (const uint8_t *)data; + size_t done = 0; + while (done < n) { + const ssize_t w = ::send(fd, p + done, n - done, MSG_NOSIGNAL); + if (w > 0) { + done += (size_t)w; + continue; + } + if (w < 0 && errno == EINTR) continue; + if (w < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) { + const int r = poll_until(fd, POLLOUT, deadline_us); + if (r == 0) { + set_err(err, "control channel write timed out"); + return false; + } + if (r < 0) { + set_err(err, errno_string("control channel poll(POLLOUT) failed")); + return false; + } + continue; + } + set_err(err, w == 0 ? std::string("control channel write returned 0") + : errno_string("control channel write failed")); + return false; + } + return true; +} + +bool read_exact_deadline_impl(int fd, void * data, size_t n, uint64_t deadline_us, + std::string * err) { + uint8_t * p = (uint8_t *)data; + size_t done = 0; + while (done < n) { + const ssize_t r = ::recv(fd, p + done, n - done, 0); + if (r > 0) { + done += (size_t)r; + continue; + } + if (r == 0) { + set_err(err, "control channel closed by peer (EOF)"); + return false; + } + if (errno == EINTR) continue; + if (errno == EAGAIN || errno == EWOULDBLOCK) { + const int pr = poll_until(fd, POLLIN, deadline_us); + if (pr == 0) { + set_err(err, "control channel read timed out"); + return false; + } + if (pr < 0) { + set_err(err, errno_string("control channel poll(POLLIN) failed")); + return false; + } + continue; + } + set_err(err, errno_string("control channel read failed")); + return false; + } + return true; +} + +bool write_frame_impl(int fd, const Frame & frame, uint64_t deadline_us, std::string * err) { + if (frame.payload.size() > kMaxPayloadBytes) { + set_err(err, "frame payload too large to send"); + return false; + } + FrameHeader h; + h.type = (uint16_t)frame.type; + h.payload_len = (uint32_t)frame.payload.size(); + uint8_t hdr[kFrameHeaderBytes]; + encode_header(h, hdr); + // One buffer so header and payload share a segment for small frames. + std::vector buf; + buf.reserve(kFrameHeaderBytes + frame.payload.size()); + buf.insert(buf.end(), hdr, hdr + kFrameHeaderBytes); + buf.insert(buf.end(), frame.payload.begin(), frame.payload.end()); + return write_exact_deadline_impl(fd, buf.data(), buf.size(), deadline_us, err); +} + +bool read_frame_impl(int fd, Frame & frame, uint64_t deadline_us, std::string * err) { + uint8_t hdr[kFrameHeaderBytes]; + if (!read_exact_deadline_impl(fd, hdr, kFrameHeaderBytes, deadline_us, err)) return false; + FrameHeader h; + if (!decode_header(hdr, h, err)) return false; + frame.type = (MsgType)h.type; + frame.payload.resize(h.payload_len); + if (h.payload_len > 0 && + !read_exact_deadline_impl(fd, frame.payload.data(), h.payload_len, deadline_us, err)) { + return false; + } + return true; +} + +// Resolve host:port to a list of addresses. `passive` selects AI_PASSIVE for +// bind on an empty host. +bool resolve(const std::string & host, int port, bool passive, struct addrinfo ** out, + std::string * err) { + struct addrinfo hints; + std::memset(&hints, 0, sizeof(hints)); + hints.ai_family = AF_UNSPEC; + hints.ai_socktype = SOCK_STREAM; + hints.ai_flags = passive ? AI_PASSIVE : 0; + if (!passive) hints.ai_flags |= AI_ADDRCONFIG; + const std::string port_str = std::to_string(port); + const char * node = host.empty() ? nullptr : host.c_str(); + const int rc = ::getaddrinfo(node, port_str.c_str(), &hints, out); + if (rc != 0) { + set_err(err, "getaddrinfo(" + (host.empty() ? std::string("*") : host) + ":" + + port_str + ") failed: " + ::gai_strerror(rc)); + return false; + } + return true; +} + +// Connect with a deadline. Returns the connected fd or -1. +int connect_with_deadline(const std::string & host, int port, uint32_t connect_ms, + std::string * err) { + struct addrinfo * ai = nullptr; + if (!resolve(host, port, /*passive=*/false, &ai, err)) return -1; + + int fd = -1; + std::string last_err = "no addresses resolved"; + for (struct addrinfo * a = ai; a != nullptr; a = a->ai_next) { + fd = ::socket(a->ai_family, a->ai_socktype, a->ai_protocol); + if (fd < 0) { + last_err = errno_string("socket() failed"); + continue; + } + if (!set_nonblocking(fd, true)) { + last_err = errno_string("fcntl(O_NONBLOCK) failed"); + ::close(fd); + fd = -1; + continue; + } + const uint64_t deadline_us = deadline_from_now_us(connect_ms); + int rc = ::connect(fd, a->ai_addr, a->ai_addrlen); + if (rc < 0 && errno == EINPROGRESS) { + const int pr = poll_until(fd, POLLOUT, deadline_us); + if (pr == 0) { + last_err = "connect timed out"; + ::close(fd); + fd = -1; + continue; + } + if (pr < 0) { + last_err = errno_string("connect poll failed"); + ::close(fd); + fd = -1; + continue; + } + int so_err = 0; + socklen_t len = sizeof(so_err); + if (::getsockopt(fd, SOL_SOCKET, SO_ERROR, &so_err, &len) < 0 || so_err != 0) { + errno = so_err != 0 ? so_err : errno; + last_err = errno_string("connect failed"); + ::close(fd); + fd = -1; + continue; + } + rc = 0; + } + if (rc < 0) { + last_err = errno_string("connect failed"); + ::close(fd); + fd = -1; + continue; + } + break; + } + ::freeaddrinfo(ai); + if (fd < 0) { + set_err(err, last_err); + return -1; + } + apply_stream_options(fd); + return fd; +} + +// ─── Peer ─────────────────────────────────────────────────────────────── +// One connected socket plus its service thread. `self_rank` is stamped into +// outgoing heartbeats; `rank_` is the remote rank. + +class SocketPeer final : public ControlPeer { +public: + SocketPeer(int fd, int remote_rank, int self_rank, const ControlTimeouts & t) + : fd_(fd), rank_(remote_rank), self_rank_(self_rank), timeouts_(t), + last_hb_us_(monotonic_us()) { + set_nonblocking(fd_, true); + } + + ~SocketPeer() override { close(); } + + // Starts the service thread. Called once after the handshake. + void start() { + std::lock_guard lk(state_mu_); + if (running_.load() || closed_.load()) return; + running_.store(true); + thread_ = std::thread([this] { service_loop(); }); + } + + int peer_rank() const override { return rank_; } + + bool send(const Frame & frame, std::string * err) override { + std::lock_guard lk(send_mu_); + if (closed_.load()) { + set_err(err, "control peer rank " + std::to_string(rank_) + " is closed: " + fail_reason()); + return false; + } + std::string e; + if (!write_frame_impl(fd_, frame, deadline_from_now_us(timeouts_.recv_ms), &e)) { + mark_failed("send " + std::string(msg_type_name(frame.type)) + " to rank " + + std::to_string(rank_) + " failed: " + e); + set_err(err, fail_reason()); + return false; + } + return true; + } + + bool recv(Frame & frame, uint32_t deadline_ms, std::string * err) override { + const uint64_t deadline_us = deadline_from_now_us(deadline_ms); + std::unique_lock lk(queue_mu_); + for (;;) { + if (!queue_.empty()) { + frame = std::move(queue_.front()); + queue_.pop_front(); + return true; + } + if (closed_.load()) { + set_err(err, "control peer rank " + std::to_string(rank_) + " is closed: " + fail_reason()); + return false; + } + if (!running_.load()) { + // Handshake phase: no service thread yet, read synchronously. + lk.unlock(); + return recv_direct(frame, deadline_us, err); + } + const int rem = remaining_ms(deadline_us); + if (rem <= 0) { + set_err(err, "timed out waiting for a frame from rank " + std::to_string(rank_)); + return false; + } + queue_cv_.wait_for(lk, std::chrono::milliseconds(rem)); + } + } + + bool alive() const override { + if (closed_.load()) return false; + const uint64_t now = monotonic_us(); + const uint64_t last = last_hb_us_.load(); + return now < last + (uint64_t)timeouts_.heartbeat_timeout_ms * 1000ull; + } + + uint64_t last_heartbeat_us() const override { return last_hb_us_.load(); } + + void close() override { + bool join = false; + { + std::lock_guard lk(state_mu_); + stop_ = true; + if (fd_ >= 0) ::shutdown(fd_, SHUT_RDWR); + join = running_.load() && thread_.joinable() && + thread_.get_id() != std::this_thread::get_id(); + } + if (join) thread_.join(); + { + // send_mu_ so no writer is mid-send on the descriptor we close. + std::lock_guard slk(send_mu_); + std::lock_guard lk(state_mu_); + running_.store(false); + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } + closed_.store(true); + } + queue_cv_.notify_all(); + } + + std::string fail_reason() const { + std::lock_guard lk(state_mu_); + return fail_reason_; + } + +private: + // Synchronous frame read used before the service thread exists. + bool recv_direct(Frame & frame, uint64_t deadline_us, std::string * err) { + for (;;) { + std::string e; + if (!read_frame_impl(fd_, frame, deadline_us, &e)) { + set_err(err, e); + return false; + } + last_hb_us_.store(monotonic_us()); + if (frame.type == MsgType::Heartbeat) continue; + return true; + } + } + + void mark_failed(const std::string & reason) { + { + std::lock_guard lk(state_mu_); + if (fail_reason_.empty()) fail_reason_ = reason; + stop_ = true; + } + closed_.store(true); + queue_cv_.notify_all(); + } + + bool send_heartbeat() { + HeartbeatMsg hb; + hb.rank = self_rank_; + hb.mono_us = monotonic_us(); + const Frame f = make_frame(hb); + std::lock_guard lk(send_mu_); + if (closed_.load()) return false; + std::string e; + if (!write_frame_impl(fd_, f, deadline_from_now_us(timeouts_.heartbeat_period_ms), &e)) { + mark_failed("heartbeat to rank " + std::to_string(rank_) + " failed: " + e); + return false; + } + return true; + } + + void service_loop() { + uint64_t next_hb_us = monotonic_us(); + for (;;) { + { + std::lock_guard lk(state_mu_); + if (stop_) break; + } + const uint64_t now = monotonic_us(); + if (now >= next_hb_us) { + if (!send_heartbeat()) break; + next_hb_us = now + (uint64_t)timeouts_.heartbeat_period_ms * 1000ull; + } + const int pr = poll_until(fd_, POLLIN, next_hb_us); + if (pr < 0) { + if (errno == EBADF) break; // closed underneath us + mark_failed("control socket to rank " + std::to_string(rank_) + " failed: " + + std::strerror(errno)); + break; + } + if (pr == 0) continue; // heartbeat due + Frame f; + std::string e; + // A frame is at least partially readable; finish it under the + // recv deadline so a stalled sender cannot pin this thread. + if (!read_frame_impl(fd_, f, deadline_from_now_us(timeouts_.recv_ms), &e)) { + mark_failed("control socket to rank " + std::to_string(rank_) + ": " + e); + break; + } + last_hb_us_.store(monotonic_us()); + if (f.type == MsgType::Heartbeat) continue; + { + std::lock_guard lk(queue_mu_); + queue_.push_back(std::move(f)); + } + queue_cv_.notify_all(); + } + // Detect a peer that stopped heartbeating without closing: the poll + // above returns on every period, so the liveness check is implicit in + // alive(); nothing more to do here. + } + + int fd_; + const int rank_; + const int self_rank_; + ControlTimeouts timeouts_; + + mutable std::mutex state_mu_; + std::atomic running_{false}; + bool stop_ = false; + std::string fail_reason_; + std::thread thread_; + + std::mutex send_mu_; + + std::mutex queue_mu_; + std::condition_variable queue_cv_; + std::deque queue_; + + std::atomic closed_{false}; + std::atomic last_hb_us_; +}; + +} // namespace + +// ─── Free helpers ─────────────────────────────────────────────────────── + +bool write_exact_deadline(int fd, const void * data, size_t n, uint32_t deadline_ms, std::string * err) { + return write_exact_deadline_impl(fd, data, n, deadline_from_now_us(deadline_ms), err); +} + +bool read_exact_deadline(int fd, void * data, size_t n, uint32_t deadline_ms, std::string * err) { + return read_exact_deadline_impl(fd, data, n, deadline_from_now_us(deadline_ms), err); +} + +bool write_frame(int fd, const Frame & frame, uint32_t deadline_ms, std::string * err) { + return write_frame_impl(fd, frame, deadline_from_now_us(deadline_ms), err); +} + +bool read_frame(int fd, Frame & frame, uint32_t deadline_ms, std::string * err) { + return read_frame_impl(fd, frame, deadline_from_now_us(deadline_ms), err); +} + +uint64_t monotonic_us() { + using namespace std::chrono; + return (uint64_t)duration_cast(steady_clock::now().time_since_epoch()).count(); +} + +std::string local_hostname(int rank) { + char buf[256]; + std::memset(buf, 0, sizeof(buf)); + if (::gethostname(buf, sizeof(buf) - 1) == 0 && buf[0] != '\0') { + return std::string(buf); + } + return "rank" + std::to_string(rank); +} + +// ─── ClusterHeadControl ───────────────────────────────────────────────── + +struct ClusterHeadControl::Impl { + int listen_fd = -1; + ControlTimeouts timeouts; + // rank -> peer (ranks 1..n_workers) + std::map> peers; + std::mutex mu; + + ~Impl() { close_all(); } + + void close_all() { + std::map> dead; + { + std::lock_guard lk(mu); + if (listen_fd >= 0) { + ::close(listen_fd); + listen_fd = -1; + } + dead.swap(peers); + } + for (auto & kv : dead) kv.second->close(); + } + + SocketPeer * peer(int rank) { + std::lock_guard lk(mu); + auto it = peers.find(rank); + return it == peers.end() ? nullptr : it->second.get(); + } +}; + +ClusterHeadControl::ClusterHeadControl() : impl_(new Impl) {} +ClusterHeadControl::~ClusterHeadControl() { close(); } + +bool ClusterHeadControl::listen(const std::string & host, int port, std::string * err) { + if (impl_->listen_fd >= 0) { + set_err(err, "control channel already listening"); + return false; + } + const bool any = host.empty() || host == "0.0.0.0" || host == "*" || host == "::"; + struct addrinfo * ai = nullptr; + if (!resolve(any ? std::string() : host, port, /*passive=*/true, &ai, err)) return false; + + int fd = -1; + std::string last_err = "no addresses resolved"; + for (struct addrinfo * a = ai; a != nullptr; a = a->ai_next) { + fd = ::socket(a->ai_family, a->ai_socktype, a->ai_protocol); + if (fd < 0) { + last_err = errno_string("socket() failed"); + continue; + } + int one = 1; + (void)::setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)); + if (a->ai_family == AF_INET6) { + int zero = 0; + (void)::setsockopt(fd, IPPROTO_IPV6, IPV6_V6ONLY, &zero, sizeof(zero)); + } + if (::bind(fd, a->ai_addr, a->ai_addrlen) < 0) { + last_err = errno_string("bind(" + host + ":" + std::to_string(port) + ") failed"); + ::close(fd); + fd = -1; + continue; + } + if (::listen(fd, 16) < 0) { + last_err = errno_string("listen() failed"); + ::close(fd); + fd = -1; + continue; + } + break; + } + ::freeaddrinfo(ai); + if (fd < 0) { + set_err(err, last_err); + return false; + } + if (!set_nonblocking(fd, true)) { + set_err(err, errno_string("fcntl(O_NONBLOCK) on listen socket failed")); + ::close(fd); + return false; + } + impl_->listen_fd = fd; + return true; +} + +bool ClusterHeadControl::accept_workers(int n_workers, const HelloMsg & head_identity, + uint32_t deadline_ms, std::vector & hellos, + std::string * err) { + if (impl_->listen_fd < 0) { + set_err(err, "accept_workers called before listen()"); + return false; + } + if (n_workers < 1) { + set_err(err, "accept_workers needs at least one worker"); + return false; + } + const uint64_t deadline_us = deadline_from_now_us(deadline_ms); + hellos.assign((size_t)n_workers, HelloMsg{}); + std::vector seen((size_t)n_workers + 1, false); + std::map> accepted; + + // Refuse a worker: send Abort with the reason, close, fail the whole + // bootstrap so the operator sees one attributable error. + auto refuse = [&](int fd, int rank, const std::string & reason) -> bool { + AbortMsg ab; + ab.rank = 0; + ab.code = 1; + ab.reason = reason; + std::string ignored; + (void)write_frame_impl(fd, make_frame(ab), deadline_from_now_us(2000), &ignored); + ::close(fd); + set_err(err, "worker rank " + std::to_string(rank) + " refused: " + reason); + for (auto & kv : accepted) kv.second->close(); + return false; + }; + + int n_accepted = 0; + while (n_accepted < n_workers) { + const int pr = poll_until(impl_->listen_fd, POLLIN, deadline_us); + if (pr == 0) { + set_err(err, "timed out waiting for workers (" + std::to_string(n_accepted) + "/" + + std::to_string(n_workers) + " connected)"); + for (auto & kv : accepted) kv.second->close(); + return false; + } + if (pr < 0) { + set_err(err, errno_string("poll on listen socket failed")); + for (auto & kv : accepted) kv.second->close(); + return false; + } + struct sockaddr_storage ss; + socklen_t sl = sizeof(ss); + const int fd = ::accept(impl_->listen_fd, (struct sockaddr *)&ss, &sl); + if (fd < 0) { + if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK) continue; + set_err(err, errno_string("accept() failed")); + for (auto & kv : accepted) kv.second->close(); + return false; + } + apply_stream_options(fd); + set_nonblocking(fd, true); + + Frame f; + std::string e; + if (!read_frame_impl(fd, f, deadline_us, &e)) { + ::close(fd); + set_err(err, "reading Hello failed: " + e); + for (auto & kv : accepted) kv.second->close(); + return false; + } + HelloMsg hello; + if (f.type != MsgType::Hello || !parse_frame(f, hello)) { + return refuse(fd, -1, "expected a well-formed Hello frame, got " + + std::string(msg_type_name(f.type))); + } + if (hello.protocol_version != kProtocolVersion) { + return refuse(fd, hello.rank, "protocol version " + std::to_string(hello.protocol_version) + + " != head " + std::to_string(kProtocolVersion)); + } + if (hello.rank < 1 || hello.rank > n_workers) { + return refuse(fd, hello.rank, "rank out of range [1, " + std::to_string(n_workers) + "]"); + } + if (seen[(size_t)hello.rank]) { + return refuse(fd, hello.rank, "duplicate rank"); + } + if (hello.size != head_identity.size) { + return refuse(fd, hello.rank, "--cluster-size " + std::to_string(hello.size) + + " != head " + std::to_string(head_identity.size)); + } + if (hello.build_sha != head_identity.build_sha) { + return refuse(fd, hello.rank, "build_sha '" + hello.build_sha + "' != head '" + + head_identity.build_sha + "'"); + } + if (hello.model_sha != head_identity.model_sha) { + return refuse(fd, hello.rank, "model_sha '" + hello.model_sha + "' != head '" + + head_identity.model_sha + "'"); + } + if (hello.placement_hash != head_identity.placement_hash) { + return refuse(fd, hello.rank, "placement_hash mismatch (worker " + + std::to_string(hello.placement_hash) + ", head " + + std::to_string(head_identity.placement_hash) + ")"); + } + seen[(size_t)hello.rank] = true; + hellos[(size_t)hello.rank - 1] = hello; + accepted[hello.rank].reset(new SocketPeer(fd, hello.rank, /*self_rank=*/0, impl_->timeouts)); + ++n_accepted; + } + + { + std::lock_guard lk(impl_->mu); + for (auto & kv : accepted) { + kv.second->start(); + impl_->peers[kv.first] = std::move(kv.second); + } + } + return true; +} + +bool ClusterHeadControl::send_to(int rank, const Frame & frame, std::string * err) { + SocketPeer * p = impl_->peer(rank); + if (!p) { + set_err(err, "no worker with rank " + std::to_string(rank)); + return false; + } + return p->send(frame, err); +} + +bool ClusterHeadControl::broadcast(const Frame & frame, std::string * err) { + std::vector targets; + { + std::lock_guard lk(impl_->mu); + for (auto & kv : impl_->peers) targets.push_back(kv.second.get()); + } + bool ok = true; + for (SocketPeer * p : targets) { + std::string e; + if (!p->send(frame, &e)) { + if (ok) set_err(err, "broadcast to rank " + std::to_string(p->peer_rank()) + " failed: " + e); + ok = false; + } + } + return ok; +} + +bool ClusterHeadControl::recv_from(int rank, Frame & frame, uint32_t deadline_ms, std::string * err) { + SocketPeer * p = impl_->peer(rank); + if (!p) { + set_err(err, "no worker with rank " + std::to_string(rank)); + return false; + } + return p->recv(frame, deadline_ms, err); +} + +bool ClusterHeadControl::gather(MsgType type, uint32_t deadline_ms, std::vector & out, + std::string * err) { + const int n = n_workers(); + out.assign((size_t)n, Frame{}); + const uint64_t deadline_us = deadline_from_now_us(deadline_ms); + for (int rank = 1; rank <= n; ++rank) { + Frame f; + std::string e; + const int rem = remaining_ms(deadline_us); + if (!recv_from(rank, f, (uint32_t)rem, &e)) { + set_err(err, "gather(" + std::string(msg_type_name(type)) + ") from rank " + + std::to_string(rank) + " failed: " + e); + return false; + } + if (f.type != type) { + if (f.type == MsgType::Abort) { + AbortMsg ab; + if (parse_frame(f, ab)) { + set_err(err, "rank " + std::to_string(rank) + " aborted: " + ab.reason); + return false; + } + } + set_err(err, "gather(" + std::string(msg_type_name(type)) + ") from rank " + + std::to_string(rank) + " got " + msg_type_name(f.type)); + return false; + } + out[(size_t)rank - 1] = std::move(f); + } + return true; +} + +int ClusterHeadControl::n_workers() const { + std::lock_guard lk(impl_->mu); + return (int)impl_->peers.size(); +} + +bool ClusterHeadControl::all_alive() const { + std::lock_guard lk(impl_->mu); + for (const auto & kv : impl_->peers) { + if (!kv.second->alive()) return false; + } + return true; +} + +std::vector ClusterHeadControl::dead_ranks() const { + std::vector out; + std::lock_guard lk(impl_->mu); + for (const auto & kv : impl_->peers) { + if (!kv.second->alive()) out.push_back(kv.first); + } + return out; +} + +void ClusterHeadControl::set_timeouts(const ControlTimeouts & t) { + impl_->timeouts = t; +} + +void ClusterHeadControl::close() { + if (impl_) impl_->close_all(); +} + +// ─── ClusterWorkerControl ─────────────────────────────────────────────── + +struct ClusterWorkerControl::Impl { + ControlTimeouts timeouts; + std::unique_ptr head; + std::mutex mu; + + ~Impl() { close_head(); } + + void close_head() { + std::unique_ptr dead; + { + std::lock_guard lk(mu); + dead = std::move(head); + } + if (dead) dead->close(); + } +}; + +ClusterWorkerControl::ClusterWorkerControl() : impl_(new Impl) {} +ClusterWorkerControl::~ClusterWorkerControl() { close(); } + +bool ClusterWorkerControl::connect_and_handshake(const std::string & head_host, int head_port, + const HelloMsg & hello, WelcomeMsg & welcome, + std::string * err) { + if (impl_->head) { + set_err(err, "worker control channel already connected"); + return false; + } + const ControlTimeouts & t = impl_->timeouts; + int fd = -1; + std::string last_err; + const uint32_t attempts = t.connect_retries == 0 ? 1u : t.connect_retries; + for (uint32_t attempt = 0; attempt < attempts && fd < 0; ++attempt) { + fd = connect_with_deadline(head_host, head_port, t.connect_ms, &last_err); + if (fd < 0 && attempt + 1 < attempts) { + const uint32_t sleep_ms = t.connect_ms < 1000 ? t.connect_ms : 1000; + std::this_thread::sleep_for(std::chrono::milliseconds(sleep_ms)); + } + } + if (fd < 0) { + set_err(err, "connect to head " + head_host + ":" + std::to_string(head_port) + + " failed after " + std::to_string(attempts) + " attempts: " + last_err); + return false; + } + set_nonblocking(fd, true); + + std::string e; + if (!write_frame_impl(fd, make_frame(hello), deadline_from_now_us(t.recv_ms), &e)) { + ::close(fd); + set_err(err, "sending Hello failed: " + e); + return false; + } + Frame f; + for (;;) { + if (!read_frame_impl(fd, f, deadline_from_now_us(t.recv_ms), &e)) { + ::close(fd); + set_err(err, "waiting for Welcome failed: " + e); + return false; + } + if (f.type == MsgType::Heartbeat) continue; + break; + } + if (f.type == MsgType::Abort) { + AbortMsg ab; + ::close(fd); + if (parse_frame(f, ab)) { + set_err(err, "head refused this worker: " + ab.reason); + } else { + set_err(err, "head refused this worker (unparseable Abort)"); + } + return false; + } + if (f.type != MsgType::Welcome || !parse_frame(f, welcome)) { + ::close(fd); + set_err(err, "expected Welcome, got " + std::string(msg_type_name(f.type))); + return false; + } + + std::unique_ptr peer(new SocketPeer(fd, /*remote_rank=*/0, hello.rank, t)); + peer->start(); + std::lock_guard lk(impl_->mu); + impl_->head = std::move(peer); + return true; +} + +bool ClusterWorkerControl::send(const Frame & frame, std::string * err) { + SocketPeer * p = nullptr; + { + std::lock_guard lk(impl_->mu); + p = impl_->head.get(); + } + if (!p) { + set_err(err, "worker control channel not connected"); + return false; + } + return p->send(frame, err); +} + +bool ClusterWorkerControl::recv(Frame & frame, uint32_t deadline_ms, std::string * err) { + SocketPeer * p = nullptr; + { + std::lock_guard lk(impl_->mu); + p = impl_->head.get(); + } + if (!p) { + set_err(err, "worker control channel not connected"); + return false; + } + return p->recv(frame, deadline_ms, err); +} + +bool ClusterWorkerControl::alive() const { + std::lock_guard lk(impl_->mu); + return impl_->head && impl_->head->alive(); +} + +void ClusterWorkerControl::set_timeouts(const ControlTimeouts & t) { + impl_->timeouts = t; +} + +void ClusterWorkerControl::close() { + if (impl_) impl_->close_head(); +} + +} // namespace luce::cluster + +#else // _WIN32 ───────────────────────────────────────────────────────────── + +namespace luce::cluster { + +namespace { +const char * kPosixOnly = "cluster control channel is POSIX-only"; +void set_posix_only(std::string * err) { + if (err) *err = kPosixOnly; +} +} // namespace + +bool write_exact_deadline(int, const void *, size_t, uint32_t, std::string * err) { + set_posix_only(err); + return false; +} + +bool read_exact_deadline(int, void *, size_t, uint32_t, std::string * err) { + set_posix_only(err); + return false; +} + +bool write_frame(int, const Frame &, uint32_t, std::string * err) { + set_posix_only(err); + return false; +} + +bool read_frame(int, Frame &, uint32_t, std::string * err) { + set_posix_only(err); + return false; +} + +uint64_t monotonic_us() { + using namespace std::chrono; + return (uint64_t)duration_cast(steady_clock::now().time_since_epoch()).count(); +} + +std::string local_hostname(int rank) { + return "rank" + std::to_string(rank); +} + +struct ClusterHeadControl::Impl { + ControlTimeouts timeouts; +}; + +ClusterHeadControl::ClusterHeadControl() : impl_(new Impl) {} +ClusterHeadControl::~ClusterHeadControl() = default; + +bool ClusterHeadControl::listen(const std::string &, int, std::string * err) { + set_posix_only(err); + return false; +} + +bool ClusterHeadControl::accept_workers(int, const HelloMsg &, uint32_t, std::vector & hellos, + std::string * err) { + hellos.clear(); + set_posix_only(err); + return false; +} + +bool ClusterHeadControl::send_to(int, const Frame &, std::string * err) { + set_posix_only(err); + return false; +} + +bool ClusterHeadControl::broadcast(const Frame &, std::string * err) { + set_posix_only(err); + return false; +} + +bool ClusterHeadControl::recv_from(int, Frame &, uint32_t, std::string * err) { + set_posix_only(err); + return false; +} + +bool ClusterHeadControl::gather(MsgType, uint32_t, std::vector & out, std::string * err) { + out.clear(); + set_posix_only(err); + return false; +} + +int ClusterHeadControl::n_workers() const { return 0; } +bool ClusterHeadControl::all_alive() const { return false; } +std::vector ClusterHeadControl::dead_ranks() const { return {}; } +void ClusterHeadControl::set_timeouts(const ControlTimeouts & t) { impl_->timeouts = t; } +void ClusterHeadControl::close() {} + +struct ClusterWorkerControl::Impl { + ControlTimeouts timeouts; +}; + +ClusterWorkerControl::ClusterWorkerControl() : impl_(new Impl) {} +ClusterWorkerControl::~ClusterWorkerControl() = default; + +bool ClusterWorkerControl::connect_and_handshake(const std::string &, int, const HelloMsg &, + WelcomeMsg &, std::string * err) { + set_posix_only(err); + return false; +} + +bool ClusterWorkerControl::send(const Frame &, std::string * err) { + set_posix_only(err); + return false; +} + +bool ClusterWorkerControl::recv(Frame &, uint32_t, std::string * err) { + set_posix_only(err); + return false; +} + +bool ClusterWorkerControl::alive() const { return false; } +void ClusterWorkerControl::set_timeouts(const ControlTimeouts & t) { impl_->timeouts = t; } +void ClusterWorkerControl::close() {} + +} // namespace luce::cluster + +#endif // _WIN32 diff --git a/server/src/cluster/cluster_control.h b/server/src/cluster/cluster_control.h new file mode 100644 index 000000000..d1a3eeea4 --- /dev/null +++ b/server/src/cluster/cluster_control.h @@ -0,0 +1,133 @@ +// cluster_control.h - TCP control channel between the head and its workers. +// +// One long-lived TCP connection per worker to the head (rank 0). Frames from +// cluster_protocol.h are written with a 12-byte header and read back with +// exact-length reads; every blocking receive takes a deadline so a dead peer +// turns into an error, never a hang. Heartbeats are sent by an internal +// thread while the channel is idle and their absence for +// `heartbeat_timeout_ms` marks the peer dead. +// +// POSIX sockets only (the whole multi-node path is Linux-only, like the +// existing backend IPC). The header itself is portable so the protocol unit +// test builds on the Windows CI leg without a control channel. +// +// Thread safety: send() may be called from one thread while recv() blocks on +// another; both are internally serialized per direction. + +#pragma once + +#include "cluster/cluster_protocol.h" + +#include +#include +#include +#include +#include + +namespace luce::cluster { + +struct ControlTimeouts { + uint32_t connect_ms = 5000; + uint32_t connect_retries = 60; // workers retry until head is up (5 min) + uint32_t recv_ms = 30000; // default deadline for a blocking recv + uint32_t heartbeat_period_ms = 1000; + uint32_t heartbeat_timeout_ms = 10000; +}; + +// A single connected peer (one worker as seen by the head, or the head as +// seen by a worker). +class ControlPeer { +public: + virtual ~ControlPeer() = default; + virtual int peer_rank() const = 0; + virtual bool send(const Frame & frame, std::string * err) = 0; + // Blocks until a non-heartbeat frame arrives or deadline_ms elapses. + // Heartbeats are consumed internally and refresh liveness. + virtual bool recv(Frame & frame, uint32_t deadline_ms, std::string * err) = 0; + virtual bool alive() const = 0; + virtual uint64_t last_heartbeat_us() const = 0; + virtual void close() = 0; +}; + +// Head side: accept exactly size-1 workers, each identifying itself with a +// HelloMsg, then fan out frames. +class ClusterHeadControl { +public: + ClusterHeadControl(); + ~ClusterHeadControl(); + ClusterHeadControl(const ClusterHeadControl &) = delete; + ClusterHeadControl & operator=(const ClusterHeadControl &) = delete; + + bool listen(const std::string & host, int port, std::string * err); + + // Waits for `n_workers` connections and their Hello frames. Validates + // rank uniqueness, range [1, n_workers], protocol version, and that + // build_sha / model_sha / placement_hash equal the head's values (any + // mismatch -> false with a rank-attributed error). On success + // hellos[rank-1] is filled and workers are addressable by rank. + bool accept_workers(int n_workers, const HelloMsg & head_identity, uint32_t deadline_ms, + std::vector & hellos, std::string * err); + + bool send_to(int rank, const Frame & frame, std::string * err); + // Sends to every worker; returns false if any send failed (err names the + // first failing rank). + bool broadcast(const Frame & frame, std::string * err); + bool recv_from(int rank, Frame & frame, uint32_t deadline_ms, std::string * err); + // Collects one frame of the given type from every worker (e.g. every + // RequestReport); out[rank-1] is filled. + bool gather(MsgType type, uint32_t deadline_ms, std::vector & out, std::string * err); + + int n_workers() const; + bool all_alive() const; + // Ranks whose heartbeat is stale; empty when healthy. + std::vector dead_ranks() const; + + void set_timeouts(const ControlTimeouts & t); + void close(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +// Worker side: connect to the head, send Hello, receive Welcome. +class ClusterWorkerControl { +public: + ClusterWorkerControl(); + ~ClusterWorkerControl(); + ClusterWorkerControl(const ClusterWorkerControl &) = delete; + ClusterWorkerControl & operator=(const ClusterWorkerControl &) = delete; + + // Connects (with retries), sends `hello`, waits for Welcome. Returns + // false on refusal (the head sends Abort with the mismatch reason). + bool connect_and_handshake(const std::string & head_host, int head_port, const HelloMsg & hello, + WelcomeMsg & welcome, std::string * err); + + bool send(const Frame & frame, std::string * err); + bool recv(Frame & frame, uint32_t deadline_ms, std::string * err); + bool alive() const; + + void set_timeouts(const ControlTimeouts & t); + void close(); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +// Utility used by both sides and by tests: exact-length socket I/O with a +// deadline. fd is a connected stream socket. Returns false on EOF, error or +// timeout (err distinguishes them). +bool write_exact_deadline(int fd, const void * data, size_t n, uint32_t deadline_ms, std::string * err); +bool read_exact_deadline(int fd, void * data, size_t n, uint32_t deadline_ms, std::string * err); +bool write_frame(int fd, const Frame & frame, uint32_t deadline_ms, std::string * err); +bool read_frame(int fd, Frame & frame, uint32_t deadline_ms, std::string * err); + +// Monotonic microseconds, shared clock for heartbeats and telemetry. +uint64_t monotonic_us(); + +// Resolve the local hostname for HelloMsg::hostname (never fails; falls back +// to "rank"). +std::string local_hostname(int rank); + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_decision_hooks.cpp b/server/src/cluster/cluster_decision_hooks.cpp new file mode 100644 index 000000000..8ec94fea7 --- /dev/null +++ b/server/src/cluster/cluster_decision_hooks.cpp @@ -0,0 +1,200 @@ +// cluster_decision_hooks.cpp - head / worker / local decision hooks. + +#include "cluster/cluster_decision_hooks.h" +#include "cluster/cluster_control.h" + + +namespace luce::cluster { + +namespace { + +void set_err(std::string * err, const std::string & msg) { + if (err) *err = msg; +} + +} // namespace + +// ─── HeadHooks ───────────────────────────────────────────────────────── + +HeadHooks::HeadHooks(ClusterHeadControl & control) : control_(control) {} + +bool HeadHooks::decide_next_token(uint64_t req, uint32_t step, const float * /*logits*/, + int /*n_vocab*/, const common::SamplerCfg & /*cfg*/, + bool /*is_eos_candidate_unused*/, + int32_t & token, uint8_t & flags, std::string * err) { + // Rank 0 samples in its decode loop, with the request's sampler and RNG. + if (token < 0) { + set_err(err, "head: no token to broadcast"); + return false; + } + + DecisionMsg msg; + msg.request_id = req; + msg.step = step; + msg.token = token; + msg.flags = flags; + + const uint64_t t0 = monotonic_us(); + const bool ok = control_.broadcast(make_frame(msg), err); + counters_.ctrl_wait_us += monotonic_us() - t0; + if (!ok) return false; + counters_.decisions++; + return true; +} + +bool HeadHooks::decide_draft(uint64_t req, uint32_t step, int pos, std::vector & tokens, + std::string * err) { + DraftMsg msg; + msg.request_id = req; + msg.step = step; + msg.pos = pos; + msg.tokens = tokens; + const uint64_t t0 = monotonic_us(); + const bool ok = control_.broadcast(make_frame(msg), err); + counters_.ctrl_wait_us += monotonic_us() - t0; + if (!ok) return false; + counters_.drafts++; + return true; +} + +bool HeadHooks::decide_accept(uint64_t req, uint32_t step, int32_t & accept, int32_t & bonus, + uint8_t & flags, std::string * err) { + AcceptMsg msg; + msg.request_id = req; + msg.step = step; + msg.accept = accept; + msg.bonus = bonus; + msg.flags = flags; + const uint64_t t0 = monotonic_us(); + const bool ok = control_.broadcast(make_frame(msg), err); + counters_.ctrl_wait_us += monotonic_us() - t0; + if (!ok) return false; + counters_.accepts++; + return true; +} + +bool HeadHooks::backend_op(const BackendOpMsg & op, std::string * err) { + const uint64_t t0 = monotonic_us(); + const bool ok = control_.broadcast(make_frame(op), err); + counters_.ctrl_wait_us += monotonic_us() - t0; + if (!ok) return false; + counters_.backend_ops++; + return true; +} + +// ─── WorkerHooks ─────────────────────────────────────────────────────── + +WorkerHooks::WorkerHooks(ClusterWorkerControl & control, const WorkerHooksConfig & cfg) + : control_(control), cfg_(cfg) {} + +bool WorkerHooks::take_pending_frame(Frame & out) { + if (!has_pending_) return false; + out = std::move(pending_); + pending_ = Frame{}; + has_pending_ = false; + return true; +} + +bool WorkerHooks::recv_expected(MsgType want, Frame & frame, std::string * err) { + if (aborted_) { + set_err(err, "worker: head aborted: " + abort_reason_); + return false; + } + if (has_pending_) { + // A frame parked by an earlier mismatch is still unconsumed; the + // main loop must handle it before the backend asks for decisions. + set_err(err, std::string("worker: unconsumed ") + msg_type_name(pending_.type) + + " frame while waiting for " + msg_type_name(want)); + return false; + } + const uint64_t t0 = monotonic_us(); + const bool ok = control_.recv(frame, cfg_.timeout_ms, err); + counters_.ctrl_wait_us += monotonic_us() - t0; + if (!ok) return false; + if (frame.type == want) return true; + + if (frame.type == MsgType::Abort) { + AbortMsg abort; + aborted_ = true; + abort_reason_ = parse_frame(frame, abort) ? abort.reason : std::string("(malformed Abort)"); + set_err(err, "worker: head aborted: " + abort_reason_); + return false; + } + pending_ = std::move(frame); + has_pending_ = true; + set_err(err, std::string("worker: expected ") + msg_type_name(want) + " but received " + + msg_type_name(pending_.type)); + return false; +} + +bool WorkerHooks::decide_next_token(uint64_t req, uint32_t step, const float * /*logits*/, + int /*n_vocab*/, const common::SamplerCfg & /*cfg*/, + bool /*is_eos_candidate_unused*/, int32_t & token, + uint8_t & flags, std::string * err) { + Frame frame; + if (!recv_expected(MsgType::Decision, frame, err)) return false; + DecisionMsg msg; + if (!parse_frame(frame, msg)) { + set_err(err, "worker: malformed Decision frame"); + return false; + } + if (msg.request_id != req || msg.step != step) { + set_err(err, "worker: Decision for req=" + std::to_string(msg.request_id) + + " step=" + std::to_string(msg.step) + ", expected req=" + + std::to_string(req) + " step=" + std::to_string(step)); + return false; + } + token = msg.token; + flags = msg.flags; + counters_.decisions++; + return true; +} + +bool WorkerHooks::decide_draft(uint64_t req, uint32_t step, int pos, std::vector & tokens, + std::string * err) { + Frame frame; + if (!recv_expected(MsgType::Draft, frame, err)) return false; + DraftMsg msg; + if (!parse_frame(frame, msg)) { + set_err(err, "worker: malformed Draft frame"); + return false; + } + if (msg.request_id != req || msg.step != step) { + set_err(err, "worker: Draft for req=" + std::to_string(msg.request_id) + + " step=" + std::to_string(msg.step) + ", expected req=" + + std::to_string(req) + " step=" + std::to_string(step)); + return false; + } + if (msg.pos != pos) { + set_err(err, "worker: Draft pos=" + std::to_string(msg.pos) + " but local pos=" + + std::to_string(pos) + " (ranks diverged)"); + return false; + } + tokens = std::move(msg.tokens); + counters_.drafts++; + return true; +} + +bool WorkerHooks::decide_accept(uint64_t req, uint32_t step, int32_t & accept, int32_t & bonus, + uint8_t & flags, std::string * err) { + Frame frame; + if (!recv_expected(MsgType::Accept, frame, err)) return false; + AcceptMsg msg; + if (!parse_frame(frame, msg)) { + set_err(err, "worker: malformed Accept frame"); + return false; + } + if (msg.request_id != req || msg.step != step) { + set_err(err, "worker: Accept for req=" + std::to_string(msg.request_id) + + " step=" + std::to_string(msg.step) + ", expected req=" + + std::to_string(req) + " step=" + std::to_string(step)); + return false; + } + accept = msg.accept; + bonus = msg.bonus; + flags = msg.flags; + counters_.accepts++; + return true; +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_decision_hooks.h b/server/src/cluster/cluster_decision_hooks.h new file mode 100644 index 000000000..0d951b769 --- /dev/null +++ b/server/src/cluster/cluster_decision_hooks.h @@ -0,0 +1,202 @@ +// cluster_decision_hooks.h - "rank 0 decides, all ranks follow" hooks for +// the DeepSeek4 generate loop. +// +// Every rank of a cluster runs the identical generate() call in lockstep. The +// only places where the ranks could diverge are host-side decisions: which +// token was sampled, whether EOS/stop/cancel ended the request, which draft +// tokens DSpark proposes and how many the target accepted. Those decisions +// are made once on the head and shipped to the workers over the control +// channel (cluster_control.h) as Decision / Draft / Accept frames. +// +// The backend calls these hooks at the points where upstream Lucebox samples +// or checks for termination. Three implementations exist: +// +// LocalHooks - cluster disabled. Every call is a no-op that returns true +// and leaves the caller's token/flags untouched, so the +// backend keeps exactly one code path and upstream behaviour +// stays byte-identical. +// HeadHooks - rank 0. The caller has already sampled (host sampler, +// deterministic RNG on rank 0 only) and computed the flags; +// the hook broadcasts DecisionMsg. +// WorkerHooks - ranks 1..N-1. Blocks on the head's frame with the +// configured deadline, validates request_id/step, and hands +// the head's token/flags back to the caller. Workers never +// sample and never evaluate EOS/cancel themselves. +// +// All methods return false and fill *err on channel failure or protocol +// violation; nothing here throws across the backend boundary. + +#pragma once + +#include "cluster/cluster_protocol.h" +#include "common/sampler.h" + +#include +#include +#include + +namespace luce::cluster { + +class ClusterHeadControl; +class ClusterWorkerControl; + +// Counters kept by every hook implementation; the worker main loop folds +// them into RequestReportMsg, the head decorator into ClusterRequestReport. +struct HookCounters { + uint64_t decisions = 0; // Decision frames sent (head) / received (worker) + uint64_t drafts = 0; + uint64_t accepts = 0; + uint64_t backend_ops = 0; + uint64_t ctrl_wait_us = 0; // time blocked on the control channel + + void reset() { *this = HookCounters{}; } +}; + +// Compose DecisionFlags from the caller's termination checks. +inline uint8_t make_decision_flags(bool eos, bool stop, bool cancel, bool budget, bool is_final) { + uint8_t f = 0; + if (eos) f |= kDecisionEos; + if (stop) f |= kDecisionStop; + if (cancel) f |= kDecisionCancel; + if (budget) f |= kDecisionBudget; + if (is_final) f |= kDecisionFinal; + return f; +} + +// True when the flags end the request on every rank. +inline bool decision_terminates(uint8_t flags) { + return (flags & (kDecisionEos | kDecisionStop | kDecisionCancel | kDecisionFinal)) != 0; +} + +struct Ds4ClusterHooks { + virtual ~Ds4ClusterHooks() = default; + + // AR decode, one call per step on every rank. + // + // Head / local: `token` and `flags` are INPUTS. The caller sampled + // `token` from `logits` with `cfg` (after the thinking budget) and + // computed `flags` with make_decision_flags; the head refuses token < 0. + // Worker: `logits`/`cfg` are ignored; `token` and `flags` are OUTPUTS + // taken from the head's DecisionMsg for (req, step). + // `is_eos_candidate_unused` is reserved (kept for signature stability). + virtual bool decide_next_token(uint64_t req, uint32_t step, + const float * logits, int n_vocab, + const common::SamplerCfg & cfg, + bool is_eos_candidate_unused, + int32_t & token, uint8_t & flags, + std::string * err) = 0; + + // DSpark: head broadcasts the drafted block `tokens` (seed + + // candidates, tokens[0] at target position `pos`); worker fills it. + virtual bool decide_draft(uint64_t req, uint32_t step, int pos, + std::vector & tokens, std::string * err) = 0; + + // DSpark: head broadcasts accept count, bonus token and flags; + // worker fills them. + virtual bool decide_accept(uint64_t req, uint32_t step, + int32_t & accept, int32_t & bonus, uint8_t & flags, + std::string * err) = 0; + + // Replicated backend mutation. Head: broadcast to all workers. Worker: + // no-op (the worker main loop already executed the op it received). + virtual bool backend_op(const BackendOpMsg & op, std::string * err) = 0; + + virtual bool is_head() const = 0; + // False for LocalHooks: the backend uses this to keep upstream cancel / + // early-return semantics when no cluster is configured. + virtual bool is_cluster() const = 0; + + const HookCounters & counters() const { return counters_; } + void reset_counters() { counters_.reset(); } + + // The request id the head decorator / worker main loop assigned to the + // generate() call currently in flight. The backend passes it back on + // every call so frames can be validated. + void set_current_request(uint64_t id) { current_request_ = id; } + uint64_t current_request() const { return current_request_; } + +protected: + HookCounters counters_; + uint64_t current_request_ = 0; +}; + +// ─── LocalHooks ──────────────────────────────────────────────────────── +class LocalHooks final : public Ds4ClusterHooks { +public: + bool decide_next_token(uint64_t, uint32_t, const float *, int, const common::SamplerCfg &, + bool, int32_t &, uint8_t &, std::string *) override { return true; } + bool decide_draft(uint64_t, uint32_t, int, std::vector &, std::string *) override { + return true; + } + bool decide_accept(uint64_t, uint32_t, int32_t &, int32_t &, uint8_t &, std::string *) override { + return true; + } + bool backend_op(const BackendOpMsg &, std::string *) override { return true; } + bool is_head() const override { return true; } + bool is_cluster() const override { return false; } +}; + +// ─── HeadHooks ───────────────────────────────────────────────────────── +class HeadHooks final : public Ds4ClusterHooks { +public: + // `control` must outlive the hooks. + explicit HeadHooks(ClusterHeadControl & control); + + bool decide_next_token(uint64_t req, uint32_t step, const float * logits, int n_vocab, + const common::SamplerCfg & cfg, bool is_eos_candidate_unused, + int32_t & token, uint8_t & flags, std::string * err) override; + bool decide_draft(uint64_t req, uint32_t step, int pos, std::vector & tokens, + std::string * err) override; + bool decide_accept(uint64_t req, uint32_t step, int32_t & accept, int32_t & bonus, + uint8_t & flags, std::string * err) override; + bool backend_op(const BackendOpMsg & op, std::string * err) override; + bool is_head() const override { return true; } + bool is_cluster() const override { return true; } + +private: + ClusterHeadControl & control_; +}; + +// ─── WorkerHooks ─────────────────────────────────────────────────────── +struct WorkerHooksConfig { + uint32_t timeout_ms = 30000; // deadline for every blocking recv +}; + +class WorkerHooks final : public Ds4ClusterHooks { +public: + WorkerHooks(ClusterWorkerControl & control, const WorkerHooksConfig & cfg); + + bool decide_next_token(uint64_t req, uint32_t step, const float * logits, int n_vocab, + const common::SamplerCfg & cfg, bool is_eos_candidate_unused, + int32_t & token, uint8_t & flags, std::string * err) override; + bool decide_draft(uint64_t req, uint32_t step, int pos, std::vector & tokens, + std::string * err) override; + bool decide_accept(uint64_t req, uint32_t step, int32_t & accept, int32_t & bonus, + uint8_t & flags, std::string * err) override; + bool backend_op(const BackendOpMsg &, std::string *) override { return true; } + bool is_head() const override { return false; } + bool is_cluster() const override { return true; } + + // When a recv expected a Decision/Draft/Accept but got a different frame + // (RequestEnd because the head failed the request, Abort, Shutdown), the + // frame is parked here so the worker main loop can act on it instead of + // blocking for it a second time. Empty when nothing is parked. + bool take_pending_frame(Frame & out); + // Reason of the last Abort frame received, if any. + const std::string & abort_reason() const { return abort_reason_; } + bool aborted() const { return aborted_; } + +private: + // Receives the next non-heartbeat frame; returns false if it is not of + // `want` (the frame is parked or, for Abort, recorded). + bool recv_expected(MsgType want, Frame & frame, std::string * err); + + ClusterWorkerControl & control_; + WorkerHooksConfig cfg_; + Frame pending_; + bool has_pending_ = false; + bool aborted_ = false; + std::string abort_reason_; +}; + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_expert_placement.cpp b/server/src/cluster/cluster_expert_placement.cpp new file mode 100644 index 000000000..ad64c4a98 --- /dev/null +++ b/server/src/cluster/cluster_expert_placement.cpp @@ -0,0 +1,617 @@ +// cluster_expert_placement.cpp - N-rank routed-expert ownership. +// +// GPU-free: MoeHybridRoutingStats is only read through count() and +// ranked_experts(); no ggml call is made here so the placement unit test can +// link without a device backend. JSON I/O is hand-rolled (same policy as +// cluster_protocol.h) so the loader has no third-party dependency. + +#include "cluster/cluster_expert_placement.h" +#include "cluster/cluster_protocol.h" // fnv1a64 + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace luce::cluster { + +namespace { + +void set_err(std::string * err, const std::string & msg) { + if (err) *err = msg; +} + +bool dims_ok(int n_ranks, int n_layer, int n_expert, int n_expert_used, std::string * err) { + if (n_ranks < 1) { + set_err(err, "n_ranks must be >= 1"); + return false; + } + if (n_layer <= 0 || n_expert <= 0) { + set_err(err, "n_layer and n_expert must be positive"); + return false; + } + if (n_expert_used <= 0 || n_expert_used > n_expert) { + set_err(err, "n_expert_used must be in [1, n_expert]"); + return false; + } + if ((size_t) n_layer > (size_t) std::numeric_limits::max() / (size_t) n_expert) { + set_err(err, "n_layer * n_expert overflows"); + return false; + } + return true; +} + +// ── minimal JSON reader ──────────────────────────────────────────────── +// Accepts exactly the shape written by save_json (and by +// scripts/cluster/build_expert_placement.py): a flat object of integers plus +// one "owner" key holding an array of arrays of integers. Unknown keys are +// skipped (strings, numbers, nested arrays/objects) so a future field does +// not break older binaries. + +class JsonCursor { +public: + explicit JsonCursor(const std::string & text) : s_(text) {} + + bool parse_object(ClusterExpertPlacement & out, std::string * err) { + skip_ws(); + if (!expect('{', err)) return false; + skip_ws(); + if (peek() == '}') { ++pos_; return true; } + for (;;) { + skip_ws(); + std::string key; + if (!parse_string(key, err)) return false; + skip_ws(); + if (!expect(':', err)) return false; + skip_ws(); + if (key == "owner") { + if (!parse_owner(out, err)) return false; + } else if (key == "n_ranks" || key == "n_layer" || key == "n_expert" || + key == "n_expert_used" || key == "replicate_hot") { + long long v = 0; + if (!parse_int(v, err)) return false; + if (v < 0 || v > std::numeric_limits::max()) { + set_err(err, "placement json: value out of range for " + key); + return false; + } + if (key == "n_ranks") out.n_ranks = (int) v; + else if (key == "n_layer") out.n_layer = (int) v; + else if (key == "n_expert") out.n_expert = (int) v; + else if (key == "n_expert_used") out.n_expert_used = (int) v; + else out.replicate_hot = (int) v; + } else { + if (!skip_value(err)) return false; + } + skip_ws(); + const char c = peek(); + if (c == ',') { ++pos_; continue; } + if (c == '}') { ++pos_; return true; } + set_err(err, "placement json: expected ',' or '}'"); + return false; + } + } + +private: + char peek() const { return pos_ < s_.size() ? s_[pos_] : '\0'; } + + void skip_ws() { + while (pos_ < s_.size() && std::isspace((unsigned char) s_[pos_])) ++pos_; + } + + bool expect(char c, std::string * err) { + if (peek() != c) { + set_err(err, std::string("placement json: expected '") + c + "'"); + return false; + } + ++pos_; + return true; + } + + bool parse_string(std::string & out, std::string * err) { + if (!expect('"', err)) return false; + out.clear(); + while (pos_ < s_.size()) { + const char c = s_[pos_++]; + if (c == '"') return true; + if (c == '\\') { + if (pos_ >= s_.size()) break; + const char e = s_[pos_++]; + switch (e) { + case 'n': out.push_back('\n'); break; + case 't': out.push_back('\t'); break; + case 'r': out.push_back('\r'); break; + case 'u': + // \uXXXX: keep the raw escape; keys we care about are ASCII. + out.push_back('?'); + pos_ = std::min(s_.size(), pos_ + 4); + break; + default: out.push_back(e); break; + } + continue; + } + out.push_back(c); + } + set_err(err, "placement json: unterminated string"); + return false; + } + + bool parse_int(long long & out, std::string * err) { + skip_ws(); + const size_t start = pos_; + if (peek() == '-') ++pos_; + while (pos_ < s_.size() && std::isdigit((unsigned char) s_[pos_])) ++pos_; + if (pos_ == start || (pos_ == start + 1 && s_[start] == '-')) { + set_err(err, "placement json: expected integer"); + return false; + } + // Reject fractions/exponents: owner ids and dims are integers. + if (peek() == '.' || peek() == 'e' || peek() == 'E') { + set_err(err, "placement json: non-integer number"); + return false; + } + out = std::strtoll(s_.c_str() + start, nullptr, 10); + return true; + } + + bool parse_owner(ClusterExpertPlacement & out, std::string * err) { + if (!expect('[', err)) return false; + out.owner.clear(); + row_width_ = -1; + skip_ws(); + if (peek() == ']') { ++pos_; return true; } + for (;;) { + skip_ws(); + if (!expect('[', err)) return false; + int width = 0; + skip_ws(); + if (peek() != ']') { + for (;;) { + long long v = 0; + if (!parse_int(v, err)) return false; + if (v < kReplicated || v > std::numeric_limits::max()) { + set_err(err, "placement json: owner value out of range"); + return false; + } + out.owner.push_back((int32_t) v); + ++width; + skip_ws(); + const char c = peek(); + if (c == ',') { ++pos_; continue; } + if (c == ']') break; + set_err(err, "placement json: expected ',' or ']' in owner row"); + return false; + } + } + ++pos_; // ']' + if (row_width_ < 0) row_width_ = width; + else if (row_width_ != width) { + set_err(err, "placement json: ragged owner rows"); + return false; + } + skip_ws(); + const char c = peek(); + if (c == ',') { ++pos_; continue; } + if (c == ']') { ++pos_; return true; } + set_err(err, "placement json: expected ',' or ']' after owner row"); + return false; + } + } + + bool skip_value(std::string * err) { + skip_ws(); + const char c = peek(); + if (c == '"') { std::string tmp; return parse_string(tmp, err); } + if (c == '{' || c == '[') { + const char close = c == '{' ? '}' : ']'; + int depth = 0; + bool in_str = false; + while (pos_ < s_.size()) { + const char d = s_[pos_++]; + if (in_str) { + if (d == '\\') ++pos_; + else if (d == '"') in_str = false; + continue; + } + if (d == '"') in_str = true; + else if (d == c) ++depth; + else if (d == close && --depth == 0) return true; + } + set_err(err, "placement json: unterminated container"); + return false; + } + // number / true / false / null + while (pos_ < s_.size() && s_[pos_] != ',' && s_[pos_] != '}' && s_[pos_] != ']') ++pos_; + return true; + } + + const std::string & s_; + size_t pos_ = 0; + int row_width_ = -1; +}; + +} // namespace + +// ─── Queries ──────────────────────────────────────────────────────────── + +int32_t ClusterExpertPlacement::owner_of(int layer, int expert) const { + if (layer < 0 || layer >= n_layer || expert < 0 || expert >= n_expert) return -2; + const size_t idx = (size_t) layer * (size_t) n_expert + (size_t) expert; + return idx < owner.size() ? owner[idx] : -2; +} + +bool ClusterExpertPlacement::is_replicated(int layer, int expert) const { + return owner_of(layer, expert) == kReplicated; +} + +int32_t ClusterExpertPlacement::slot_owner(int layer, int expert, int slot) const { + const int32_t o = owner_of(layer, expert); + if (o >= 0) return o; + if (o == kReplicated && n_ranks > 0 && slot >= 0) { + return (int32_t) (((int64_t) expert + (int64_t) slot) % (int64_t) n_ranks); + } + return -2; +} + +std::vector ClusterExpertPlacement::resident_experts(int layer, int rank) const { + std::vector out; + if (layer < 0 || layer >= n_layer) return out; + for (int e = 0; e < n_expert; ++e) { + const int32_t o = owner_of(layer, e); + if (o == rank || o == kReplicated) out.push_back((int32_t) e); + } + return out; // ascending by construction +} + +int ClusterExpertPlacement::owned_count(int layer, int rank) const { + if (layer < 0 || layer >= n_layer) return 0; + int n = 0; + for (int e = 0; e < n_expert; ++e) { + if (owner_of(layer, e) == rank) ++n; + } + return n; +} + +int ClusterExpertPlacement::total_owned(int rank) const { + int n = 0; + for (int il = 0; il < n_layer; ++il) n += owned_count(il, rank); + return n; +} + +bool ClusterExpertPlacement::valid(std::string * err) const { + if (!dims_ok(n_ranks, n_layer, n_expert, n_expert_used, err)) return false; + if (replicate_hot < 0 || replicate_hot > n_expert) { + set_err(err, "replicate_hot out of range"); + return false; + } + if (owner.size() != (size_t) n_layer * (size_t) n_expert) { + set_err(err, "owner table size does not match n_layer * n_expert"); + return false; + } + std::vector owned((size_t) n_ranks); + for (int il = 0; il < n_layer; ++il) { + std::fill(owned.begin(), owned.end(), 0); + int replicated = 0; + for (int e = 0; e < n_expert; ++e) { + const int32_t o = owner[(size_t) il * (size_t) n_expert + (size_t) e]; + if (o == kReplicated) { ++replicated; continue; } + if (o < 0 || o >= n_ranks) { + set_err(err, "owner rank out of range at layer " + std::to_string(il) + + " expert " + std::to_string(e)); + return false; + } + ++owned[(size_t) o]; + } + // When there are at least n_ranks non-replicated experts, every rank + // must own at least one: a rank with no owned experts would evaluate + // only replicated routes, which is never what the operator wanted. + if (n_expert - replicated >= n_ranks && !allow_empty_shards) { + for (int r = 0; r < n_ranks; ++r) { + if (owned[(size_t) r] == 0) { + set_err(err, "rank " + std::to_string(r) + " owns no expert in layer " + + std::to_string(il)); + return false; + } + } + } + } + return true; +} + +bool ClusterExpertPlacement::matches(const common::MoeHybridConfig & cfg) const { + return n_layer == cfg.n_layer && n_expert == cfg.n_expert && + n_expert_used == cfg.n_expert_used; +} + +uint64_t ClusterExpertPlacement::hash() const { + // Field order is part of the wire contract (HelloMsg.placement_hash): + // n_ranks, n_layer, n_expert, n_expert_used as int32, then the owner + // table as int32 in layer-major order. Byte order is the host's; all + // cluster nodes run the same binary on the same ISA. + const int32_t dims[4] = { (int32_t) n_ranks, (int32_t) n_layer, + (int32_t) n_expert, (int32_t) n_expert_used }; + uint64_t h = fnv1a64(dims, sizeof(dims)); + if (!owner.empty()) { + h = fnv1a64(owner.data(), owner.size() * sizeof(int32_t), h); + } + return h; +} + +bool ClusterExpertPlacement::to_rank_placement(int rank, common::MoeHybridPlacement & out, + std::string * err) const { + if (!valid(err)) return false; + if (rank < 0 || rank >= n_ranks) { + set_err(err, "rank out of range"); + return false; + } + out = common::MoeHybridPlacement{}; + out.n_layer = n_layer; + out.n_expert = n_expert; + out.n_expert_used = n_expert_used; + out.hot_counts.assign((size_t) n_layer, 0); + out.hot_expert_ids.assign((size_t) n_layer, {}); + int total = 0; + for (int il = 0; il < n_layer; ++il) { + out.hot_expert_ids[(size_t) il] = resident_experts(il, rank); + out.hot_counts[(size_t) il] = (int) out.hot_expert_ids[(size_t) il].size(); + total += out.hot_counts[(size_t) il]; + } + out.total_hot = total; + return out.valid(err); +} + +void ClusterExpertPlacement::slot_validity(int rank, int n_slots, std::vector & out) const { + out.clear(); + if (n_slots <= 0 || n_layer <= 0 || n_expert <= 0) return; + out.assign((size_t) n_layer * (size_t) n_slots * (size_t) n_expert, 0); + for (int il = 0; il < n_layer; ++il) { + for (int e = 0; e < n_expert; ++e) { + const int32_t o = owner_of(il, e); + if (o != rank && o != kReplicated) continue; + for (int t = 0; t < n_slots; ++t) { + const bool mine = (o == rank) || (slot_owner(il, e, t) == rank); + out[((size_t) il * (size_t) n_slots + (size_t) t) * (size_t) n_expert + (size_t) e] = + mine ? 1 : 0; + } + } + } +} + +// ─── Builders ─────────────────────────────────────────────────────────── + +bool ClusterExpertPlacement::build_uniform(int n_ranks, + const common::MoeHybridConfig & cfg, + int replicate_hot, + const common::MoeHybridRoutingStats * stats, + ClusterExpertPlacement & out, + std::string * err) { + if (!dims_ok(n_ranks, cfg.n_layer, cfg.n_expert, cfg.n_expert_used, err)) return false; + if (replicate_hot < 0) { + set_err(err, "replicate_hot must be >= 0"); + return false; + } + if (stats && replicate_hot > 0 && !stats->matches(cfg)) { + set_err(err, "routing stats do not match the model dimensions"); + return false; + } + const int effective_replicate = (stats && replicate_hot > 0) ? replicate_hot : 0; + if (effective_replicate > cfg.n_expert - n_ranks) { + set_err(err, "replicate_hot leaves fewer experts than ranks"); + return false; + } + + out = ClusterExpertPlacement{}; + out.n_ranks = n_ranks; + out.n_layer = cfg.n_layer; + out.n_expert = cfg.n_expert; + out.n_expert_used = cfg.n_expert_used; + out.replicate_hot = effective_replicate; + out.owner.assign((size_t) cfg.n_layer * (size_t) cfg.n_expert, 0); + + for (int il = 0; il < cfg.n_layer; ++il) { + int32_t * row = out.owner.data() + (size_t) il * (size_t) cfg.n_expert; + for (int e = 0; e < cfg.n_expert; ++e) row[e] = (int32_t) (e % n_ranks); + if (effective_replicate > 0) { + const std::vector ranked = stats->ranked_experts(il); + for (int i = 0; i < effective_replicate && i < (int) ranked.size(); ++i) { + row[ranked[(size_t) i]] = kReplicated; + } + } + } + // Replicating the hottest experts can in principle strip a rank of every + // owned expert in a layer (e.g. n_expert == n_ranks + k); valid() catches + // that so the operator lowers replicate_hot instead of running lopsided. + return out.valid(err); +} + +bool ClusterExpertPlacement::build_balanced(int n_ranks, + const common::MoeHybridRoutingStats & stats, + int replicate_hot, + ClusterExpertPlacement & out, + std::string * err) { + if (!stats.valid(err)) return false; + if (!dims_ok(n_ranks, stats.n_layer, stats.n_expert, stats.n_expert_used, err)) return false; + if (replicate_hot < 0) { + set_err(err, "replicate_hot must be >= 0"); + return false; + } + if (replicate_hot > stats.n_expert - n_ranks) { + set_err(err, "replicate_hot leaves fewer experts than ranks"); + return false; + } + + out = ClusterExpertPlacement{}; + out.n_ranks = n_ranks; + out.n_layer = stats.n_layer; + out.n_expert = stats.n_expert; + out.n_expert_used = stats.n_expert_used; + out.replicate_hot = replicate_hot; + out.owner.assign((size_t) stats.n_layer * (size_t) stats.n_expert, 0); + + std::vector load((size_t) n_ranks); + std::vector owned((size_t) n_ranks); + for (int il = 0; il < stats.n_layer; ++il) { + int32_t * row = out.owner.data() + (size_t) il * (size_t) stats.n_expert; + // ranked_experts is count-descending with the expert id as tie-break, + // so identical inputs produce identical placements on every rank. + const std::vector ranked = stats.ranked_experts(il); + if ((int) ranked.size() != stats.n_expert) { + set_err(err, "routing stats ranking failed at layer " + std::to_string(il)); + return false; + } + for (int i = 0; i < replicate_hot; ++i) row[ranked[(size_t) i]] = kReplicated; + + const int remaining = stats.n_expert - replicate_hot; + const int cap = (remaining + n_ranks - 1) / n_ranks; // ceil(remaining / n_ranks) + std::fill(load.begin(), load.end(), 0); + std::fill(owned.begin(), owned.end(), 0); + for (int i = replicate_hot; i < stats.n_expert; ++i) { + const int e = ranked[(size_t) i]; + int best = -1; + for (int r = 0; r < n_ranks; ++r) { + if (owned[(size_t) r] >= cap) continue; + if (best < 0 || load[(size_t) r] < load[(size_t) best]) best = r; + } + if (best < 0) { + set_err(err, "balanced placement: no rank below the memory cap"); + return false; + } + row[e] = (int32_t) best; + load[(size_t) best] += stats.count(il, e); + ++owned[(size_t) best]; + } + } + return out.valid(err); +} + +// ─── JSON ─────────────────────────────────────────────────────────────── + +bool ClusterExpertPlacement::save_json(const std::string & path, std::string * err) const { + if (!valid(err)) return false; + std::ofstream f(path); + if (!f) { + set_err(err, "failed to open placement file for writing: " + path); + return false; + } + f << "{\n" + << " \"n_ranks\": " << n_ranks << ",\n" + << " \"n_layer\": " << n_layer << ",\n" + << " \"n_expert\": " << n_expert << ",\n" + << " \"n_expert_used\": " << n_expert_used << ",\n" + << " \"replicate_hot\": " << replicate_hot << ",\n" + << " \"owner\": [\n"; + for (int il = 0; il < n_layer; ++il) { + f << " ["; + const int32_t * row = owner.data() + (size_t) il * (size_t) n_expert; + for (int e = 0; e < n_expert; ++e) { + if (e) f << ','; + f << row[e]; + } + f << (il + 1 < n_layer ? "],\n" : "]\n"); + } + f << " ]\n}\n"; + if (!f) { + set_err(err, "failed to write placement file: " + path); + return false; + } + return true; +} + +bool ClusterExpertPlacement::load_json(const std::string & path, ClusterExpertPlacement & out, + std::string * err) { + std::ifstream f(path, std::ios::binary); + if (!f) { + set_err(err, "failed to open placement file: " + path); + return false; + } + std::stringstream ss; + ss << f.rdbuf(); + const std::string text = ss.str(); + + ClusterExpertPlacement parsed; + JsonCursor cursor(text); + if (!cursor.parse_object(parsed, err)) return false; + if (parsed.owner.size() != (size_t) parsed.n_layer * (size_t) parsed.n_expert) { + set_err(err, "placement json: owner table does not match n_layer x n_expert"); + return false; + } + // Files are the diagnostic entry point: tolerate empty shards, but say so. + parsed.allow_empty_shards = true; + if (!parsed.valid(err)) return false; + for (int r = 0; r < parsed.n_ranks; ++r) { + for (int il = 0; il < parsed.n_layer; ++il) { + if (parsed.owned_count(il, r) == 0) { + std::fprintf(stderr, + "[cluster-placement] warning: %s leaves rank %d with no owned " + "expert in layer %d (diagnostic placement; every route of that " + "layer is evaluated elsewhere)\n", + path.c_str(), r, il); + break; + } + } + } + out = std::move(parsed); + return true; +} + +// ─── Log summary ──────────────────────────────────────────────────────── + +std::string ClusterExpertPlacement::describe(const common::MoeHybridRoutingStats * stats) const { + std::ostringstream os; + os << "placement ranks=" << n_ranks << " layers=" << n_layer + << " experts=" << n_expert << " top_k=" << n_expert_used + << " replicate_hot=" << replicate_hot << " hash=0x" << std::hex << hash() << std::dec; + int replicated_total = 0; + for (size_t i = 0; i < owner.size(); ++i) { + if (owner[i] == kReplicated) ++replicated_total; + } + os << " replicated_total=" << replicated_total << " owned_per_rank=["; + for (int r = 0; r < n_ranks; ++r) { + if (r) os << ','; + os << total_owned(r); + } + os << "]"; + + if (stats && stats->matches(n_layer, n_expert, n_expert_used)) { + // Expected per-rank activation share: owned counts plus an equal + // 1/N share of every replicated expert's activations. The max/mean + // ratio over all layers is the load imbalance the all-reduce waits on. + double worst = 0.0; + double sum_ratio = 0.0; + int layers_counted = 0; + std::vector load((size_t) n_ranks); + for (int il = 0; il < n_layer; ++il) { + std::fill(load.begin(), load.end(), 0.0); + double total = 0.0; + for (int e = 0; e < n_expert; ++e) { + const double c = (double) stats->count(il, e); + total += c; + const int32_t o = owner_of(il, e); + if (o == kReplicated) { + for (int r = 0; r < n_ranks; ++r) load[(size_t) r] += c / (double) n_ranks; + } else if (o >= 0) { + load[(size_t) o] += c; + } + } + if (total <= 0.0) continue; + const double mean = total / (double) n_ranks; + const double mx = *std::max_element(load.begin(), load.end()); + const double ratio = mean > 0.0 ? mx / mean : 1.0; + worst = std::max(worst, ratio); + sum_ratio += ratio; + ++layers_counted; + } + if (layers_counted > 0) { + char buf[96]; + std::snprintf(buf, sizeof(buf), " imbalance(max/mean) worst=%.3f avg=%.3f", + worst, sum_ratio / (double) layers_counted); + os << buf; + } + } + return os.str(); +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_expert_placement.h b/server/src/cluster/cluster_expert_placement.h new file mode 100644 index 000000000..0a4b55727 --- /dev/null +++ b/server/src/cluster/cluster_expert_placement.h @@ -0,0 +1,110 @@ +// cluster_expert_placement.h - which rank owns which routed expert. +// +// Generalizes the two-owner hot/cold split of moe_hybrid_placement.h to N +// ranks. Every (layer, expert) has exactly one owner rank, or is marked +// kReplicated: replicated experts are resident on every rank and the rank +// that computes a replicated expert for token slot t is (expert + t) % N, so +// each route is still evaluated exactly once cluster-wide. Ranks evaluate +// their local experts as owner partial sums (the existing masked-route path +// of build_moe_hybrid_ffn_graph with cold owner None) and the cluster +// all-reduce sums the partials. +// +// The placement is built once at startup (uniform, balanced from routing +// stats, or loaded from JSON), hashed, and the hash is exchanged in Hello so +// ranks with different placements refuse to form a cluster. + +#pragma once + +#include "common/moe_hybrid_placement.h" +#include "common/moe_hybrid_routing_stats.h" + +#include +#include +#include +#include + +namespace luce::cluster { + +inline constexpr int32_t kReplicated = -1; + +struct ClusterExpertPlacement { + int n_ranks = 0; + int n_layer = 0; + int n_expert = 0; + int n_expert_used = 0; + int replicate_hot = 0; // replicated experts per layer (informational) + // Diagnostics only: set by load_json so a hand-written file may leave a + // rank with no owned expert in a layer (e.g. rank 0 owns everything and + // rank 1 nothing, which must reproduce the single-node output exactly if + // the cluster path is correct). Builders never set it. + bool allow_empty_shards = false; + + // Flattened [layer * n_expert + expert] -> owner rank, or kReplicated. + std::vector owner; + + // ── Queries ────────────────────────────────────────────────────────── + int32_t owner_of(int layer, int expert) const; + bool is_replicated(int layer, int expert) const; + // Rank that evaluates (layer, expert) for token slot `slot` (0-based + // position inside the current batch). Owned experts: the owner; + // replicated experts: (expert + slot) % n_ranks. + int32_t slot_owner(int layer, int expert, int slot) const; + // Experts resident on `rank` for `layer`: owned + replicated, ascending. + std::vector resident_experts(int layer, int rank) const; + // Number of owned (non-replicated) experts of `rank` in `layer`. + int owned_count(int layer, int rank) const; + // Sum over layers of owned_count. + int total_owned(int rank) const; + + bool valid(std::string * err = nullptr) const; + bool matches(const common::MoeHybridConfig & cfg) const; + uint64_t hash() const; + + // Rank-local view for the existing hybrid machinery: hot set = resident + // experts of `rank`, cold set = everything else (never materialized, + // cold owner None). hot_expert_ids are sorted ascending so the LUT order + // is deterministic across ranks. + bool to_rank_placement(int rank, common::MoeHybridPlacement & out, + std::string * err = nullptr) const; + + // Per-(layer, slot) LUT rows for replicated experts: for every expert e + // resident on `rank`, valid[layer][slot][e] says whether this rank + // evaluates e at that slot. Owned experts are always valid; replicated + // ones only when slot_owner == rank. n_slots is the verify width / batch + // size. Output is [n_layer][n_slots][n_expert] flattened, 0/1. + void slot_validity(int rank, int n_slots, std::vector & out) const; + + // ── Builders ───────────────────────────────────────────────────────── + // expert e of every layer -> rank e % n_ranks. When replicate_hot > 0 and + // stats is non-null, the replicate_hot hottest experts of each layer + // (by stats) become kReplicated instead. + static bool build_uniform(int n_ranks, + const common::MoeHybridConfig & cfg, + int replicate_hot, + const common::MoeHybridRoutingStats * stats, + ClusterExpertPlacement & out, + std::string * err = nullptr); + + // Longest-processing-time greedy: per layer, experts sorted by hotness + // descending, each assigned to the rank with the smallest accumulated + // hotness, subject to |owned_count(rank) - n_expert/n_ranks| <= 1 so no + // rank's shard grows beyond one expert of the mean (memory balance). The + // replicate_hot hottest experts are marked kReplicated first. + static bool build_balanced(int n_ranks, + const common::MoeHybridRoutingStats & stats, + int replicate_hot, + ClusterExpertPlacement & out, + std::string * err = nullptr); + + // JSON: {"n_ranks":N,"n_layer":L,"n_expert":E,"n_expert_used":K, + // "replicate_hot":R,"owner":[[...E ints...] x L]} + bool save_json(const std::string & path, std::string * err = nullptr) const; + static bool load_json(const std::string & path, ClusterExpertPlacement & out, + std::string * err = nullptr); + + // Human-readable summary for the startup log: per-rank owned counts, + // replicated count, expected imbalance from stats if given. + std::string describe(const common::MoeHybridRoutingStats * stats = nullptr) const; +}; + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_head_backend.cpp b/server/src/cluster/cluster_head_backend.cpp new file mode 100644 index 000000000..89178568d --- /dev/null +++ b/server/src/cluster/cluster_head_backend.cpp @@ -0,0 +1,632 @@ +// cluster_head_backend.cpp - rank-0 ModelBackend decorator. + +#include "cluster/cluster_head_backend.h" + +#include "cluster/cluster_identity.h" +#include "deepseek4/deepseek4_cluster.h" + +#include +#include +#include +#include +#include +#include + +namespace luce::cluster { + +using common::DaemonIO; +using common::GenerateErrorCode; +using common::GenerateRequest; +using common::GenerateResult; +using common::ParkTarget; + +namespace { + +constexpr int kAbortCodeControl = 1; // control channel failure +constexpr int kAbortCodeCollective = 2; // RCCL failure +// Watchdog poll period. Heartbeats are 1 s with a 10 s timeout, so this only +// bounds how quickly the abort follows the detection, not the detection. +constexpr int kWatchdogPeriodMs = 250; +// Time between aborting the cluster and exiting, so the failed request's HTTP +// error response is written before the supervisor takes the process. +constexpr int kFatalExitGraceMs = 2000; +constexpr int kAbortCodeRequest = 3; // a rank failed the request + +// Self-test sizes: 1000 x 64 KiB, 43 x 16 MiB. +constexpr int kSelftestIters = 1000; +constexpr size_t kSelftestSmallN = 64u * 1024u / sizeof(float); +constexpr size_t kSelftestLargeN = 16u * 1024u * 1024u / sizeof(float); + +std::string hex64(uint64_t v) { + char buf[17]; + std::snprintf(buf, sizeof(buf), "%016llx", (unsigned long long) v); + return buf; +} + +uint64_t now_us() { + return (uint64_t) std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +// Bootstrap deadline: model loading on the workers can take minutes, so the +// accept phase waits much longer than a per-collective timeout. Workers retry +// their connect for ControlTimeouts::connect_retries * connect_ms (5 min by +// default), and this matches that window. +uint32_t bootstrap_deadline_ms(const ClusterConfig & cfg) { + return std::max(cfg.timeout_ms, 300000u); +} + +} // namespace + +// ─── Identity ─────────────────────────────────────────────────────────── + +uint64_t backend_placement_hash(const common::ModelBackend & backend) { + const auto * p = dynamic_cast(&backend); + return p ? p->cluster_placement_hash() : 0; +} + +// ─── Request mapping ──────────────────────────────────────────────────── + +// GenerateRequest -> RequestMsg. Fields the wire format does not carry (each +// would need a kProtocolVersion bump): +// SamplerCfg.rep_window, freq_pen, pres_pen -> affect the AR-vs-spec +// routing (needs_logit_processing) on workers; today the worker derives +// do_sample from temperature > 0 and never sees the penalties. Harmless +// because the head computes decode_mode from its own full sampler and +// the worker only ever narrows spec to AR, never the other way round. +// (snap_slot / snap_pos and the restore slot travel since protocol 2.) +// BudgetHook.hard_limit_remaining -> only close_token_ids travel +// (stop_token_ids); the worker gets hard_limit_remaining = 0, which is +// fine because the head's override arrives inside the token itself. +// hint_tokens / stall_* pointers, observer -> spec-only, not replicated. +RequestMsg ClusterHeadBackend::build_request_msg(uint64_t request_id, const GenerateRequest & req, + int restore_slot, int restore_kv_offset, + bool spec_available) { + RequestMsg m; + m.request_id = request_id; + m.prompt_tokens = req.prompt; + m.n_gen = req.n_gen; + m.max_ctx = 0; // not part of GenerateRequest; workers use their own --max-ctx + m.temperature = req.sampler.temp; + m.top_p = req.sampler.top_p; + m.top_k = req.sampler.top_k; + m.min_p = 0.0f; // SamplerCfg has no min_p + m.repeat_penalty = req.sampler.rep_pen; + m.seed = req.sampler.seed; + const bool budget_requires_ar = !req.budget_hook.close_token_ids.empty(); + const bool sampling_requires_ar = req.sampler.needs_logit_processing(); + m.decode_mode = (spec_available && !req.force_ar_decode && !budget_requires_ar && + !sampling_requires_ar) + ? DecodeMode::Speculative : DecodeMode::Autoregressive; + m.force_ar = req.force_ar_decode; + // Restore and inline save are independent: the HTTP layer may resume + // from one slot and check-point this request's longer prefix into + // another. Both have to travel, or the workers' slots drift away from the + // head's and a later restore silently resumes different KV. + m.restore_slot = restore_slot; + m.kv_offset = restore_slot >= 0 ? restore_kv_offset : 0; + m.snapshot_slot = req.snap_slot; + m.snapshot_pos = req.snap_pos; + m.stop_token_ids = req.budget_hook.close_token_ids; + return m; +} + +// ─── Lifecycle ────────────────────────────────────────────────────────── + +ClusterHeadBackend::ClusterHeadBackend(std::unique_ptr inner, + const ClusterConfig & cfg, + std::string model_path, + int device) + : inner_(std::move(inner)), cfg_(cfg), model_path_(std::move(model_path)), device_(device) {} + +ClusterHeadBackend::~ClusterHeadBackend() { + if (!shut_down_) shutdown(); +} + +bool ClusterHeadBackend::init() { + if (!inner_) { + std::fprintf(stderr, "[cluster] head: no inner backend\n"); + return false; + } + const std::string cfg_err = cfg_.validate(); + if (!cfg_err.empty() || !cfg_.is_head()) { + std::fprintf(stderr, "[cluster] head: invalid cluster config: %s\n", + cfg_err.empty() ? "rank is not 0" : cfg_err.c_str()); + return false; + } + // server_main already exported the NCCL_* environment before the + // factory ran; nothing to repeat here. + + std::string err; + ControlTimeouts timeouts; + timeouts.recv_ms = cfg_.timeout_ms; + control_.set_timeouts(timeouts); + if (!control_.listen(cfg_.head_host, cfg_.head_port, &err)) { + std::fprintf(stderr, "[cluster] head: listen %s:%d failed: %s\n", + cfg_.head_host.c_str(), cfg_.head_port, err.c_str()); + return false; + } + + // The placement hash comes from the inner backend's built placement + // (backend_factory called set_cluster(&cfg, nullptr) before init()). + const uint64_t placement_hash = backend_placement_hash(*inner_); + if (placement_hash == 0) { + std::fprintf(stderr, "[cluster] head: inner backend has no cluster placement; " + "set_cluster() must run before init()\n"); + control_.close(); + return false; + } + identity_ = cluster_identity(0, cfg_.size, model_path_, placement_hash); + std::fprintf(stderr, "[cluster] head rank 0/%d identity build=%s model=%s placement=%s\n", + cfg_.size, identity_.build_sha.c_str(), identity_.model_sha.c_str(), + hex64(identity_.placement_hash).c_str()); + // Launch scripts can wait for this exact stdout prefix before they start + // the workers. + std::printf("cluster: listening on %s:%d, waiting for %d worker(s)\n", + cfg_.head_host.c_str(), cfg_.head_port, cfg_.size - 1); + std::fflush(stdout); + + const uint64_t t0 = now_us(); + if (!control_.accept_workers(cfg_.size - 1, identity_, bootstrap_deadline_ms(cfg_), hellos_, &err)) { + std::fprintf(stderr, "[cluster] head: worker handshake failed: %s\n", err.c_str()); + control_.close(); + return false; + } + for (const HelloMsg & h : hellos_) { + std::fprintf(stderr, "[cluster] head: rank %d joined from %s\n", h.rank, h.hostname.c_str()); + } + std::printf("cluster: all %d ranks connected\n", cfg_.size); + std::fflush(stdout); + + RcclUniqueId uid{}; + if (!rccl_get_unique_id(uid, &err)) { + std::fprintf(stderr, "[cluster] head: rccl_get_unique_id failed: %s\n", err.c_str()); + fail_cluster("rccl_get_unique_id: " + err, kAbortCodeCollective, 0); + return false; + } + + WelcomeMsg welcome; + welcome.size = cfg_.size; + welcome.rccl_unique_id = uid; + welcome.placement_hash = identity_.placement_hash; + welcome.timeout_ms = cfg_.timeout_ms; + welcome.allreduce_dtype = (uint8_t) cfg_.allreduce_dtype; + welcome.shared_expert = (uint8_t) cfg_.shared_expert; + welcome.rank_hostnames.resize((size_t) cfg_.size); + welcome.rank_hostnames[0] = identity_.hostname; + for (const HelloMsg & h : hellos_) { + if (h.rank > 0 && h.rank < cfg_.size) welcome.rank_hostnames[(size_t) h.rank] = h.hostname; + } + if (!control_.broadcast(make_frame(welcome), &err)) { + std::fprintf(stderr, "[cluster] head: Welcome broadcast failed: %s\n", err.c_str()); + fail_cluster("welcome: " + err, kAbortCodeControl, 0); + return false; + } + + RcclCommInit init; + init.rank = 0; + init.size = cfg_.size; + init.unique_id = uid; + init.device = device_; // local HIP ordinal the inner backend runs on + init.timeout_ms = cfg_.timeout_ms; + init.blocking = false; + comm_ = create_rccl_cluster_comm(init, &err); + if (!comm_) { + std::fprintf(stderr, "[cluster] head: RCCL communicator failed: %s\n", err.c_str()); + fail_cluster("rccl init: " + err, kAbortCodeCollective, 0); + return false; + } + std::fprintf(stderr, "[cluster] head: %s communicator up, %d ranks, bootstrap %.2f s\n", + comm_->backend_name(), comm_->size(), (now_us() - t0) / 1e6); + std::printf("cluster: rccl communicator ready\n"); + std::fflush(stdout); + + if (cfg_.selftest) { + const bool ok = run_cluster_selftest(*comm_, init.device, kSelftestIters, kSelftestSmallN, + kSelftestLargeN, &err); + std::fprintf(stderr, "[cluster] head: selftest %s%s%s\n", ok ? "PASSED" : "FAILED", + err.empty() ? "" : ": ", err.c_str()); + ShutdownMsg bye; + bye.reason = ok ? 0 : 1; + control_.broadcast(make_frame(bye), nullptr); + control_.close(); + comm_.reset(); + std::fflush(nullptr); + std::exit(ok ? 0 : 1); + } + + hooks_ = std::make_unique(control_); + + // Second set_cluster call: attaches the communicator to the + // runtime the factory created before init(); placement-defining fields + // must be unchanged, which they are (same ClusterConfig). + common::IClusterParticipant * part = cluster_participant(); + if (!part || !part->cluster_attach(&cfg_, comm_.get())) { + fail_cluster("cluster_attach(comm) refused by the backend", kAbortCodeCollective, 0); + return false; + } + part->cluster_set_hooks(hooks_.get()); + + if (!comm_->barrier(&err)) { + std::fprintf(stderr, "[cluster] head: post-init barrier failed: %s\n", err.c_str()); + fail_cluster("barrier: " + err, kAbortCodeCollective, 0); + return false; + } + initialized_ = true; + start_watchdog(); + return true; +} + +void ClusterHeadBackend::print_ready_banner() const { + inner_->print_ready_banner(); + std::printf("[cluster] head ready: size=%d workers=%d rccl=%s\n", + cfg_.size, control_.n_workers(), rccl_version_string().c_str()); + std::fflush(stdout); +} + +// ─── Failure ──────────────────────────────────────────────────────────── + +void ClusterHeadBackend::fail_cluster(const std::string & reason, int code, uint64_t request_id) { + if (fatal_) return; + fatal_ = true; + std::fprintf(stderr, "[cluster] head: FATAL (%d): %s\n", code, reason.c_str()); + AbortMsg abort; + abort.rank = 0; + abort.code = code; + abort.request_id = request_id; + abort.reason = reason; + control_.broadcast(make_frame(abort), nullptr); + if (comm_) comm_->abort(); + control_.close(); + std::fprintf(stderr, "[cluster] head: exiting so the supervisor restarts all ranks\n"); + std::fflush(nullptr); + std::_Exit(3); +} + +// ─── Peer-failure watchdog ────────────────────────────────────────────── + +void ClusterHeadBackend::request_fatal_exit(const std::string & reason, int code, + uint64_t request_id) { + if (exit_scheduled_.exchange(true)) return; + std::fprintf(stderr, "[cluster] head: FATAL (%d): %s\n", code, reason.c_str()); + fatal_ = true; + // Abort the communicator FIRST: that is what makes a collective that is + // already enqueued on the GPU return an error instead of waiting for a + // rank that will never arrive. Without it the forward never returns. + if (comm_) comm_->abort(); + AbortMsg abort; + abort.rank = 0; + abort.code = code; + abort.request_id = request_id; + abort.reason = reason; + control_.broadcast(make_frame(abort), nullptr); + // Give the request thread time to unwind and the HTTP layer time to write + // the error response, then hand over to the supervisor. Detached because + // this may run from the watchdog itself. + std::thread([reason]() { + std::this_thread::sleep_for(std::chrono::milliseconds(kFatalExitGraceMs)); + std::fprintf(stderr, + "[cluster] head: exiting so the supervisor restarts all ranks (%s)\n", + reason.c_str()); + std::fflush(nullptr); + std::_Exit(3); + }).detach(); +} + +void ClusterHeadBackend::watchdog_loop() { + while (!watchdog_stop_.load()) { + std::this_thread::sleep_for(std::chrono::milliseconds(kWatchdogPeriodMs)); + if (watchdog_stop_.load()) break; + // Only meaningful while a forward can be blocked in a collective. + if (!request_in_flight_.load() || fatal_) continue; + + std::string err; + const std::vector dead = control_.dead_ranks(); + const bool async_failed = comm_ && comm_->async_error(&err); + if (dead.empty() && !async_failed) continue; + + std::string reason; + if (!dead.empty()) { + reason = "rank"; + for (int r : dead) reason += " " + std::to_string(r); + reason += " stopped responding during a request"; + } else { + reason = "communicator error during a request: " + err; + } + request_fatal_exit(reason, dead.empty() ? kAbortCodeCollective : kAbortCodeControl, + last_report_.request_id); + } +} + +void ClusterHeadBackend::start_watchdog() { + if (watchdog_.joinable()) return; + watchdog_stop_.store(false); + watchdog_ = std::thread([this] { watchdog_loop(); }); +} + +void ClusterHeadBackend::stop_watchdog() { + watchdog_stop_.store(true); + if (watchdog_.joinable()) watchdog_.join(); +} + +// ─── Generation ───────────────────────────────────────────────────────── + +bool ClusterHeadBackend::spec_available() const { + // Rank 0 decides for the whole cluster whether a request decodes + // speculatively. A worker that happens to lack a drafter must not quietly + // fall back to AR while the head speculates, so it rejects the request + // instead (cluster_worker_main.cpp). + const common::IClusterParticipant * part = cluster_participant(); + return part && part->cluster_spec_decode_ready(); +} + +GenerateResult ClusterHeadBackend::generate_impl(const GenerateRequest & req, const DaemonIO & io) { + return run_request(req, io, -1); +} + +GenerateResult ClusterHeadBackend::restore_and_generate_impl(int slot, const GenerateRequest & req, + const DaemonIO & io) { + return run_request(req, io, slot); +} + +GenerateResult ClusterHeadBackend::run_request(const GenerateRequest & req, const DaemonIO & io, + int restore_slot) { + GenerateResult result; + if (!initialized_ || fatal_) { + result.fail(GenerateErrorCode::DecodeFailed, + fatal_ ? "cluster aborted" : "cluster not initialized"); + return result; + } + const uint64_t request_id = next_request_id_++; + const int restore_kv_offset = restore_slot >= 0 ? inner_->snapshot_cur_pos(restore_slot) : 0; + const RequestMsg msg = build_request_msg(request_id, req, restore_slot, restore_kv_offset, + spec_available()); + + hooks_->set_current_request(request_id); + hooks_->reset_counters(); + if (comm_) comm_->reset_stats(); + last_report_ = ClusterRequestReport{}; + last_report_.request_id = request_id; + last_report_.size = cfg_.size; + + // A failure from here on ends the cluster through request_fatal_exit, which + // exits only after the HTTP layer has written this request's error. + std::string err; + if (!control_.broadcast(make_frame(msg), &err)) { + request_fatal_exit("request broadcast: " + err, kAbortCodeControl, request_id); + result.fail(GenerateErrorCode::DecodeFailed, "cluster: request broadcast failed: " + err); + return result; + } + + // Rank 0 decodes in the mode it broadcast, as the workers do; on its own + // the backend would also speculate on sampled and budgeted requests. + std::optional ar_req; + if (request_decodes_ar(msg) && !req.force_ar_decode) { + ar_req.emplace(req); + ar_req->force_ar_decode = true; + } + const GenerateRequest & local_req = ar_req ? *ar_req : req; + + const auto compute_t0 = std::chrono::steady_clock::now(); + request_in_flight_.store(true); + result = restore_slot >= 0 ? inner_->restore_and_generate_impl(restore_slot, local_req, io) + : inner_->generate_impl(local_req, io); + request_in_flight_.store(false); + // Rank 0's own prefill+decode wall time, the counterpart of the workers' + // compute_us. Comparing them across ranks is how an imbalanced expert + // placement or a slow node shows up. + head_compute_us_ = (uint64_t) std::chrono::duration_cast( + std::chrono::steady_clock::now() - compute_t0).count(); + + // The decode loop reports channel failures as DecodeFailed with the + // message already logged; a dead worker also shows up here. + if (!result.ok() && (!control_.all_alive() || (comm_ && comm_->async_error(&err)))) { + request_fatal_exit("request " + std::to_string(request_id) + " failed: " + + std::string(result.error_detail()) + (err.empty() ? "" : " / " + err), + kAbortCodeRequest, request_id); + return result; + } + + RequestEndMsg end; + end.request_id = request_id; + end.steps = (uint32_t) result.tokens.size(); + end.status = result.ok() ? 0 : (int32_t) result.error->code; + if (!control_.broadcast(make_frame(end), &err)) { + request_fatal_exit("RequestEnd broadcast: " + err, kAbortCodeControl, request_id); + if (result.ok()) result.fail(GenerateErrorCode::DecodeFailed, "cluster: RequestEnd failed"); + return result; + } + if (!gather_reports(request_id, end.steps)) { + request_fatal_exit("RequestReport gather: " + last_report_.error, kAbortCodeControl, + request_id); + if (result.ok()) result.fail(GenerateErrorCode::DecodeFailed, "cluster: report gather failed"); + return result; + } + return result; +} + +bool ClusterHeadBackend::gather_reports(uint64_t request_id, uint32_t steps) { + last_report_.steps = steps; + last_report_.ranks.assign((size_t) cfg_.size, RequestReportMsg{}); + RequestReportMsg & head = last_report_.ranks[0]; + head.request_id = request_id; + head.rank = 0; + head.steps = steps; + head.compute_us = head_compute_us_; + if (comm_) { + const ClusterCommStats s = comm_->stats(); + head.allreduce_calls = s.allreduce_calls; + head.allreduce_bytes = s.allreduce_bytes; + head.allreduce_wait_us = s.allreduce_wait_us; + } + head.ctrl_wait_us = hooks_->counters().ctrl_wait_us; + const uint64_t head_device_bytes = cluster_device_bytes_in_use(device_); + if (head_device_bytes > peak_device_bytes_) peak_device_bytes_ = head_device_bytes; + head.peak_device_bytes = peak_device_bytes_; + last_report_.head_ctrl_wait_us = head.ctrl_wait_us; + + std::vector frames; + std::string err; + if (!control_.gather(MsgType::RequestReport, cfg_.timeout_ms, frames, &err)) { + last_report_.error = err; + return false; + } + for (size_t i = 0; i < frames.size(); i++) { + RequestReportMsg r; + if (!parse_frame(frames[i], r)) { + last_report_.error = "malformed RequestReport from rank " + std::to_string(i + 1); + return false; + } + if (r.request_id != request_id) { + last_report_.error = "RequestReport from rank " + std::to_string(r.rank) + + " is for request " + std::to_string(r.request_id); + return false; + } + if (r.rank <= 0 || r.rank >= cfg_.size) { + last_report_.error = "RequestReport with bad rank " + std::to_string(r.rank); + return false; + } + last_report_.ranks[(size_t) r.rank] = r; + } + last_report_.complete = true; + if (cluster_env_trace()) { + for (const RequestReportMsg & r : last_report_.ranks) { + std::fprintf(stderr, + "[cluster] req=%llu rank=%d steps=%u compute_us=%llu allreduce=%llu calls/%llu B/%llu us " + "ctrl_wait_us=%llu\n", + (unsigned long long) request_id, r.rank, r.steps, (unsigned long long) r.compute_us, + (unsigned long long) r.allreduce_calls, (unsigned long long) r.allreduce_bytes, + (unsigned long long) r.allreduce_wait_us, (unsigned long long) r.ctrl_wait_us); + } + } + return true; +} + +// ─── Replicated mutations ─────────────────────────────────────────────── + +bool ClusterHeadBackend::broadcast_op(BackendOpKind kind, std::vector args) { + if (!initialized_ || fatal_) return false; + BackendOpMsg op; + op.request_id = 0; // process-level; ops never run while a request is in flight (single stream) + op.kind = kind; + op.args = std::move(args); + std::string err; + if (!hooks_->backend_op(op, &err)) { + request_fatal_exit("BackendOp broadcast: " + err, kAbortCodeControl, 0); + return false; + } + return true; +} + +bool ClusterHeadBackend::park(ParkTarget target) { + if (!broadcast_op(BackendOpKind::Park, {(int64_t) target})) return false; + return inner_->park(target); +} + +bool ClusterHeadBackend::unpark(ParkTarget target) { + if (!broadcast_op(BackendOpKind::Unpark, {(int64_t) target})) return false; + return inner_->unpark(target); +} + +bool ClusterHeadBackend::is_target_parked() const { return inner_->is_target_parked(); } + +bool ClusterHeadBackend::snapshot_save(int slot) { + if (!broadcast_op(BackendOpKind::SnapshotSave, {slot})) return false; + if (inner_->snapshot_save(slot)) return true; + // The workers have taken the snapshot and this rank has not. Leaving it + // there would make a later restore resume from a slot only some ranks + // hold, which diverges silently. Drop it everywhere instead. + std::fprintf(stderr, + "[cluster] head: snapshot_save(%d) failed here; freeing the slot on " + "every rank so no rank restores from it\n", slot); + broadcast_op(BackendOpKind::SnapshotFree, {slot}); + return false; +} + +void ClusterHeadBackend::snapshot_free(int slot) { + broadcast_op(BackendOpKind::SnapshotFree, {slot}); + inner_->snapshot_free(slot); +} + +bool ClusterHeadBackend::snapshot_used(int slot) const { return inner_->snapshot_used(slot); } +int ClusterHeadBackend::snapshot_cur_pos(int slot) const { return inner_->snapshot_cur_pos(slot); } + +common::ModelBackend::SnapshotRef ClusterHeadBackend::snapshot_ref(int slot) const { + return inner_->snapshot_ref(slot); +} + +bool ClusterHeadBackend::snapshot_adopt(int slot, ggml_context * ctx, ggml_backend_buffer_t buf, + int cur_pos, int32_t last_tok) { + // Adopting a deserialized snapshot on the head only would leave the + // workers without it; refusing keeps every rank identical. + (void) slot; (void) ctx; (void) buf; (void) cur_pos; (void) last_tok; + std::fprintf(stderr, "[cluster] head: snapshot_adopt is not replicated; refusing\n"); + return false; +} + +common::ModelBackend::CompressResult ClusterHeadBackend::compress(const CompressRequest & req) { + return inner_->compress(req); +} + +std::vector ClusterHeadBackend::compress_batch( + const std::vector & requests) { + return inner_->compress_batch(requests); +} + +bool ClusterHeadBackend::handle_compress(const std::string & line, const DaemonIO & io) { + // DeepSeek4Backend::handle_compress is unsupported and mutates nothing, + // so no BackendOp is sent; revisit if PFlash lands for DS4. + return inner_->handle_compress(line, io); +} + +void ClusterHeadBackend::free_drafter() { + broadcast_op(BackendOpKind::FreeDrafter, {}); + inner_->free_drafter(); +} + +bool ClusterHeadBackend::try_handle_command(const std::string & line, const DaemonIO & io) { + return inner_->try_handle_command(line, io); +} + +bool ClusterHeadBackend::supports_dflash_spec_decode() const { + return inner_->supports_dflash_spec_decode(); +} +common::DFlashTarget * ClusterHeadBackend::dflash_target() { return inner_->dflash_target(); } +void ClusterHeadBackend::release_scratch() { inner_->release_scratch(); } +bool ClusterHeadBackend::supports_remote_draft() const { return inner_->supports_remote_draft(); } +bool ClusterHeadBackend::supports_kvflash() const { return inner_->supports_kvflash(); } +bool ClusterHeadBackend::supports_mixed_backend_layer_split() const { + return inner_->supports_mixed_backend_layer_split(); +} +bool ClusterHeadBackend::set_routing_collector(common::MoeRoutingCollector * collector) { + return inner_->set_routing_collector(collector); +} +const common::MoeHybridRoutingStats * ClusterHeadBackend::get_routing_stats() const { + return inner_->get_routing_stats(); +} +bool ClusterHeadBackend::spark_wants_bootstrap() const { return inner_->spark_wants_bootstrap(); } +bool ClusterHeadBackend::spark_bootstrap_finalize(const std::string & profile_path) { + return inner_->spark_bootstrap_finalize(profile_path); +} + +void ClusterHeadBackend::shutdown() { + if (shut_down_) return; + shut_down_ = true; + stop_watchdog(); + if (initialized_ && !fatal_) { + ShutdownMsg bye; + bye.reason = 0; + control_.broadcast(make_frame(bye), nullptr); + } + if (inner_) { + if (common::IClusterParticipant * part = cluster_participant()) { + part->cluster_set_hooks(nullptr); + } + inner_->shutdown(); + } + control_.close(); + comm_.reset(); + hooks_.reset(); +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_head_backend.h b/server/src/cluster/cluster_head_backend.h new file mode 100644 index 000000000..1ba943817 --- /dev/null +++ b/server/src/cluster/cluster_head_backend.h @@ -0,0 +1,180 @@ +// cluster_head_backend.h - rank-0 ModelBackend decorator for cluster runs. +// +// The HTTP server talks to one ModelBackend. In a cluster that backend is +// this decorator: it owns the real backend, the control channel to +// the workers, the RCCL communicator and the HeadHooks. Before every +// generate() it broadcasts a RequestMsg so the workers run the identical +// call; every state-changing ModelBackend virtual (snapshot save/free, +// park/unpark, free_drafter) is broadcast as a BackendOpMsg before it is +// executed locally; after generate() it broadcasts RequestEnd and gathers +// one RequestReport per worker. +// +// Failure policy: any control-channel or collective failure broadcasts +// Abort, aborts the communicator and ends the process with a non-zero +// status so a supervisor restarts all ranks together. + +#pragma once + +#include "cluster/cluster_comm.h" +#include "cluster/cluster_config.h" +#include "cluster/cluster_control.h" +#include "cluster/cluster_decision_hooks.h" +#include "cluster/cluster_protocol.h" +#include "common/cluster_participant.h" +#include "common/model_backend.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace luce::cluster { + +// Per-request cluster telemetry merged on the head. ranks[0] is the head +// itself (only steps / ctrl_wait_us / allreduce_* from its own comm stats are +// filled), ranks[r] for r >= 1 is the worker's RequestReportMsg. +struct ClusterRequestReport { + uint64_t request_id = 0; + int size = 0; + uint32_t steps = 0; + bool complete = false; // every worker reported + std::vector ranks; + uint64_t head_ctrl_wait_us = 0; + std::string error; // non-empty when gathering failed +}; + +// Placement hash for HelloMsg: the inner backend's built placement +// (Ds4ClusterRuntime::placement.hash()) or 0 when no cluster runtime exists. +uint64_t backend_placement_hash(const common::ModelBackend & backend); + +class ClusterHeadBackend final : public common::ModelBackend { +public: + // `device` is the local HIP ordinal the inner backend runs on + // (BackendArgs::device.gpu); RCCL binds the communicator to it. + ClusterHeadBackend(std::unique_ptr inner, + const ClusterConfig & cfg, + std::string model_path, + int device); + ~ClusterHeadBackend() override; + + ClusterHeadBackend(const ClusterHeadBackend &) = delete; + ClusterHeadBackend & operator=(const ClusterHeadBackend &) = delete; + + // Bootstrap: listen, accept size-1 workers, distribute the RCCL unique + // id, create the communicator, attach it to the inner backend + // (set_cluster second call) and install the hooks. The inner backend + // must already be init()ed with set_cluster(&cfg, nullptr) called before + // its init() (backend_factory does both) so the Hello carries the real + // placement hash. When cfg.selftest is set this runs run_cluster_selftest + // and exits the process with its status (the operator asked for exactly + // that). + bool init(); + + // ── ModelBackend ───────────────────────────────────────────────── + void print_ready_banner() const override; + bool park(common::ParkTarget target) override; + bool unpark(common::ParkTarget target) override; + bool is_target_parked() const override; + common::GenerateResult generate_impl(const common::GenerateRequest & req, + const common::DaemonIO & io) override; + common::SeqEngine * seq_engine() override { return nullptr; } + MemoryReport memory_report() const override { return inner_->memory_report(); } + int snapshot_granularity() const override { return inner_->snapshot_granularity(); } + size_t snapshot_bytes_estimate(int tokens) const override { + return inner_->snapshot_bytes_estimate(tokens); + } + bool snapshot_save(int slot) override; + void snapshot_free(int slot) override; + bool snapshot_used(int slot) const override; + int snapshot_cur_pos(int slot) const override; + common::GenerateResult restore_and_generate_impl(int slot, + const common::GenerateRequest & req, + const common::DaemonIO & io) override; + SnapshotRef snapshot_ref(int slot) const override; + bool snapshot_adopt(int slot, ggml_context * ctx, ggml_backend_buffer_t buf, int cur_pos, + int32_t last_tok) override; + CompressResult compress(const CompressRequest & req) override; + std::vector compress_batch(const std::vector & requests) override; + bool handle_compress(const std::string & line, const common::DaemonIO & io) override; + void free_drafter() override; + bool try_handle_command(const std::string & line, const common::DaemonIO & io) override; + bool supports_dflash_spec_decode() const override; + common::DFlashTarget * dflash_target() override; + void release_scratch() override; + bool supports_remote_draft() const override; + bool supports_kvflash() const override; + bool supports_mixed_backend_layer_split() const override; + bool set_routing_collector(common::MoeRoutingCollector * collector) override; + const common::MoeHybridRoutingStats * get_routing_stats() const override; + bool spark_wants_bootstrap() const override; + bool spark_bootstrap_finalize(const std::string & profile_path) override; + void shutdown() override; + + // Maps a GenerateRequest (plus optional restore slot) onto the wire + // descriptor. Exposed for the unit test; see the .cpp for the exact + // field mapping and the fields the protocol cannot carry yet. + static RequestMsg build_request_msg(uint64_t request_id, const common::GenerateRequest & req, + int restore_slot, int restore_kv_offset, + bool spec_available); + +private: + common::GenerateResult run_request(const common::GenerateRequest & req, + const common::DaemonIO & io, + int restore_slot); + bool broadcast_op(BackendOpKind kind, std::vector args); + bool gather_reports(uint64_t request_id, uint32_t steps); + + // A dead peer used to be noticed only between forwards. With the + // in-graph all-reduce nothing on the host waits for a + // collective any more, so a rank that dies mid-forward leaves this one + // blocked inside HIP forever (observed: a killed worker hung the head + // indefinitely, no HTTP answer). This thread polls the control channel's + // heartbeats and the communicator's async error while a request is in + // flight and calls ncclCommAbort, which makes the pending collective + // return an error so the forward can fail and the client gets an answer. + void watchdog_loop(); + void start_watchdog(); + void stop_watchdog(); + // Abort the cluster, let the in-flight request fail, and exit after a + // grace period so the HTTP response is written before the supervisor + // takes over. Idempotent. + void request_fatal_exit(const std::string & reason, int code, uint64_t request_id); + // Broadcast Abort, abort the communicator, mark fatal, maybe exit at once. + // For start-up failures; a request thread uses request_fatal_exit. + void fail_cluster(const std::string & reason, int code, uint64_t request_id); + bool spec_available() const; + + // The inner backend seen as a cluster participant. Null would mean the + // factory handed us an architecture the gates should have rejected. + common::IClusterParticipant * cluster_participant() const { + return dynamic_cast(inner_.get()); + } + + std::unique_ptr inner_; + ClusterConfig cfg_; + std::string model_path_; + int device_ = 0; + // High-water mark of device memory in use on rank 0, sampled at request + // ends; the workers keep the same mark and ship it in RequestReport. + uint64_t peak_device_bytes_ = 0; + uint64_t head_compute_us_ = 0; + std::thread watchdog_; + std::atomic watchdog_stop_{false}; + std::atomic request_in_flight_{false}; + std::atomic exit_scheduled_{false}; + ClusterHeadControl control_; + std::unique_ptr comm_; + std::unique_ptr hooks_; + std::vector hellos_; + HelloMsg identity_; + uint64_t next_request_id_ = 1; + ClusterRequestReport last_report_; + bool initialized_ = false; + bool fatal_ = false; + bool shut_down_ = false; +}; + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_identity.h b/server/src/cluster/cluster_identity.h new file mode 100644 index 000000000..017e70a31 --- /dev/null +++ b/server/src/cluster/cluster_identity.h @@ -0,0 +1,57 @@ +// cluster_identity.h - the HelloMsg identity every rank presents. +// +// Head and workers must compute these fields with the same code, otherwise +// the head's accept_workers() refuses ranks that are in fact identical. +// Header-only so server_main.cpp (selftest), ClusterHeadBackend and +// run_cluster_worker share one definition without a link dependency. +// +// build_sha : LUCE_BUILD_SHA compile-time define (CMake sets it +// from `git describe` when LUCE_CLUSTER=ON), else +// "unknown" - then only model and placement are compared. +// model_sha : fnv1a64 hex of the model path. All ranks must load the +// same file from the same path (the launcher enforces the +// same mount layout); hashing the 98 GB file itself at +// startup is not an option. +// placement_hash : ClusterExpertPlacement::hash() of the rank's built +// placement, or 0 before a placement exists (selftest). + +#pragma once + +#include "cluster/cluster_control.h" +#include "cluster/cluster_protocol.h" + +#include +#include +#include + +namespace luce::cluster { + +inline std::string cluster_build_sha() { +#ifdef LUCE_BUILD_SHA + return LUCE_BUILD_SHA; +#else + return "unknown"; +#endif +} + +inline std::string cluster_model_sha(const std::string & model_path) { + char hex[32]; + std::snprintf(hex, sizeof(hex), "%016llx", + (unsigned long long) fnv1a64(model_path.data(), model_path.size())); + return hex; +} + +inline HelloMsg cluster_identity(int rank, int size, const std::string & model_path, + uint64_t placement_hash) { + HelloMsg h; + h.rank = rank; + h.size = size; + h.protocol_version = kProtocolVersion; + h.build_sha = cluster_build_sha(); + h.model_sha = cluster_model_sha(model_path); + h.placement_hash = placement_hash; + h.hostname = local_hostname(rank); + return h; +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_protocol.cpp b/server/src/cluster/cluster_protocol.cpp new file mode 100644 index 000000000..a7e176dac --- /dev/null +++ b/server/src/cluster/cluster_protocol.cpp @@ -0,0 +1,452 @@ +// cluster_protocol.cpp - hand-rolled little-endian wire encoding for the +// control channel (cluster_protocol.h). No exceptions: every decode returns +// false on truncation, over-long lengths or malformed enums. + +#include "cluster/cluster_protocol.h" + +#include + +namespace luce::cluster { + +// ─── ByteWriter ───────────────────────────────────────────────────────── + +void ByteWriter::u8(uint8_t v) { + buf_.push_back(v); +} + +void ByteWriter::u16(uint16_t v) { + buf_.push_back((uint8_t)(v & 0xFFu)); + buf_.push_back((uint8_t)((v >> 8) & 0xFFu)); +} + +void ByteWriter::u32(uint32_t v) { + for (int i = 0; i < 4; ++i) { + buf_.push_back((uint8_t)((v >> (8 * i)) & 0xFFu)); + } +} + +void ByteWriter::u64(uint64_t v) { + for (int i = 0; i < 8; ++i) { + buf_.push_back((uint8_t)((v >> (8 * i)) & 0xFFu)); + } +} + +void ByteWriter::i32(int32_t v) { + u32((uint32_t)v); +} + +void ByteWriter::i64(int64_t v) { + u64((uint64_t)v); +} + +void ByteWriter::f32(float v) { + uint32_t bits = 0; + static_assert(sizeof(bits) == sizeof(v), "float must be 32-bit"); + std::memcpy(&bits, &v, sizeof(bits)); + u32(bits); +} + +void ByteWriter::str(const std::string & s) { + u32((uint32_t)s.size()); + buf_.insert(buf_.end(), s.begin(), s.end()); +} + +void ByteWriter::bytes(const void * data, size_t n) { + u32((uint32_t)n); + if (n > 0) { + const uint8_t * p = (const uint8_t *)data; + buf_.insert(buf_.end(), p, p + n); + } +} + +void ByteWriter::vec_i32(const std::vector & v) { + u32((uint32_t)v.size()); + for (int32_t x : v) i32(x); +} + +void ByteWriter::vec_i64(const std::vector & v) { + u32((uint32_t)v.size()); + for (int64_t x : v) i64(x); +} + +void ByteWriter::vec_u64(const std::vector & v) { + u32((uint32_t)v.size()); + for (uint64_t x : v) u64(x); +} + +// ─── ByteReader ───────────────────────────────────────────────────────── + +bool ByteReader::u8(uint8_t & v) { + if (failed_ || size_ - pos_ < 1) { failed_ = true; return false; } + v = data_[pos_++]; + return true; +} + +bool ByteReader::u16(uint16_t & v) { + if (failed_ || size_ - pos_ < 2) { failed_ = true; return false; } + v = (uint16_t)data_[pos_] | (uint16_t)((uint16_t)data_[pos_ + 1] << 8); + pos_ += 2; + return true; +} + +bool ByteReader::u32(uint32_t & v) { + if (failed_ || size_ - pos_ < 4) { failed_ = true; return false; } + uint32_t r = 0; + for (int i = 0; i < 4; ++i) { + r |= (uint32_t)data_[pos_ + (size_t)i] << (8 * i); + } + pos_ += 4; + v = r; + return true; +} + +bool ByteReader::u64(uint64_t & v) { + if (failed_ || size_ - pos_ < 8) { failed_ = true; return false; } + uint64_t r = 0; + for (int i = 0; i < 8; ++i) { + r |= (uint64_t)data_[pos_ + (size_t)i] << (8 * i); + } + pos_ += 8; + v = r; + return true; +} + +bool ByteReader::i32(int32_t & v) { + uint32_t u = 0; + if (!u32(u)) return false; + v = (int32_t)u; + return true; +} + +bool ByteReader::i64(int64_t & v) { + uint64_t u = 0; + if (!u64(u)) return false; + v = (int64_t)u; + return true; +} + +bool ByteReader::f32(float & v) { + uint32_t bits = 0; + if (!u32(bits)) return false; + std::memcpy(&v, &bits, sizeof(v)); + return true; +} + +bool ByteReader::str(std::string & s, size_t max_len) { + uint32_t n = 0; + if (!u32(n)) return false; + if (n > max_len || size_ - pos_ < n) { failed_ = true; return false; } + s.assign((const char *)data_ + pos_, n); + pos_ += n; + return true; +} + +bool ByteReader::bytes(std::vector & out, size_t max_len) { + uint32_t n = 0; + if (!u32(n)) return false; + if (n > max_len || size_ - pos_ < n) { failed_ = true; return false; } + out.assign(data_ + pos_, data_ + pos_ + n); + pos_ += n; + return true; +} + +bool ByteReader::vec_i32(std::vector & v, size_t max_count) { + uint32_t n = 0; + if (!u32(n)) return false; + if (n > max_count || size_ - pos_ < (size_t)n * 4) { failed_ = true; return false; } + v.resize(n); + for (uint32_t i = 0; i < n; ++i) { + if (!i32(v[i])) return false; + } + return true; +} + +bool ByteReader::vec_i64(std::vector & v, size_t max_count) { + uint32_t n = 0; + if (!u32(n)) return false; + if (n > max_count || size_ - pos_ < (size_t)n * 8) { failed_ = true; return false; } + v.resize(n); + for (uint32_t i = 0; i < n; ++i) { + if (!i64(v[i])) return false; + } + return true; +} + +bool ByteReader::vec_u64(std::vector & v, size_t max_count) { + uint32_t n = 0; + if (!u32(n)) return false; + if (n > max_count || size_ - pos_ < (size_t)n * 8) { failed_ = true; return false; } + v.resize(n); + for (uint32_t i = 0; i < n; ++i) { + if (!u64(v[i])) return false; + } + return true; +} + +// ─── Header ───────────────────────────────────────────────────────────── + +void encode_header(const FrameHeader & h, uint8_t out[kFrameHeaderBytes]) { + ByteWriter w; + w.u32(h.magic); + w.u16(h.version); + w.u16(h.type); + w.u32(h.payload_len); + std::memcpy(out, w.data().data(), kFrameHeaderBytes); +} + +bool decode_header(const uint8_t in[kFrameHeaderBytes], FrameHeader & h, std::string * err) { + ByteReader r(in, kFrameHeaderBytes); + FrameHeader tmp; + if (!r.u32(tmp.magic) || !r.u16(tmp.version) || !r.u16(tmp.type) || !r.u32(tmp.payload_len)) { + if (err) *err = "frame header truncated"; + return false; + } + if (tmp.magic != kFrameMagic) { + if (err) *err = "bad frame magic"; + return false; + } + if (tmp.version != kProtocolVersion) { + if (err) { + *err = "protocol version mismatch (peer " + std::to_string(tmp.version) + + ", local " + std::to_string(kProtocolVersion) + ")"; + } + return false; + } + if (tmp.payload_len > kMaxPayloadBytes) { + if (err) *err = "frame payload too large (" + std::to_string(tmp.payload_len) + " bytes)"; + return false; + } + h = tmp; + return true; +} + +const char * msg_type_name(MsgType type) { + switch (type) { + case MsgType::Hello: return "Hello"; + case MsgType::Welcome: return "Welcome"; + case MsgType::Heartbeat: return "Heartbeat"; + case MsgType::Abort: return "Abort"; + case MsgType::Request: return "Request"; + case MsgType::Decision: return "Decision"; + case MsgType::Draft: return "Draft"; + case MsgType::Accept: return "Accept"; + case MsgType::RequestReport: return "RequestReport"; + case MsgType::BackendOp: return "BackendOp"; + case MsgType::RequestEnd: return "RequestEnd"; + case MsgType::Shutdown: return "Shutdown"; + } + return "Unknown"; +} + +// ─── Messages ─────────────────────────────────────────────────────────── + +void HelloMsg::encode(ByteWriter & w) const { + w.i32(rank); + w.i32(size); + w.u16(protocol_version); + w.str(build_sha); + w.str(model_sha); + w.u64(placement_hash); + w.str(hostname); +} + +bool HelloMsg::decode(ByteReader & r, HelloMsg & out) { + return r.i32(out.rank) && r.i32(out.size) && r.u16(out.protocol_version) && + r.str(out.build_sha, 4096) && r.str(out.model_sha, 4096) && + r.u64(out.placement_hash) && r.str(out.hostname, 1024); +} + +void WelcomeMsg::encode(ByteWriter & w) const { + w.i32(size); + w.bytes(rccl_unique_id.data(), rccl_unique_id.size()); + w.u64(placement_hash); + w.u32(timeout_ms); + w.u8(allreduce_dtype); + w.u8(shared_expert); + w.u32((uint32_t)rank_hostnames.size()); + for (const std::string & h : rank_hostnames) w.str(h); +} + +bool WelcomeMsg::decode(ByteReader & r, WelcomeMsg & out) { + if (!r.i32(out.size)) return false; + std::vector id; + if (!r.bytes(id, kRcclUniqueIdBytes)) return false; + if (id.size() != kRcclUniqueIdBytes) return false; + std::memcpy(out.rccl_unique_id.data(), id.data(), kRcclUniqueIdBytes); + if (!r.u64(out.placement_hash) || !r.u32(out.timeout_ms) || !r.u8(out.allreduce_dtype) || + !r.u8(out.shared_expert)) { + return false; + } + uint32_t n = 0; + if (!r.u32(n)) return false; + if (n > 256) return false; + out.rank_hostnames.resize(n); + for (uint32_t i = 0; i < n; ++i) { + if (!r.str(out.rank_hostnames[i], 1024)) return false; + } + return true; +} + +void HeartbeatMsg::encode(ByteWriter & w) const { + w.i32(rank); + w.u64(mono_us); +} + +bool HeartbeatMsg::decode(ByteReader & r, HeartbeatMsg & out) { + return r.i32(out.rank) && r.u64(out.mono_us); +} + +void AbortMsg::encode(ByteWriter & w) const { + w.i32(rank); + w.i32(code); + w.u64(request_id); + w.str(reason); +} + +bool AbortMsg::decode(ByteReader & r, AbortMsg & out) { + return r.i32(out.rank) && r.i32(out.code) && r.u64(out.request_id) && + r.str(out.reason, 1 << 16); +} + +void RequestMsg::encode(ByteWriter & w) const { + w.u64(request_id); + w.vec_i32(prompt_tokens); + w.i32(n_gen); + w.i32(max_ctx); + w.f32(temperature); + w.f32(top_p); + w.i32(top_k); + w.f32(min_p); + w.f32(repeat_penalty); + w.u64(seed); + w.u8((uint8_t)decode_mode); + w.u8(force_ar ? 1 : 0); + w.i32(restore_slot); + w.i32(kv_offset); + w.i32(snapshot_slot); + w.i32(snapshot_pos); + w.vec_i32(stop_token_ids); +} + +bool RequestMsg::decode(ByteReader & r, RequestMsg & out) { + uint8_t mode = 0; + uint8_t force = 0; + if (!r.u64(out.request_id) || !r.vec_i32(out.prompt_tokens) || !r.i32(out.n_gen) || + !r.i32(out.max_ctx) || !r.f32(out.temperature) || !r.f32(out.top_p) || + !r.i32(out.top_k) || !r.f32(out.min_p) || !r.f32(out.repeat_penalty) || + !r.u64(out.seed) || !r.u8(mode) || !r.u8(force) || !r.i32(out.restore_slot) || + !r.i32(out.kv_offset) || !r.i32(out.snapshot_slot) || !r.i32(out.snapshot_pos) || + !r.vec_i32(out.stop_token_ids, 1 << 16)) { + return false; + } + if (mode > (uint8_t)DecodeMode::Speculative) return false; + if (force > 1) return false; + out.decode_mode = (DecodeMode)mode; + out.force_ar = force != 0; + return true; +} + +void DecisionMsg::encode(ByteWriter & w) const { + w.u64(request_id); + w.u32(step); + w.i32(token); + w.u8(flags); +} + +bool DecisionMsg::decode(ByteReader & r, DecisionMsg & out) { + return r.u64(out.request_id) && r.u32(out.step) && r.i32(out.token) && r.u8(out.flags); +} + +void DraftMsg::encode(ByteWriter & w) const { + w.u64(request_id); + w.u32(step); + w.i32(pos); + w.vec_i32(tokens); +} + +bool DraftMsg::decode(ByteReader & r, DraftMsg & out) { + return r.u64(out.request_id) && r.u32(out.step) && r.i32(out.pos) && + r.vec_i32(out.tokens, 1 << 12); +} + +void AcceptMsg::encode(ByteWriter & w) const { + w.u64(request_id); + w.u32(step); + w.i32(accept); + w.i32(bonus); + w.u8(flags); +} + +bool AcceptMsg::decode(ByteReader & r, AcceptMsg & out) { + return r.u64(out.request_id) && r.u32(out.step) && r.i32(out.accept) && + r.i32(out.bonus) && r.u8(out.flags); +} + +void RequestReportMsg::encode(ByteWriter & w) const { + w.u64(request_id); + w.i32(rank); + w.u32(steps); + w.u64(compute_us); + w.u64(allreduce_calls); + w.u64(allreduce_bytes); + w.u64(allreduce_wait_us); + w.u64(ctrl_wait_us); + w.u64(peak_device_bytes); +} + +bool RequestReportMsg::decode(ByteReader & r, RequestReportMsg & out) { + return r.u64(out.request_id) && r.i32(out.rank) && r.u32(out.steps) && + r.u64(out.compute_us) && r.u64(out.allreduce_calls) && + r.u64(out.allreduce_bytes) && r.u64(out.allreduce_wait_us) && + r.u64(out.ctrl_wait_us) && r.u64(out.peak_device_bytes); +} + +void BackendOpMsg::encode(ByteWriter & w) const { + w.u64(request_id); + w.u8((uint8_t)kind); + w.vec_i64(args); +} + +bool BackendOpMsg::decode(ByteReader & r, BackendOpMsg & out) { + uint8_t kind = 0; + if (!r.u64(out.request_id) || !r.u8(kind) || !r.vec_i64(out.args, 1 << 12)) { + return false; + } + if (kind < (uint8_t)BackendOpKind::SnapshotSave || kind > (uint8_t)BackendOpKind::HandleCompress) { + return false; + } + out.kind = (BackendOpKind)kind; + return true; +} + +void RequestEndMsg::encode(ByteWriter & w) const { + w.u64(request_id); + w.u32(steps); + w.i32(status); +} + +bool RequestEndMsg::decode(ByteReader & r, RequestEndMsg & out) { + return r.u64(out.request_id) && r.u32(out.steps) && r.i32(out.status); +} + +void ShutdownMsg::encode(ByteWriter & w) const { + w.i32(reason); +} + +bool ShutdownMsg::decode(ByteReader & r, ShutdownMsg & out) { + return r.i32(out.reason); +} + +// ─── Hash ─────────────────────────────────────────────────────────────── + +uint64_t fnv1a64(const void * data, size_t n, uint64_t seed) { + const uint8_t * p = (const uint8_t *)data; + uint64_t h = seed; + for (size_t i = 0; i < n; ++i) { + h ^= (uint64_t)p[i]; + h *= 0x100000001b3ull; + } + return h; +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_protocol.h b/server/src/cluster/cluster_protocol.h new file mode 100644 index 000000000..5943d46f8 --- /dev/null +++ b/server/src/cluster/cluster_protocol.h @@ -0,0 +1,354 @@ +// cluster_protocol.h - control-channel wire protocol between cluster ranks. +// +// The control channel is a plain TCP connection from every worker to the head +// (rank 0). It carries small, latency-tolerant messages: bootstrap (Hello / +// Welcome with the RCCL unique id), one RequestDescriptor per generation, the +// per-step decisions the head makes on behalf of all ranks (Decision for AR, +// Draft + Accept for DSpark), backend mutations that must happen on every +// rank (BackendOp), heartbeats, aborts and end-of-request telemetry. +// +// Bulk tensor traffic (the per-layer partial-sum reduction) never travels on +// this channel; that is RCCL over RoCE (cluster_comm.h). +// +// Encoding: every frame is a 12-byte little-endian header followed by the +// payload. Integers are little-endian; strings and vectors carry a u32 length +// prefix. The encoding is deliberately hand-rolled so the protocol test can +// run without ggml, HIP or RCCL and so a wire dump stays readable with xxd. +// +// Compatibility rule: bump kProtocolVersion on any field change; the head +// refuses a Hello whose protocol version differs. + +#pragma once + +#include +#include +#include +#include +#include + +namespace luce::cluster { + +inline constexpr uint32_t kFrameMagic = 0x4C43424Cu; // "LBCL" +inline constexpr uint16_t kProtocolVersion = 3; +inline constexpr size_t kFrameHeaderBytes = 12; +inline constexpr size_t kMaxPayloadBytes = 64ull * 1024ull * 1024ull; +inline constexpr size_t kRcclUniqueIdBytes = 128; // == NCCL_UNIQUE_ID_BYTES + +enum class MsgType : uint16_t { + Hello = 1, // worker -> head, once after connect + Welcome = 2, // head -> worker, carries the RCCL unique id + Heartbeat = 3, // both directions, every ~1 s while idle + Abort = 4, // any -> all, request or process is being torn down + Request = 5, // head -> workers, start of one generation + Decision = 6, // head -> workers, AR: sampled token + stop flags + Draft = 7, // head -> workers, DSpark: proposed tokens for verify + Accept = 8, // head -> workers, DSpark: accepted count + bonus token + RequestReport = 10, // worker -> head, per-request telemetry + BackendOp = 11, // head -> workers, snapshot/park/etc. replicated call + RequestEnd = 12, // head -> workers, generation finished (normal path) + Shutdown = 13, // head -> workers, orderly process exit +}; + +const char * msg_type_name(MsgType type); + +struct FrameHeader { + uint32_t magic = kFrameMagic; + uint16_t version = kProtocolVersion; + uint16_t type = 0; + uint32_t payload_len = 0; +}; + +struct Frame { + MsgType type = MsgType::Heartbeat; + std::vector payload; +}; + +// ─── Byte-level (de)serialization ─────────────────────────────────────── + +class ByteWriter { +public: + void u8(uint8_t v); + void u16(uint16_t v); + void u32(uint32_t v); + void u64(uint64_t v); + void i32(int32_t v); + void i64(int64_t v); + void f32(float v); + void str(const std::string & s); // u32 length + bytes + void bytes(const void * data, size_t n); // u32 length + bytes + void vec_i32(const std::vector & v); // u32 count + i32... + void vec_i64(const std::vector & v); + void vec_u64(const std::vector & v); + + const std::vector & data() const { return buf_; } + std::vector take() { return std::move(buf_); } + size_t size() const { return buf_.size(); } + +private: + std::vector buf_; +}; + +class ByteReader { +public: + ByteReader(const uint8_t * data, size_t size) : data_(data), size_(size) {} + explicit ByteReader(const std::vector & v) : data_(v.data()), size_(v.size()) {} + + bool u8(uint8_t & v); + bool u16(uint16_t & v); + bool u32(uint32_t & v); + bool u64(uint64_t & v); + bool i32(int32_t & v); + bool i64(int64_t & v); + bool f32(float & v); + bool str(std::string & s, size_t max_len = 1 << 20); + bool bytes(std::vector & out, size_t max_len = kMaxPayloadBytes); + bool vec_i32(std::vector & v, size_t max_count = 1 << 22); + bool vec_i64(std::vector & v, size_t max_count = 1 << 20); + bool vec_u64(std::vector & v, size_t max_count = 1 << 20); + + bool at_end() const { return pos_ == size_; } + size_t remaining() const { return size_ - pos_; } + bool failed() const { return failed_; } + +private: + const uint8_t * data_; + size_t size_; + size_t pos_ = 0; + bool failed_ = false; +}; + +// Header encode/decode. decode_header validates magic and version and +// rejects payload_len > kMaxPayloadBytes. +void encode_header(const FrameHeader & h, uint8_t out[kFrameHeaderBytes]); +bool decode_header(const uint8_t in[kFrameHeaderBytes], FrameHeader & h, std::string * err); + +// ─── Messages ─────────────────────────────────────────────────────────── +// Every message has encode()/decode() and a static type(). decode() returns +// false on truncated or malformed payloads and never throws. + +struct HelloMsg { + int32_t rank = -1; + int32_t size = 0; + uint16_t protocol_version = kProtocolVersion; + std::string build_sha; // git describe of the binary + std::string model_sha; // sha256 of the target GGUF header or path hash + uint64_t placement_hash = 0; // ClusterExpertPlacement::hash() + std::string hostname; + + static MsgType type() { return MsgType::Hello; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, HelloMsg & out); +}; + +struct WelcomeMsg { + int32_t size = 0; + std::array rccl_unique_id{}; + uint64_t placement_hash = 0; + uint32_t timeout_ms = 0; + uint8_t allreduce_dtype = 0; // AllreduceDType as u8 + uint8_t shared_expert = 0; // SharedExpertMode as u8 + std::vector rank_hostnames; // index = rank + + static MsgType type() { return MsgType::Welcome; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, WelcomeMsg & out); +}; + +struct HeartbeatMsg { + int32_t rank = -1; + uint64_t mono_us = 0; // sender's monotonic clock + + static MsgType type() { return MsgType::Heartbeat; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, HeartbeatMsg & out); +}; + +struct AbortMsg { + int32_t rank = -1; + int32_t code = 0; // 0 = operator shutdown, >0 = failure class + uint64_t request_id = 0; // 0 = process-level + std::string reason; + + static MsgType type() { return MsgType::Abort; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, AbortMsg & out); +}; + +// Decode mode requested for one generation. +enum class DecodeMode : uint8_t { + Autoregressive = 0, + Speculative = 1, // DSpark: head drafts, all ranks verify +}; + +// Everything a worker needs to run the identical generate() call. Field +// names mirror the head's GenerateRequest/SamplerCfg; the head decorator +// fills this before invoking the inner DeepSeek4Backend::generate. +struct RequestMsg { + uint64_t request_id = 0; + std::vector prompt_tokens; + int32_t n_gen = 0; + int32_t max_ctx = 0; + // Sampler configuration (rank 0 is the only rank that samples; workers + // still need these so the AR-vs-spec routing decision is identical). + float temperature = 0.0f; + float top_p = 1.0f; + int32_t top_k = 0; + float min_p = 0.0f; + float repeat_penalty = 1.0f; + uint64_t seed = 0; + DecodeMode decode_mode = DecodeMode::Autoregressive; + bool force_ar = false; + // Prefix-cache state. The two halves are independent: a request may + // resume from `restore_slot` (its KV holds `kv_offset` tokens) and still + // be asked to write an inline snapshot into `snapshot_slot` once prefill + // reaches `snapshot_pos`. -1 means "none" for both slots. + int32_t restore_slot = -1; + int32_t kv_offset = 0; + int32_t snapshot_slot = -1; + int32_t snapshot_pos = -1; + std::vector stop_token_ids; // budget-hook close tokens etc. + + static MsgType type() { return MsgType::Request; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, RequestMsg & out); +}; + +// Whether every rank decodes this request autoregressively. The workers +// rebuild their request from the message and the head decodes in the mode it +// broadcast, so a backend that would speculate on its own (sampling, a +// thinking budget) never leaves the ranks in two different decode loops. +inline bool request_decodes_ar(const RequestMsg & m) { + return m.force_ar || m.decode_mode == DecodeMode::Autoregressive; +} + +enum DecisionFlags : uint8_t { + kDecisionEos = 1u << 0, + kDecisionStop = 1u << 1, // stop string / stop token matched on head + kDecisionCancel = 1u << 2, // client disconnected / cancelled + kDecisionBudget = 1u << 3, // budget hook forced a close token + kDecisionFinal = 1u << 4, // last decision of this request +}; + +// AR decode: the token rank 0 sampled for this step, plus termination flags. +struct DecisionMsg { + uint64_t request_id = 0; + uint32_t step = 0; + int32_t token = -1; + uint8_t flags = 0; + + static MsgType type() { return MsgType::Decision; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, DecisionMsg & out); +}; + +// DSpark: tokens the head drafted for this verify step. tokens.size() == q +// (seed + q-1 drafted), pos is the target position of tokens[0]. +struct DraftMsg { + uint64_t request_id = 0; + uint32_t step = 0; + int32_t pos = 0; + std::vector tokens; + + static MsgType type() { return MsgType::Draft; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, DraftMsg & out); +}; + +// DSpark: verification outcome computed on rank 0. accept = number of +// drafted tokens accepted, bonus = the target's token after the accepted +// prefix (becomes next step's seed), flags as DecisionFlags. +struct AcceptMsg { + uint64_t request_id = 0; + uint32_t step = 0; + int32_t accept = 0; + int32_t bonus = -1; + uint8_t flags = 0; + + static MsgType type() { return MsgType::Accept; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, AcceptMsg & out); +}; + +// Per-request telemetry sent by every worker after RequestEnd; the head +// merges it into usage.timings.cluster. +struct RequestReportMsg { + uint64_t request_id = 0; + int32_t rank = -1; + uint32_t steps = 0; + uint64_t compute_us = 0; + uint64_t allreduce_calls = 0; + uint64_t allreduce_bytes = 0; + uint64_t allreduce_wait_us = 0; + uint64_t ctrl_wait_us = 0; + uint64_t peak_device_bytes = 0; + + static MsgType type() { return MsgType::RequestReport; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, RequestReportMsg & out); +}; + +// Backend mutations that must be mirrored on every rank so the replicated +// KV/snapshot state stays identical. +enum class BackendOpKind : uint8_t { + SnapshotSave = 1, + SnapshotFree = 2, + SnapshotRestore = 3, + Park = 4, + Unpark = 5, + FreeDrafter = 6, + ResetRequestState = 7, + HandleCompress = 8, +}; + +struct BackendOpMsg { + uint64_t request_id = 0; + BackendOpKind kind = BackendOpKind::ResetRequestState; + std::vector args; // op-specific scalar arguments + + static MsgType type() { return MsgType::BackendOp; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, BackendOpMsg & out); +}; + +struct RequestEndMsg { + uint64_t request_id = 0; + uint32_t steps = 0; + int32_t status = 0; // 0 = ok, otherwise GenerateErrorCode as int + + static MsgType type() { return MsgType::RequestEnd; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, RequestEndMsg & out); +}; + +struct ShutdownMsg { + int32_t reason = 0; + + static MsgType type() { return MsgType::Shutdown; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, ShutdownMsg & out); +}; + +// ─── Frame helpers ────────────────────────────────────────────────────── + +template +Frame make_frame(const Msg & msg) { + ByteWriter w; + msg.encode(w); + Frame f; + f.type = Msg::type(); + f.payload = w.take(); + return f; +} + +template +bool parse_frame(const Frame & f, Msg & out) { + if (f.type != Msg::type()) return false; + ByteReader r(f.payload); + if (!Msg::decode(r, out)) return false; + return r.at_end(); +} + +// FNV-1a 64-bit, used for placement and model-path hashes. +uint64_t fnv1a64(const void * data, size_t n, uint64_t seed = 0xcbf29ce484222325ull); + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_telemetry.h b/server/src/cluster/cluster_telemetry.h new file mode 100644 index 000000000..94ce22cf7 --- /dev/null +++ b/server/src/cluster/cluster_telemetry.h @@ -0,0 +1,51 @@ +// cluster_telemetry.h - per-step timing counters shared by the head decorator, +// the worker loop and the DeepSeek4 cluster graph path. +// +// Header-only and dependency-free so both the model code and the decision +// hooks / worker main can accumulate into the same struct; the head folds the +// per-rank totals into RequestReportMsg. + +#pragma once + +#include +#include + +namespace luce::cluster { + +struct ClusterStepTelemetry { + uint64_t compute_us = 0; // local graph compute (host wall time) + uint64_t allreduce_calls = 0; + uint64_t allreduce_bytes = 0; // payload bytes handed to RCCL + uint64_t allreduce_wait_us = 0; // host time blocked on collectives + uint64_t ctrl_wait_us = 0; // host time blocked on the control channel + + void add(const ClusterStepTelemetry & o) { + compute_us += o.compute_us; + allreduce_calls += o.allreduce_calls; + allreduce_bytes += o.allreduce_bytes; + allreduce_wait_us += o.allreduce_wait_us; + ctrl_wait_us += o.ctrl_wait_us; + } + + void reset() { *this = ClusterStepTelemetry{}; } +}; + +// RAII: adds the elapsed wall time (microseconds, steady clock) to `acc` +// when the scope ends. +class ScopedUsTimer { +public: + explicit ScopedUsTimer(uint64_t & acc) + : acc_(acc), start_(std::chrono::steady_clock::now()) {} + ScopedUsTimer(const ScopedUsTimer &) = delete; + ScopedUsTimer & operator=(const ScopedUsTimer &) = delete; + ~ScopedUsTimer() { + const auto end = std::chrono::steady_clock::now(); + acc_ += (uint64_t)std::chrono::duration_cast(end - start_).count(); + } + +private: + uint64_t & acc_; + std::chrono::steady_clock::time_point start_; +}; + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_worker_main.cpp b/server/src/cluster/cluster_worker_main.cpp new file mode 100644 index 000000000..a741a16bb --- /dev/null +++ b/server/src/cluster/cluster_worker_main.cpp @@ -0,0 +1,480 @@ +// cluster_worker_main.cpp - lockstep worker loop for ranks 1..N-1. + +#include "cluster/cluster_worker_main.h" + +#include "cluster/cluster_comm.h" +#include "cluster/cluster_config.h" +#include "cluster/cluster_control.h" +#include "cluster/cluster_decision_hooks.h" +#include "cluster/cluster_head_backend.h" // backend_placement_hash +#include "cluster/cluster_identity.h" +#include "cluster/cluster_protocol.h" +#include "common/backend_factory.h" +#include "common/cluster_participant.h" +#include "common/model_backend.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace luce::cluster { + +using common::DaemonIO; +using common::GenerateRequest; +using common::GenerateResult; + +namespace { + +constexpr int kExitOk = 0; +constexpr int kExitConfig = 2; +constexpr int kExitAborted = 3; +constexpr int kExitBackend = 4; +constexpr int kExitProtocol = 5; + +constexpr int kSelftestIters = 1000; +constexpr size_t kSelftestSmallN = 64u * 1024u / sizeof(float); +constexpr size_t kSelftestLargeN = 16u * 1024u * 1024u / sizeof(float); + +uint64_t now_us() { + return (uint64_t) std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()).count(); +} + +// RequestMsg -> GenerateRequest. Inverse of ClusterHeadBackend:: +// build_request_msg; see the list there of the fields that do not travel. +GenerateRequest request_from_msg(const RequestMsg & m) { + GenerateRequest req; + req.prompt = m.prompt_tokens; + req.n_gen = m.n_gen; + req.sampler.temp = m.temperature; + req.sampler.top_p = m.top_p; + req.sampler.top_k = m.top_k; + req.sampler.rep_pen = m.repeat_penalty; + req.sampler.seed = m.seed; + req.do_sample = m.temperature > 0.0f; + req.stream = false; + req.force_ar_decode = request_decodes_ar(m); + // Since protocol 2 both halves travel, so this rank writes the same + // inline snapshot at the same position as the head. + req.snap_slot = m.snapshot_slot; + req.snap_pos = m.snapshot_pos; + req.budget_hook.close_token_ids = m.stop_token_ids; + req.budget_hook.hard_limit_remaining = 0; // the head's override arrives inside the token + return req; +} + +struct WorkerState { + ClusterConfig cfg; + std::unique_ptr backend_owner; + // The backend as the worker replays requests into it, and the same object + // seen as a cluster participant. Both are borrowed from backend_owner. + common::ModelBackend * model = nullptr; + common::IClusterParticipant * part = nullptr; + ClusterWorkerControl control; + std::unique_ptr comm; + std::unique_ptr hooks; + uint64_t requests_served = 0; + int device = -1; // local HIP ordinal + // Watchdog: with the in-graph all-reduce nothing on the + // host waits for a collective, so a dead head would leave this rank + // blocked inside HIP forever. The thread aborts the communicator, which + // makes the pending collective return an error. + std::thread watchdog; + std::atomic watchdog_stop{false}; + std::atomic request_in_flight{false}; + // High-water mark of device memory in use, sampled at request ends. + uint64_t peak_device_bytes = 0; +}; + +void send_abort(WorkerState & st, int code, uint64_t request_id, const std::string & reason) { + std::fprintf(stderr, "[cluster] worker %d: ABORT (%d): %s\n", st.cfg.rank, code, reason.c_str()); + AbortMsg abort; + abort.rank = st.cfg.rank; + abort.code = code; + abort.request_id = request_id; + abort.reason = reason; + st.control.send(make_frame(abort), nullptr); + if (st.comm) st.comm->abort(); +} + +void teardown(WorkerState & st) { + st.watchdog_stop.store(true); + if (st.watchdog.joinable()) st.watchdog.join(); + if (st.part) st.part->cluster_set_hooks(nullptr); + if (st.backend_owner) st.backend_owner->shutdown(); + st.control.close(); + st.comm.reset(); + st.hooks.reset(); +} + +// Executes one BackendOpMsg on the local backend. Returns false on an +// argument error (the head never sends those; treat as protocol failure). +bool apply_backend_op(WorkerState & st, const BackendOpMsg & op) { + auto arg = [&](size_t i, int64_t def) { return i < op.args.size() ? op.args[i] : def; }; + switch (op.kind) { + case BackendOpKind::SnapshotSave: + if (!st.model->snapshot_save((int) arg(0, -1))) { + // A rank without the snapshot would silently resume different KV + // on the next restore. Fail the run instead. + std::fprintf(stderr, + "[cluster] worker %d: snapshot_save(%lld) failed; the slot would " + "differ from the head's\n", + st.cfg.rank, (long long) arg(0, -1)); + return false; + } + return true; + case BackendOpKind::SnapshotFree: + st.model->snapshot_free((int) arg(0, -1)); + return true; + case BackendOpKind::SnapshotRestore: + // Restore is carried by RequestMsg (snapshot_slot + kv_offset), not + // as a standalone op; accept and ignore for forward compatibility. + return true; + case BackendOpKind::Park: + if (!st.model->park((common::ParkTarget) arg(0, (int64_t) common::ParkTarget::Empty))) { + std::fprintf(stderr, "[cluster] worker %d: park failed\n", st.cfg.rank); + } + return true; + case BackendOpKind::Unpark: + if (!st.model->unpark((common::ParkTarget) arg(0, (int64_t) common::ParkTarget::Empty))) { + std::fprintf(stderr, "[cluster] worker %d: unpark failed\n", st.cfg.rank); + return false; + } + return true; + case BackendOpKind::FreeDrafter: + st.model->free_drafter(); + return true; + case BackendOpKind::ResetRequestState: + st.model->release_scratch(); + return true; + case BackendOpKind::HandleCompress: + // DeepSeek4Backend::handle_compress is unsupported upstream; nothing to mirror. + return true; + } + return false; +} + +void worker_watchdog_loop(WorkerState & st) { + constexpr int kPeriodMs = 250; + while (!st.watchdog_stop.load()) { + std::this_thread::sleep_for(std::chrono::milliseconds(kPeriodMs)); + if (st.watchdog_stop.load()) break; + if (!st.request_in_flight.load()) continue; + std::string err; + const bool head_gone = !st.control.alive(); + const bool async_failed = st.comm && st.comm->async_error(&err); + if (!head_gone && !async_failed) continue; + std::fprintf(stderr, + "[cluster] worker %d: %s during a request; aborting the communicator\n", + st.cfg.rank, + head_gone ? "the head stopped responding" + : ("communicator error: " + err).c_str()); + if (st.comm) st.comm->abort(); + // Leave the exit to the request path: the aborted collective makes the + // forward fail, handle_request returns, and the loop tears down. + break; + } +} + +// Returns the exit code to use, or -1 to keep serving. +int handle_request(WorkerState & st, const RequestMsg & msg) { + const uint64_t t0 = now_us(); + GenerateRequest req = request_from_msg(msg); + DaemonIO io; // discarding sink: no fd, no callback, never cancels + io.stream_fd = -1; + + st.hooks->set_current_request(msg.request_id); + st.hooks->reset_counters(); + if (st.comm) st.comm->reset_stats(); + st.request_in_flight.store(true); + struct InFlightGuard { + WorkerState & st; + ~InFlightGuard() { st.request_in_flight.store(false); } + } in_flight_guard{st}; + + // The head speculates only when it broadcast DecodeMode::Speculative. + // A worker without a drafter would silently decode AR instead and then + // wait for a Decision frame that never comes, so fail the run loudly. + if (msg.decode_mode == DecodeMode::Speculative && !st.part->cluster_spec_decode_ready()) { + send_abort(st, kExitProtocol, msg.request_id, + "head requested speculative decode but this rank has no DSpark drafter " + "(pass the same --draft / LUCE_DS4_SPEC settings to every rank)"); + return kExitProtocol; + } + + GenerateResult result; + if (msg.restore_slot >= 0) { + result = st.model->restore_and_generate_impl(msg.restore_slot, req, io); + } else { + result = st.model->generate_impl(req, io); + } + const uint64_t compute_us = now_us() - t0; + + // The head always follows the request with RequestEnd; if our decode + // loop already consumed it (head failed the request early) it is parked. + Frame frame; + std::string err; + bool have_end = false; + if (st.hooks->take_pending_frame(frame)) { + have_end = frame.type == MsgType::RequestEnd; + } else if (st.hooks->aborted()) { + std::fprintf(stderr, "[cluster] worker %d: head aborted request %llu: %s\n", st.cfg.rank, + (unsigned long long) msg.request_id, st.hooks->abort_reason().c_str()); + return kExitAborted; + } else if (st.control.recv(frame, st.cfg.timeout_ms, &err)) { + have_end = frame.type == MsgType::RequestEnd; + } else { + send_abort(st, kExitProtocol, msg.request_id, "no RequestEnd: " + err); + return kExitProtocol; + } + + if (!have_end) { + if (frame.type == MsgType::Abort) { + AbortMsg abort; + parse_frame(frame, abort); + std::fprintf(stderr, "[cluster] worker %d: head aborted: %s\n", st.cfg.rank, + abort.reason.c_str()); + return kExitAborted; + } + if (frame.type == MsgType::Shutdown) return kExitOk; + send_abort(st, kExitProtocol, msg.request_id, + std::string("expected RequestEnd, got ") + msg_type_name(frame.type)); + return kExitProtocol; + } + RequestEndMsg end; + if (!parse_frame(frame, end) || end.request_id != msg.request_id) { + send_abort(st, kExitProtocol, msg.request_id, "RequestEnd for a different request"); + return kExitProtocol; + } + + if (!result.ok() && end.status == 0) { + // Head succeeded, we did not: replicated state has diverged. + send_abort(st, kExitBackend, msg.request_id, + "worker generate failed while head succeeded: " + std::string(result.error_detail())); + return kExitBackend; + } + if (result.ok() && (uint32_t) result.tokens.size() != end.steps) { + send_abort(st, kExitBackend, msg.request_id, + "token count " + std::to_string(result.tokens.size()) + " != head " + + std::to_string(end.steps)); + return kExitBackend; + } + + RequestReportMsg report; + report.request_id = msg.request_id; + report.rank = st.cfg.rank; + report.steps = (uint32_t) result.tokens.size(); + report.compute_us = compute_us; + if (st.comm) { + const ClusterCommStats s = st.comm->stats(); + report.allreduce_calls = s.allreduce_calls; + report.allreduce_bytes = s.allreduce_bytes; + report.allreduce_wait_us = s.allreduce_wait_us; + } + report.ctrl_wait_us = st.hooks->counters().ctrl_wait_us; + const uint64_t device_bytes = cluster_device_bytes_in_use(st.device); + if (device_bytes > st.peak_device_bytes) st.peak_device_bytes = device_bytes; + report.peak_device_bytes = st.peak_device_bytes; + if (!st.control.send(make_frame(report), &err)) { + std::fprintf(stderr, "[cluster] worker %d: RequestReport send failed: %s\n", st.cfg.rank, + err.c_str()); + return kExitProtocol; + } + st.requests_served++; + if (cluster_env_trace()) { + std::fprintf(stderr, "[cluster] worker %d: req=%llu %zu tokens in %.3f s (ctrl wait %.1f ms)\n", + st.cfg.rank, (unsigned long long) msg.request_id, result.tokens.size(), + compute_us / 1e6, report.ctrl_wait_us / 1e3); + } + return -1; +} + +} // namespace + +int run_cluster_worker(const ClusterConfig & cfg, + std::unique_ptr backend, + const std::string & model_path, + int device) { + WorkerState st; + st.cfg = cfg; + const std::string cfg_err = st.cfg.validate(); + if (!cfg_err.empty() || !st.cfg.is_worker()) { + std::fprintf(stderr, "[cluster] worker: invalid cluster config: %s\n", + cfg_err.empty() ? "rank must be >= 1" : cfg_err.c_str()); + return kExitConfig; + } + // NCCL_* environment: server_main exported it before the backend loaded. + st.backend_owner = std::move(backend); + if (!st.backend_owner) { + std::fprintf(stderr, "[cluster] worker %d: no backend\n", st.cfg.rank); + return kExitBackend; + } + st.model = st.backend_owner.get(); + st.part = dynamic_cast(st.model); + if (!st.part) { + std::fprintf(stderr, "[cluster] worker %d: this architecture does not " + "implement IClusterParticipant, so it cannot shard\n", + st.cfg.rank); + teardown(st); + return kExitConfig; + } + + // Join the cluster. Identity fields are computed by cluster_identity.h on + // every rank; the placement hash is this rank's built placement. + const uint64_t placement_hash = backend_placement_hash(*st.model); + if (placement_hash == 0) { + std::fprintf(stderr, "[cluster] worker %d: backend has no cluster placement\n", st.cfg.rank); + teardown(st); + return kExitBackend; + } + const HelloMsg hello = cluster_identity(st.cfg.rank, st.cfg.size, model_path, placement_hash); + ControlTimeouts timeouts; + timeouts.recv_ms = st.cfg.timeout_ms; + st.control.set_timeouts(timeouts); + WelcomeMsg welcome; + std::string err; + std::fprintf(stderr, "[cluster] worker %d/%d connecting to %s:%d (build=%s model=%s)\n", + st.cfg.rank, st.cfg.size, st.cfg.head_host.c_str(), st.cfg.head_port, + hello.build_sha.c_str(), hello.model_sha.c_str()); + if (!st.control.connect_and_handshake(st.cfg.head_host, st.cfg.head_port, hello, welcome, &err)) { + std::fprintf(stderr, "[cluster] worker %d: handshake failed: %s\n", st.cfg.rank, err.c_str()); + teardown(st); + return kExitProtocol; + } + if (welcome.size != st.cfg.size) { + std::fprintf(stderr, "[cluster] worker %d: head reports size %d, we were started with %d\n", + st.cfg.rank, welcome.size, st.cfg.size); + teardown(st); + return kExitConfig; + } + // The head's values win for knobs that must agree cluster-wide. + // shared_expert is placement-defining (set_cluster refuses a change after + // init), so a mismatch is an operator error rather than a silent override. + st.cfg.timeout_ms = welcome.timeout_ms ? welcome.timeout_ms : st.cfg.timeout_ms; + st.cfg.allreduce_dtype = (AllreduceDType) welcome.allreduce_dtype; + if ((SharedExpertMode) welcome.shared_expert != st.cfg.shared_expert) { + std::fprintf(stderr, "[cluster] worker %d: --cluster-shared-expert differs from the head " + "(%s vs %s)\n", st.cfg.rank, + shared_expert_mode_name(st.cfg.shared_expert), + shared_expert_mode_name((SharedExpertMode) welcome.shared_expert)); + teardown(st); + return kExitConfig; + } + + RcclCommInit init; + init.rank = st.cfg.rank; + init.size = st.cfg.size; + init.unique_id = welcome.rccl_unique_id; + init.device = device; // local HIP ordinal of the primary GPU + st.device = init.device; + init.timeout_ms = st.cfg.timeout_ms; + init.blocking = false; + st.comm = create_rccl_cluster_comm(init, &err); + if (!st.comm) { + std::fprintf(stderr, "[cluster] worker %d: RCCL communicator failed: %s\n", st.cfg.rank, err.c_str()); + send_abort(st, kExitBackend, 0, "rccl init: " + err); + teardown(st); + return kExitBackend; + } + std::fprintf(stderr, "[cluster] worker %d: %s communicator up (%d ranks)\n", st.cfg.rank, + st.comm->backend_name(), st.comm->size()); + st.watchdog = std::thread([&st] { worker_watchdog_loop(st); }); + + if (st.cfg.selftest) { + const bool ok = run_cluster_selftest(*st.comm, init.device, kSelftestIters, kSelftestSmallN, + kSelftestLargeN, &err); + std::fprintf(stderr, "[cluster] worker %d: selftest %s%s%s\n", st.cfg.rank, + ok ? "PASSED" : "FAILED", err.empty() ? "" : ": ", err.c_str()); + Frame f; + st.control.recv(f, st.cfg.timeout_ms, nullptr); // head's Shutdown + teardown(st); + return ok ? kExitOk : 1; + } + + WorkerHooksConfig hcfg; + hcfg.timeout_ms = st.cfg.timeout_ms; + st.hooks = std::make_unique(st.control, hcfg); + // Second set_cluster call: attach the communicator to the + // runtime the factory created before init(). + if (!st.part->cluster_attach(&st.cfg, st.comm.get())) { + send_abort(st, kExitBackend, 0, "set_cluster(comm) refused by the backend"); + teardown(st); + return kExitBackend; + } + st.part->cluster_set_hooks(st.hooks.get()); + + if (!st.comm->barrier(&err)) { + std::fprintf(stderr, "[cluster] worker %d: post-init barrier failed: %s\n", st.cfg.rank, err.c_str()); + send_abort(st, kExitBackend, 0, "barrier: " + err); + teardown(st); + return kExitBackend; + } + std::fprintf(stderr, "[cluster] worker %d ready\n", st.cfg.rank); + + // ── Main loop ───────────────────────────────────────────────────── + int exit_code = kExitOk; + for (;;) { + Frame frame; + if (!st.hooks->take_pending_frame(frame)) { + if (!st.control.recv(frame, st.cfg.timeout_ms, &err)) { + // Idle periods between requests are unbounded; a recv deadline + // is only fatal when the heartbeat says the head is gone. + if (st.control.alive()) continue; + std::fprintf(stderr, "[cluster] worker %d: head connection lost: %s\n", st.cfg.rank, + err.c_str()); + if (st.comm) st.comm->abort(); + exit_code = kExitAborted; + break; + } + } + if (frame.type == MsgType::Request) { + RequestMsg msg; + if (!parse_frame(frame, msg)) { + send_abort(st, kExitProtocol, 0, "malformed Request frame"); + exit_code = kExitProtocol; + break; + } + const int rc = handle_request(st, msg); + if (rc >= 0) { exit_code = rc; break; } + } else if (frame.type == MsgType::BackendOp) { + BackendOpMsg op; + if (!parse_frame(frame, op) || !apply_backend_op(st, op)) { + send_abort(st, kExitProtocol, op.request_id, "BackendOp failed"); + exit_code = kExitProtocol; + break; + } + } else if (frame.type == MsgType::RequestEnd) { + // Stale RequestEnd (e.g. after a head-side prefill failure that + // our decode never reached); nothing to do. + continue; + } else if (frame.type == MsgType::Shutdown) { + ShutdownMsg bye; + parse_frame(frame, bye); + std::fprintf(stderr, "[cluster] worker %d: shutdown (reason %d) after %llu request(s)\n", + st.cfg.rank, bye.reason, (unsigned long long) st.requests_served); + exit_code = kExitOk; + break; + } else if (frame.type == MsgType::Abort) { + AbortMsg abort; + parse_frame(frame, abort); + std::fprintf(stderr, "[cluster] worker %d: head abort (%d): %s\n", st.cfg.rank, abort.code, + abort.reason.c_str()); + if (st.comm) st.comm->abort(); + exit_code = kExitAborted; + break; + } else { + send_abort(st, kExitProtocol, 0, std::string("unexpected frame ") + msg_type_name(frame.type)); + exit_code = kExitProtocol; + break; + } + } + teardown(st); + return exit_code; +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_worker_main.h b/server/src/cluster/cluster_worker_main.h new file mode 100644 index 000000000..fe23a1a59 --- /dev/null +++ b/server/src/cluster/cluster_worker_main.h @@ -0,0 +1,22 @@ +// cluster_worker_main.h - lockstep worker loop for ranks 1..N-1. +// server_main builds the backend through the +// regular factory (the DeepSeek4 backend loads only this rank's expert share) +// and hands it over; the worker joins the head, brings up RCCL and replays +// every request the head broadcasts. Returns the process exit code. + +#pragma once + +#include "cluster/cluster_config.h" +#include "common/model_backend.h" + +#include +#include + +namespace luce::cluster { + +int run_cluster_worker(const ClusterConfig & cfg, + std::unique_ptr backend, + const std::string & model_path, + int device); + +} // namespace luce::cluster diff --git a/server/src/cluster/fast_reduce.cpp b/server/src/cluster/fast_reduce.cpp new file mode 100644 index 000000000..e851b85c0 --- /dev/null +++ b/server/src/cluster/fast_reduce.cpp @@ -0,0 +1,842 @@ +#include "cluster/fast_reduce.h" + +// The RDMA fabric is built only with -DLUCE_CLUSTER=ON (HIP + libibverbs). +// Without it, FastReduce never comes up and callers keep RCCL, the same way +// cluster_comm.cpp compiles its RCCL half under LUCE_CLUSTER_RCCL. +#ifdef LUCE_CLUSTER + +#include +#include +#include + +#include "cluster/cluster_control.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace luce::cluster { + +// The kernels live in fast_reduce_kernels.cu (LANGUAGE HIP). A kernel that +// writes an exchange row also raises the row's flag from its last block. +void fast_reduce_launch_pack(const float * src, void * dst, int n, int bf16, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream); +void fast_reduce_launch_add(float * dst, const void * peer, int n, int bf16, hipStream_t stream); +void fast_reduce_launch_hyb_combine(const float * cold, const float * hot, void * own_row, int n, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream); +void fast_reduce_launch_hyb_final(const void * own_row, const void * peer_row, float * dst, int n, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream); + +namespace { + +// What one rank has to tell every other rank to be written into. +constexpr int kMaxRanks = 8; + +struct Endpoint { + // qpn[r] is the queue pair this rank created to talk to rank r. Every pair + // needs its own number, so one per rank is not enough: rank 2 must connect + // to the queue pair rank 1 made *for rank 2*, not to whichever one rank 1 + // happened to make first. + uint32_t qpn[kMaxRanks] = {0}; + uint32_t psn = 0; + // One key per region. The payload and the flags are separate memory + // regions -- deliberately, so a flag write cannot be reordered behind a + // payload write -- and a remote write carries the key of the region it + // lands in. Using one key for both is not a mismatch the sender notices: + // the receiver drops it as an access violation and the peer simply waits. + uint32_t rkey_data = 0; + uint32_t rkey_flag = 0; + uint64_t data_addr = 0; // base of the receive region + uint64_t flag_addr = 0; // base of the flag region + uint8_t gid[16] = {0}; +}; + +bool set_nodelay(int fd) { + int one = 1; + return setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one)) == 0; +} + +// A star: rank 0 collects every endpoint and hands the full table back. Small +// and synchronous, which is what a one-off exchange should be. +bool exchange(const FastReduce::Config & cfg, const Endpoint & mine, + std::vector & all, std::string * err) { + all.assign((size_t) cfg.size, Endpoint{}); + all[(size_t) cfg.rank] = mine; + if (cfg.size <= 1) return true; + + auto fail = [&](const char * what) { + if (err) *err = std::string("fast-reduce bootstrap: ") + what + ": " + std::strerror(errno); + return false; + }; + auto fail_io = [&](const char * what, const std::string & why) { + if (err) *err = std::string("fast-reduce bootstrap: ") + what + ": " + why; + return false; + }; + + if (cfg.rank == 0) { + const int lst = ::socket(AF_INET, SOCK_STREAM, 0); + if (lst < 0) return fail("socket"); + int one = 1; + setsockopt(lst, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)); + sockaddr_in a{}; + a.sin_family = AF_INET; + a.sin_addr.s_addr = INADDR_ANY; + a.sin_port = htons((uint16_t) cfg.bootstrap_port); + if (::bind(lst, (sockaddr *) &a, sizeof(a)) != 0) { ::close(lst); return fail("bind"); } + if (::listen(lst, cfg.size) != 0) { ::close(lst); return fail("listen"); } + + // Every wait is bounded by timeout_ms: a worker that died during its + // model load must fail rank 0's load, not park it forever. + std::vector peers((size_t) cfg.size, -1); + auto close_peers = [&]() { + for (int & fd : peers) { + if (fd >= 0) ::close(fd); + fd = -1; + } + }; + std::string ioerr; + for (int i = 1; i < cfg.size; ++i) { + pollfd pfd{lst, POLLIN, 0}; + const int ready = ::poll(&pfd, 1, (int) cfg.timeout_ms); + if (ready <= 0) { + if (ready == 0) errno = ETIMEDOUT; + ::close(lst); close_peers(); + return fail("accept"); + } + const int fd = ::accept(lst, nullptr, nullptr); + if (fd < 0) { ::close(lst); close_peers(); return fail("accept"); } + set_nodelay(fd); + int32_t r = -1; + Endpoint e{}; + if (!read_exact_deadline(fd, &r, sizeof(r), cfg.timeout_ms, &ioerr) || + !read_exact_deadline(fd, &e, sizeof(e), cfg.timeout_ms, &ioerr) || + r <= 0 || r >= cfg.size || peers[(size_t) r] >= 0) { + ::close(fd); ::close(lst); close_peers(); + return fail_io("read endpoint", ioerr.empty() ? "bad or duplicate rank" : ioerr); + } + all[(size_t) r] = e; + peers[(size_t) r] = fd; + } + ::close(lst); + for (int i = 1; i < cfg.size; ++i) { + if (!write_exact_deadline(peers[(size_t) i], all.data(), sizeof(Endpoint) * all.size(), + cfg.timeout_ms, &ioerr)) { + close_peers(); + return fail_io("write table", ioerr); + } + } + close_peers(); + return true; + } + + // Workers retry: rank 0 may still be loading when they get here. + int fd = -1; + for (int attempt = 0; attempt < 600; ++attempt) { + fd = ::socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) return fail("socket"); + sockaddr_in a{}; + a.sin_family = AF_INET; + a.sin_port = htons((uint16_t) cfg.bootstrap_port); + if (::inet_pton(AF_INET, cfg.bootstrap_host.c_str(), &a.sin_addr) != 1) { + ::close(fd); + if (err) *err = "fast-reduce bootstrap: bad host " + cfg.bootstrap_host; + return false; + } + if (::connect(fd, (sockaddr *) &a, sizeof(a)) == 0) break; + ::close(fd); + fd = -1; + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + if (fd < 0) return fail("connect"); + set_nodelay(fd); + const int32_t r = cfg.rank; + std::string ioerr; + if (!write_exact_deadline(fd, &r, sizeof(r), cfg.timeout_ms, &ioerr) || + !write_exact_deadline(fd, &mine, sizeof(mine), cfg.timeout_ms, &ioerr) || + !read_exact_deadline(fd, all.data(), sizeof(Endpoint) * all.size(), cfg.timeout_ms, &ioerr)) { + ::close(fd); + return fail_io("exchange", ioerr); + } + ::close(fd); + return true; +} + +} // namespace + +// A switch is on when set to anything but empty or "0". +static bool fr_env_on(const char * name) { + const char * v = std::getenv(name); + return v && *v && std::strcmp(v, "0") != 0; +} + +struct FastReduceImpl { + // Last-block counters for in-kernel publishes, per device (a kernel counts + // in its own device memory): [slot][kind], kinds 0 attention row, 1 hybrid + // export, 2 hybrid combine, 3 hybrid final. Zeroed on the first stream that + // needs them, ahead of its first publish. + static constexpr int kPubKinds = 4; + std::mutex pub_counter_mu; + std::map pub_counters; + uint32_t * pub_counter(hipStream_t stm, int slot, int kind); + FastReduce::Config cfg; + bool up = false; + + ibv_context * ctx = nullptr; + ibv_pd * pd = nullptr; + ibv_cq * cq = nullptr; + ibv_mr * mr_data = nullptr; + ibv_mr * mr_flag = nullptr; + std::vector qps; // one per peer, null at own rank + std::vector peers; + + // Pinned, NIC-registered, GPU-mapped. One region for payloads and one for + // flags, so a flag write cannot be reordered behind a payload write on a + // different memory region. + float * data_host = nullptr; // [slots][size][max_elems] + float * data_dev = nullptr; + // Two flag arrays, not one. The NIC writes `arrival` by DMA and only the + // host reads it: a DMA write into system memory is coherent with a CPU + // read, but nothing invalidates the GPU's caches for it, so a kernel + // spinning on a NIC-written word can spin on a stale line forever. The + // host progress thread sees the arrival and then writes `visible` with an + // ordinary store, which is the path the probe measured at 0.79 us. This is + // the same reason a general collective has a proxy thread at all. + uint32_t * flag_host = nullptr; // [slots][size] arrival | [slots][size] visible + // | [slots] pub | [slots] hybrid export | [slots] hybrid final + uint32_t * flag_dev = nullptr; + + + // Payload length per reduction, on its own much larger ring. It cannot + // share the slot ring: it is written when the host *enqueues* a reduction + // and the host enqueues a whole step ahead of the GPU, so with one ring the + // entry the progress thread still needs is overwritten a hundred + // reductions before it reads it. + static constexpr int kPendingRing = 8192; + std::vector pending_n; + std::atomic seq{0}; + std::atomic stop{false}; + // Set by the stall watchdog: every later submit becomes a no-op so the GPU + // never parks on a peer that is gone, and the process exits once idle. + std::atomic failed{false}; + // Hybrid exchange buffers, pinned: [part 0 = the dGPU's export | part 1 = the final sum][slots][hyb_elems]. + float * hyb_host = nullptr; + float * hyb_dev = nullptr; + int hyb_elems = 0; + // Exchange index -> sequence number: whichever of an exchange's calls the + // host issues first takes the next sequence number for it. + static constexpr int kHybRing = 4096; + std::vector> hyb_ring; + float * hyb_part(int part, int slot) const { + return hyb_dev + ((size_t) part * (size_t) cfg.slots + (size_t) slot) * (size_t) hyb_elems; + } + std::thread progress; + + size_t slot_stride() const { return (size_t) cfg.max_elems * (size_t) cfg.size; } + // My own row inside a slot is never written by a peer; peers write theirs. + float * recv_row(int slot, int peer) { + return data_host + (size_t) slot * slot_stride() + (size_t) peer * cfg.max_elems; + } + size_t arrival_off(int slot, int peer) const { + return (size_t) slot * (size_t) cfg.size + (size_t) peer; + } + size_t visible_off(int slot, int peer) const { + return (size_t) cfg.slots * (size_t) cfg.size + arrival_off(slot, peer); + } + size_t pub_off(int slot) const { + return 2u * (size_t) cfg.slots * (size_t) cfg.size + (size_t) slot; + } + uint32_t * arrival_flag(int slot, int peer) { return flag_host + arrival_off(slot, peer); } + uint32_t * visible_flag(int slot, int peer) { return flag_host + visible_off(slot, peer); } + uint32_t * pub_flag(int slot) { return flag_host + pub_off(slot); } + size_t hybh_off(int slot) const { return pub_off(0) + (size_t) cfg.slots + (size_t) slot; } + size_t hybf_off(int slot) const { return hybh_off(0) + (size_t) cfg.slots + (size_t) slot; } + size_t flag_count() const { + return 2u * (size_t) cfg.size * (size_t) cfg.slots + 3u * (size_t) cfg.slots; + } +}; + +FastReduce::FastReduce() : p_(new FastReduceImpl()) {} +uint32_t * FastReduceImpl::pub_counter(hipStream_t stm, int slot, int kind) { + int dev = 0; + if (hipStreamGetDevice(stm, &dev) != hipSuccess) return nullptr; + std::lock_guard lock(pub_counter_mu); + auto it = pub_counters.find(dev); + if (it == pub_counters.end()) { + int prev = 0; + if (hipGetDevice(&prev) != hipSuccess || hipSetDevice(dev) != hipSuccess) return nullptr; + uint32_t * ctr = nullptr; + const size_t bytes = (size_t) cfg.slots * kPubKinds * sizeof(uint32_t); + const bool ok = hipMalloc((void **) &ctr, bytes) == hipSuccess && + hipMemsetAsync(ctr, 0, bytes, stm) == hipSuccess; // ahead of this stream's first publish + (void) hipSetDevice(prev); + if (!ok) return nullptr; + it = pub_counters.emplace(dev, ctr).first; + } + return it->second + (size_t) slot * kPubKinds + (size_t) kind; +} + +FastReduce::~FastReduce() { shutdown(); } + +bool FastReduce::ok() const { return p_ && p_->up; } + +bool FastReduce::init(const Config & cfg, std::string * err) { + auto & s = *p_; + s.cfg = cfg; + if (cfg.size <= 1) { s.up = false; return true; } + if (cfg.size != 2) { + if (err) *err = "fast-reduce: a two-rank exchange, " + std::to_string(cfg.size) + " ranks keep RCCL"; + return false; + } + + auto bail = [&](const std::string & what) { + if (err) *err = "fast-reduce: " + what; + shutdown(); + return false; + }; + + // ── device ──────────────────────────────────────────────────────────── + int n_dev = 0; + ibv_device ** devs = ibv_get_device_list(&n_dev); + if (!devs) return bail("ibv_get_device_list failed"); + ibv_device * dev = nullptr; + for (int i = 0; i < n_dev; ++i) { + if (cfg.hca.empty() || cfg.hca == ibv_get_device_name(devs[i])) { + dev = devs[i]; + break; + } + } + if (!dev) { ibv_free_device_list(devs); return bail("no HCA named " + cfg.hca); } + s.ctx = ibv_open_device(dev); + ibv_free_device_list(devs); + if (!s.ctx) return bail("ibv_open_device failed"); + + s.pd = ibv_alloc_pd(s.ctx); + if (!s.pd) return bail("ibv_alloc_pd failed"); + // Two work requests per peer per reduction, a step's worth in flight. + s.cq = ibv_create_cq(s.ctx, 4096, nullptr, nullptr, 0); + if (!s.cq) return bail("ibv_create_cq failed"); + + // ── buffers ─────────────────────────────────────────────────────────── + const size_t data_elems = (size_t) cfg.slots * s.slot_stride(); + if (hipHostMalloc((void **) &s.data_host, data_elems * sizeof(float), + hipHostMallocDefault) != hipSuccess) { + return bail("hipHostMalloc for the payload region failed"); + } + std::memset(s.data_host, 0, data_elems * sizeof(float)); + if (hipHostMalloc((void **) &s.flag_host, s.flag_count() * sizeof(uint32_t), + hipHostMallocDefault) != hipSuccess) { + return bail("hipHostMalloc for the flag region failed"); + } + std::memset(s.flag_host, 0, s.flag_count() * sizeof(uint32_t)); + if (fr_env_on("LUCE_CLUSTER_HYBRID_EXCHANGE")) { + // Decode widths only: 5120 x 6 lanes fits; prefill keeps its own path. + s.hyb_elems = std::min(cfg.max_elems, 32768); + const size_t bytes = (size_t) 2 * (size_t) cfg.slots * (size_t) s.hyb_elems * sizeof(float); + float * h = nullptr; + if (hipHostMalloc((void **) &h, bytes, hipHostMallocDefault) != hipSuccess) { + return bail("hipHostMalloc for the hybrid exchange failed"); + } + s.hyb_host = h; + if (hipHostGetDevicePointer((void **) &s.hyb_dev, s.hyb_host, 0) != hipSuccess) { + return bail("hipHostGetDevicePointer for the hybrid exchange failed"); + } + s.hyb_ring.assign(FastReduceImpl::kHybRing, {UINT64_MAX, 0}); + std::fprintf(stderr, "[fast-reduce] hybrid exchange buffers: %.1f MiB pinned\n", bytes / 1048576.0); + } + + if (hipHostGetDevicePointer((void **) &s.data_dev, s.data_host, 0) != hipSuccess || + hipHostGetDevicePointer((void **) &s.flag_dev, s.flag_host, 0) != hipSuccess) { + return bail("hipHostGetDevicePointer failed"); + } + + s.mr_data = ibv_reg_mr(s.pd, s.data_host, data_elems * sizeof(float), + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + s.mr_flag = ibv_reg_mr(s.pd, s.flag_host, s.flag_count() * sizeof(uint32_t), + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE); + if (!s.mr_data || !s.mr_flag) return bail("ibv_reg_mr failed (memlock limit?)"); + + // ── queue pairs ─────────────────────────────────────────────────────── + s.qps.assign((size_t) cfg.size, nullptr); + for (int r = 0; r < cfg.size; ++r) { + if (r == cfg.rank) continue; + ibv_qp_init_attr ia{}; + ia.send_cq = s.cq; + ia.recv_cq = s.cq; + ia.qp_type = IBV_QPT_RC; + ia.cap.max_send_wr = 4096; + ia.cap.max_recv_wr = 1; + ia.cap.max_send_sge = 1; + ia.cap.max_recv_sge = 1; + ia.cap.max_inline_data = 64; // the flag write goes inline + s.qps[(size_t) r] = ibv_create_qp(s.pd, &ia); + if (!s.qps[(size_t) r]) return bail("ibv_create_qp failed"); + + ibv_qp_attr a{}; + a.qp_state = IBV_QPS_INIT; + a.pkey_index = 0; + a.port_num = (uint8_t) cfg.ib_port; + a.qp_access_flags = IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_LOCAL_WRITE; + if (ibv_modify_qp(s.qps[(size_t) r], &a, + IBV_QP_STATE | IBV_QP_PKEY_INDEX | IBV_QP_PORT | + IBV_QP_ACCESS_FLAGS) != 0) { + return bail("ibv_modify_qp to INIT failed"); + } + } + + // The memory keys and addresses are the same for every peer; only the + // queue pair numbers differ. + Endpoint mine{}; + mine.rkey_data = s.mr_data->rkey; + mine.rkey_flag = s.mr_flag->rkey; + mine.data_addr = (uint64_t) (uintptr_t) s.data_host; + mine.flag_addr = (uint64_t) (uintptr_t) s.flag_host; + mine.psn = 0; + ibv_gid gid{}; + if (ibv_query_gid(s.ctx, (uint8_t) cfg.ib_port, cfg.gid_index, &gid) != 0) { + return bail("ibv_query_gid failed"); + } + std::memcpy(mine.gid, gid.raw, 16); + + if (cfg.size > kMaxRanks) return bail("more ranks than the endpoint carries"); + for (int r = 0; r < cfg.size; ++r) { + if (r == cfg.rank) continue; + mine.qpn[r] = s.qps[(size_t) r]->qp_num; + } + + std::vector table; + if (!exchange(cfg, mine, table, err)) { shutdown(); return false; } + s.peers = table; + + // LUCE_CLUSTER_FAST_REDUCE_MTU: the path MTU (default 4096, what a direct + // ConnectX link runs). + const char * mtu_env = std::getenv("LUCE_CLUSTER_FAST_REDUCE_MTU"); + const int mtu = mtu_env ? std::atoi(mtu_env) : 4096; + + for (int r = 0; r < cfg.size; ++r) { + if (r == cfg.rank) continue; + ibv_qp_attr a{}; + a.qp_state = IBV_QPS_RTR; + a.path_mtu = mtu >= 4096 ? IBV_MTU_4096 : (mtu >= 2048 ? IBV_MTU_2048 : IBV_MTU_1024); + a.dest_qp_num = s.peers[(size_t) r].qpn[cfg.rank]; + a.rq_psn = 0; + a.max_dest_rd_atomic = 1; + a.min_rnr_timer = 12; + a.ah_attr.is_global = 1; + a.ah_attr.port_num = (uint8_t) cfg.ib_port; + a.ah_attr.grh.hop_limit = 64; + a.ah_attr.grh.sgid_index = (uint8_t) cfg.gid_index; + std::memcpy(a.ah_attr.grh.dgid.raw, s.peers[(size_t) r].gid, 16); + if (ibv_modify_qp(s.qps[(size_t) r], &a, + IBV_QP_STATE | IBV_QP_AV | IBV_QP_PATH_MTU | + IBV_QP_DEST_QPN | IBV_QP_RQ_PSN | + IBV_QP_MAX_DEST_RD_ATOMIC | IBV_QP_MIN_RNR_TIMER) != 0) { + return bail("ibv_modify_qp to RTR failed"); + } + ibv_qp_attr b{}; + b.qp_state = IBV_QPS_RTS; + b.timeout = 14; + b.retry_cnt = 7; + b.rnr_retry = 7; + b.sq_psn = 0; + b.max_rd_atomic = 1; + if (ibv_modify_qp(s.qps[(size_t) r], &b, + IBV_QP_STATE | IBV_QP_TIMEOUT | IBV_QP_RETRY_CNT | + IBV_QP_RNR_RETRY | IBV_QP_SQ_PSN | IBV_QP_MAX_QP_RD_ATOMIC) != 0) { + return bail("ibv_modify_qp to RTS failed"); + } + } + + s.pending_n.assign((size_t) FastReduceImpl::kPendingRing, 0); + + // ── progress ────────────────────────────────────────────────────────── + // LUCE_CLUSTER_FAST_REDUCE_STALL_S: how long a reduction may wait before + // the rank releases its GPU and exits (0 turns the watchdog off). + const char * stall_env = std::getenv("LUCE_CLUSTER_FAST_REDUCE_STALL_S"); + const double stall_s = stall_env ? std::atof(stall_env) : 20.0; + + s.up = true; + s.progress = std::thread([&s, stall_s]() { + uint64_t next = 1; + std::vector wc(64); + // Stall release: every flag a stream can be parked on for reduction q, + // including the hybrid exchange's export and final words, so the GPUs + // drain instead of needing a reset. + auto release_one = [&s](uint64_t q) { + const int qs = (int) (q % (uint64_t) s.cfg.slots); + for (int r = 0; r < s.cfg.size; ++r) { + if (r != s.cfg.rank) *(volatile uint32_t *) s.visible_flag(qs, r) = (uint32_t) q; + } + if (s.hyb_host) { + *(volatile uint32_t *) (s.flag_host + s.hybh_off(qs)) = (uint32_t) q; + *(volatile uint32_t *) (s.flag_host + s.hybf_off(qs)) = (uint32_t) q; + } + }; + // A rank that cannot finish reduction `first` (stalled, or its RDMA + // write failed) releases every reduction the host has handed its GPU + // so the queue drains, makes later submits no-ops, and exits once the + // GPU is idle: exiting with a parked queue costs a GPU reset, and the + // peer's own watchdog ends the other rank. + auto release_and_exit = [&s, &release_one](uint64_t first) { + s.failed.store(true); + // Nothing reports a finished reduction, so keep releasing for the + // whole grace period; late submits are released too. + const auto drain_t0 = std::chrono::steady_clock::now(); + while (std::chrono::duration(std::chrono::steady_clock::now() - drain_t0).count() < 15.0) { + const uint64_t last = std::max(first, s.seq.load()); + for (uint64_t q = first; q <= last; ++q) release_one(q); + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + std::fprintf(stderr, "[fast-reduce] rank %d: GPU released, exiting\n", s.cfg.rank); + std::fflush(nullptr); + std::_Exit(4); + }; + // A broken connection also starves the reductions already published: + // any of the last ring's worth may still wait on a peer's flag. + auto oldest_in_flight = [&s](uint64_t q) { + return q > (uint64_t) s.cfg.slots ? q - (uint64_t) s.cfg.slots + 1 : (uint64_t) 1; + }; + while (!s.stop.load(std::memory_order_relaxed)) { + const int slot = (int) (next % (uint64_t) s.cfg.slots); + volatile uint32_t * pub = s.pub_flag(slot); + uint64_t spins = 0; + bool moaned = false; + const auto wait_t0 = std::chrono::steady_clock::now(); + // The unpublished-reduction clock starts when the host hands the + // reduction out, not when this thread starts waiting: an idle + // server waits here between requests and that time is no stall. + bool handed_out = false; + auto handed_t0 = wait_t0; + while (*pub != (uint32_t) next) { + if (s.stop.load(std::memory_order_relaxed)) return; + if (++spins < 200000ull) continue; + spins = 0; + std::this_thread::yield(); + if (!handed_out && s.seq.load() >= next) { + handed_out = true; + handed_t0 = std::chrono::steady_clock::now(); + } + if (handed_out && stall_s > 0.0) { + const double w = std::chrono::duration( + std::chrono::steady_clock::now() - handed_t0).count(); + if (w > stall_s) { // handed out, never published + std::fprintf(stderr, + "[fast-reduce] rank %d: STALL, reduction %llu unpublished for %.0f s " + "with %llu queued after it; releasing and exiting\n", + s.cfg.rank, (unsigned long long) next, w, + (unsigned long long) (s.seq.load() - next)); + release_and_exit(next); + } + } + if (next <= 1) continue; + // The GPU published reduction p = next-1 and is waiting on the + // peer's flag for it. Arrived long ago is the normal idle case; + // published but not arrived means a peer never delivered. + const uint64_t p = next - 1; + const int ps = (int) (p % (uint64_t) s.cfg.slots); + bool arrived = true; + for (int r = 0; r < s.cfg.size; ++r) { + if (r == s.cfg.rank) continue; + if (*(volatile uint32_t *) s.visible_flag(ps, r) != (uint32_t) p) arrived = false; + } + if (arrived) continue; + const double waited = std::chrono::duration( + std::chrono::steady_clock::now() - wait_t0).count(); + if (waited > 5.0 && !moaned) { + moaned = true; + std::string vis; + for (int r = 0; r < s.cfg.size; ++r) { + if (r == s.cfg.rank) continue; + vis += " r" + std::to_string(r) + "=" + + std::to_string(*(volatile uint32_t *) s.visible_flag(ps, r)); + } + std::fprintf(stderr, + "[fast-reduce] rank %d: reduction %llu published %.1f s ago, " + "not arrived; peer flags%s (want %llu); host submitted %llu\n", + s.cfg.rank, (unsigned long long) p, waited, vis.c_str(), + (unsigned long long) p, + (unsigned long long) s.seq.load()); + } + if (stall_s <= 0.0 || waited < stall_s) continue; + const uint64_t last = s.seq.load(); + std::fprintf(stderr, + "[fast-reduce] rank %d: STALL at reduction %llu for %.0f s " + "(peer gone?); releasing %llu pending reduction(s) and exiting\n", + s.cfg.rank, (unsigned long long) p, waited, + (unsigned long long) (last >= p ? last - p + 1 : 1)); + release_and_exit(p); + } + int n = s.pending_n[(size_t) (next % (uint64_t) FastReduceImpl::kPendingRing)]; + const bool exact_row = n < 0; // submit(..., exact): f32 even on a bf16 wire + if (exact_row) n = -n; + for (int r = 0; r < s.cfg.size; ++r) { + if (r == s.cfg.rank) continue; + const Endpoint & pe = s.peers[(size_t) r]; + + ibv_sge sge{}; + sge.addr = (uint64_t) (uintptr_t) (s.data_host + + (size_t) slot * s.slot_stride() + (size_t) s.cfg.rank * s.cfg.max_elems); + const bool bf16_row = !exact_row; + sge.length = (uint32_t) ((size_t) n * (bf16_row ? 2u : sizeof(float))); + sge.lkey = s.mr_data->lkey; + + ibv_send_wr w{}; + w.wr_id = next; + w.sg_list = &sge; + w.num_sge = 1; + w.opcode = IBV_WR_RDMA_WRITE; + w.send_flags = 0; + w.wr.rdma.remote_addr = pe.data_addr + + ((size_t) slot * s.slot_stride() + (size_t) s.cfg.rank * s.cfg.max_elems) * sizeof(float); + w.wr.rdma.rkey = pe.rkey_data; + + // The flag rides the same queue pair, so reliable-connection + // ordering puts it after the payload. That ordering is the + // whole protocol: a flag that could overtake its data would + // hand the peer a half-written buffer. + const uint32_t flagv = (uint32_t) next; + ibv_sge fsge{}; + fsge.addr = (uint64_t) (uintptr_t) &flagv; + fsge.length = sizeof(uint32_t); + fsge.lkey = 0; + + ibv_send_wr fw{}; + fw.wr_id = next; + fw.sg_list = &fsge; + fw.num_sge = 1; + fw.opcode = IBV_WR_RDMA_WRITE; + fw.send_flags = IBV_SEND_INLINE | IBV_SEND_SIGNALED; + // Straight at the word the peer's stream is waiting on. When + // a shader was the waiter this had to go through the host: a + // DMA write into system memory is coherent with a CPU read but + // nothing invalidates the GPU's caches for it. The command + // processor is not behind those caches, so the hop is gone. + fw.wr.rdma.remote_addr = pe.flag_addr + + s.visible_off(slot, s.cfg.rank) * sizeof(uint32_t); + fw.wr.rdma.rkey = pe.rkey_flag; + + w.next = &fw; + ibv_send_wr * bad = nullptr; + if (ibv_post_send(s.qps[(size_t) r], &w, &bad) != 0) { + std::fprintf(stderr, "[fast-reduce] rank %d: ibv_post_send failed (seq %llu); " + "releasing and exiting\n", s.cfg.rank, (unsigned long long) next); + release_and_exit(oldest_in_flight(next)); + } + } + // Drain completions so the send queue does not fill. A failed + // write is not recoverable here: the peer is already waiting on a + // flag that will never arrive, so this rank releases its own GPU + // and exits, and the peer's watchdog ends the other rank. + int done = 0; + do { + done = ibv_poll_cq(s.cq, (int) wc.size(), wc.data()); + for (int i = 0; i < done; ++i) { + if (wc[(size_t) i].status != IBV_WC_SUCCESS) { + std::fprintf(stderr, + "[fast-reduce] write failed: %s (seq %llu); releasing and exiting\n", + ibv_wc_status_str(wc[(size_t) i].status), + (unsigned long long) wc[(size_t) i].wr_id); + release_and_exit(oldest_in_flight(next)); + } + } + } while (done > 0); + + // Nothing to carry across any more: the peer's own stream is + // waiting on the word this rank's NIC wrote. + if (next <= 3 || next % 4096 == 0) { + std::fprintf(stderr, "[fast-reduce] seq %llu delivered (%d floats)\n", + (unsigned long long) next, n); + } + ++next; + } + }); + + return true; +} + +bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { + auto & s = *p_; + if (s.failed.load(std::memory_order_relaxed)) return true; // stalled: drain, then exit + if (!s.up || n == 0 || n > (size_t) s.cfg.max_elems) return false; + + const uint64_t seq = s.seq.fetch_add(1, std::memory_order_relaxed) + 1; + const int slot = (int) (seq % (uint64_t) s.cfg.slots); + // The progress thread reads this once it sees the flag, and the flag is + // raised further down the stream -- so writing it here is ordered ahead of + // any possible read. An exact row (f32 on a bf16 wire) is stored negated. + s.pending_n[(size_t) (seq % (uint64_t) FastReduceImpl::kPendingRing)] = exact ? -(int) n : (int) n; + + // The wait is a stream operation, not a kernel spin. On gfx1151 a kernel + // spinning on a flag the host raised needs a system-scope acquire, which + // invalidates the L2 the next layer is about to read; hipStreamWaitValue32 + // is handled by the command processor, outside the shader caches. + // + // No slot-reuse wait is needed: the ranks advance in lockstep (a rank + // publishes reduction q only after it added the peer's row of q - 1), so a + // slot comes round again only after both ranks consumed it `slots` + // reductions earlier, and the stream orders this GPU's writes behind its + // reads. + // + // Pack straight into the pinned row the NIC sends from (the kernel raises + // the flag), wait for the peer's row, add it straight out of pinned memory. + auto stm = (hipStream_t) stream; + const int wire_bf16 = exact ? 0 : 1; + const int peer = 1 - s.cfg.rank; + float * own_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) s.cfg.rank * s.cfg.max_elems; + const float * peer_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) peer * s.cfg.max_elems; + uint32_t * ctr = s.pub_counter(stm, slot, 0); + if (!ctr) return false; + fast_reduce_launch_pack(data, own_row, (int) n, wire_bf16, + s.flag_dev + s.pub_off(slot), (uint32_t) seq, ctr, stm); + if (hipStreamWaitValue32(stm, s.flag_dev + s.visible_off(slot, peer), (uint32_t) seq, + hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; + fast_reduce_launch_add(data, peer_row, (int) n, wire_bf16, stm); + return true; +} + +bool FastReduce::hyb_available() const { + const auto & s = *p_; + return s.up && s.hyb_host != nullptr && !s.failed.load(std::memory_order_relaxed); +} + +// The sequence number of hybrid exchange k, taken by its first call. +static uint64_t hyb_seq_for(FastReduceImpl & s, uint64_t k, size_t n) { + auto & e = s.hyb_ring[(size_t) (k % (uint64_t) FastReduceImpl::kHybRing)]; + if (e.first == k) return e.second; + const uint64_t seq = s.seq.fetch_add(1, std::memory_order_relaxed) + 1; + s.pending_n[(size_t) (seq % (uint64_t) FastReduceImpl::kPendingRing)] = (int) n; + e = {k, seq}; + return seq; +} + +bool FastReduce::hyb_export(const float * data, size_t n, void * stream, uint64_t k) { + auto & s = *p_; + if (s.failed.load(std::memory_order_relaxed)) return true; + if (!s.hyb_host || n == 0 || n > (size_t) s.hyb_elems) return false; + const uint64_t seq = hyb_seq_for(s, k, n); + const int slot = (int) (seq % (uint64_t) s.cfg.slots); + auto stm = (hipStream_t) stream; + uint32_t * ctr = s.pub_counter(stm, slot, 1); + if (!ctr) return false; + fast_reduce_launch_pack(data, s.hyb_part(0, slot), (int) n, 0, + s.flag_dev + s.hybh_off(slot), (uint32_t) seq, ctr, stm); + return true; +} + +bool FastReduce::hyb_combine(const float * data, size_t n, void * stream, uint64_t k) { + auto & s = *p_; + if (s.failed.load(std::memory_order_relaxed)) return true; + if (!s.hyb_host || n == 0 || n > (size_t) s.hyb_elems) return false; + const uint64_t seq = hyb_seq_for(s, k, n); + const int slot = (int) (seq % (uint64_t) s.cfg.slots); + const int peer = 1 - s.cfg.rank; + float * own_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) s.cfg.rank * s.cfg.max_elems; + const float * peer_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) peer * s.cfg.max_elems; + auto stm = (hipStream_t) stream; + if (hipStreamWaitValue32(stm, s.flag_dev + s.hybh_off(slot), (uint32_t) seq, + hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; + uint32_t * ctr_pub = s.pub_counter(stm, slot, 2); + uint32_t * ctr_fin = s.pub_counter(stm, slot, 3); + if (!ctr_pub || !ctr_fin) return false; + fast_reduce_launch_hyb_combine(data, s.hyb_part(0, slot), own_row, (int) n, + s.flag_dev + s.pub_off(slot), (uint32_t) seq, ctr_pub, stm); + if (hipStreamWaitValue32(stm, s.flag_dev + s.visible_off(slot, peer), (uint32_t) seq, + hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; + fast_reduce_launch_hyb_final(own_row, peer_row, s.hyb_part(1, slot), (int) n, + s.flag_dev + s.hybf_off(slot), (uint32_t) seq, ctr_fin, stm); + return true; +} + +bool FastReduce::hyb_import(float * data, size_t n, void * stream, uint64_t k) { + auto & s = *p_; + if (s.failed.load(std::memory_order_relaxed)) return true; + if (!s.hyb_host || n == 0 || n > (size_t) s.hyb_elems) return false; + const uint64_t seq = hyb_seq_for(s, k, n); + const int slot = (int) (seq % (uint64_t) s.cfg.slots); + auto stm = (hipStream_t) stream; + if (hipStreamWaitValue32(stm, s.flag_dev + s.hybf_off(slot), (uint32_t) seq, + hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; + return hipMemcpyAsync(data, s.hyb_part(1, slot), n * sizeof(float), hipMemcpyHostToDevice, stm) == hipSuccess; +} + +bool FastReduce::submit_from_root(float * data, size_t n, void * stream, int root, bool exact) { + auto & s = *p_; + if (!s.up || n == 0 || n > (size_t) s.cfg.max_elems) return false; + if (s.cfg.rank != root && + hipMemsetAsync(data, 0, n * sizeof(float), (hipStream_t) stream) != hipSuccess) { + return false; + } + return submit(data, n, stream, exact); +} + +void FastReduce::shutdown() { + if (!p_) return; + auto & s = *p_; + s.stop.store(true); + if (s.progress.joinable()) s.progress.join(); + s.up = false; + + for (ibv_qp * q : s.qps) if (q) ibv_destroy_qp(q); + s.qps.clear(); + if (s.mr_data) { ibv_dereg_mr(s.mr_data); s.mr_data = nullptr; } + if (s.mr_flag) { ibv_dereg_mr(s.mr_flag); s.mr_flag = nullptr; } + if (s.cq) { ibv_destroy_cq(s.cq); s.cq = nullptr; } + if (s.pd) { ibv_dealloc_pd(s.pd); s.pd = nullptr; } + if (s.ctx) { ibv_close_device(s.ctx); s.ctx = nullptr; } + + for (auto & [dev, ctr] : s.pub_counters) { + int prev = 0; + if (hipGetDevice(&prev) == hipSuccess && hipSetDevice(dev) == hipSuccess) { + (void) hipFree(ctr); + (void) hipSetDevice(prev); + } + } + s.pub_counters.clear(); + if (s.data_host) { (void) hipHostFree(s.data_host); s.data_host = nullptr; } + if (s.flag_host) { (void) hipHostFree(s.flag_host); s.flag_host = nullptr; } + if (s.hyb_host) { (void) hipHostFree(s.hyb_host); s.hyb_host = nullptr; } +} + +} // namespace luce::cluster + +#else // !LUCE_CLUSTER + +namespace luce::cluster { + +struct FastReduceImpl {}; + +FastReduce::FastReduce() = default; +FastReduce::~FastReduce() = default; + +bool FastReduce::init(const Config &, std::string * err) { + if (err) *err = "fast reduce not compiled in (build with -DLUCE_CLUSTER=ON on HIP)"; + return false; +} + +void FastReduce::shutdown() {} +bool FastReduce::ok() const { return false; } +bool FastReduce::submit(float *, size_t, void *, bool) { return false; } +bool FastReduce::submit_from_root(float *, size_t, void *, int, bool) { return false; } +bool FastReduce::hyb_available() const { return false; } +bool FastReduce::hyb_export(const float *, size_t, void *, uint64_t) { return false; } +bool FastReduce::hyb_combine(const float *, size_t, void *, uint64_t) { return false; } +bool FastReduce::hyb_import(float *, size_t, void *, uint64_t) { return false; } + +} // namespace luce::cluster + +#endif // LUCE_CLUSTER diff --git a/server/src/cluster/fast_reduce.h b/server/src/cluster/fast_reduce.h new file mode 100644 index 000000000..1b3f0adb3 --- /dev/null +++ b/server/src/cluster/fast_reduce.h @@ -0,0 +1,105 @@ +// fast_reduce.h - a lean all-reduce for the decode path's tiny payloads. +// +// WHY THIS EXISTS. A decode step makes about a hundred reductions of a few +// kilobytes each, and a general-purpose collective pays a fixed cost per call +// that dwarfs what the wire needs for such a payload. So this replaces the +// collective for exactly the case it is bad at: a few +// kilobytes, latency-bound, on the critical path of every layer. Each rank +// writes its partial straight into every peer's pre-registered buffer and adds +// what arrives. There is no algorithm selection, no protocol negotiation and +// no proxy handshake -- one RDMA write per peer and a flag. +// +// WHAT IT ASSUMES. Two ranks, and pinned host memory that both the NIC and +// the GPU reach: the NIC writes the peer's row into it and the GPU reads it +// there, so no GPUDirect is needed. It also assumes the ranks run +// the same graph in lockstep, so that the n-th reduction on one rank is the +// n-th on all of them -- which is what the sequence numbers check rather than +// trust. +// +// WHAT IT IS NOT. Not a general collective: two ranks, sum of f32 (bf16 on +// the wire unless a row asks to be exact), small sizes, no groups. Anything +// outside that keeps using RCCL. +// +// Include convention: #include "cluster/fast_reduce.h" + +#pragma once + +#include +#include +#include +#include + +namespace luce::cluster { + +// Opaque so verbs headers stay out of every translation unit that reduces. +struct FastReduceImpl; + +class FastReduce { +public: + FastReduce(); + ~FastReduce(); + + FastReduce(const FastReduce &) = delete; + FastReduce & operator=(const FastReduce &) = delete; + + struct Config { + int rank = 0; + int size = 1; + std::string hca; // e.g. "rocep197s0f3" + int ib_port = 1; + int gid_index = 1; // RoCE v2 + std::string bootstrap_host; // rank 0's address + int bootstrap_port = 9500; + // Sized for a plain decode step (10240 floats is the largest it + // reduces); DeepSeek V4 sizes both for its verify batch. The ring is + // deep so the GPU can run a whole step ahead of the host. Prefill's + // much larger reductions fall back to RCCL. + int max_elems = 16384; // 64 KiB + int slots = 4096; // a step has ~97 + uint32_t timeout_ms = 30000; + }; + + // Brings the fabric up for exactly two ranks (any other size returns false + // so the caller keeps RCCL): registers the buffers, exchanges endpoints over a + // plain TCP mesh of its own, and connects one queue pair per peer. The + // bootstrap is separate from the cluster control channel on purpose -- + // this has to be able to fail without taking the run with it. + bool init(const Config & cfg, std::string * err); + + void shutdown(); + + bool ok() const; + + // Enqueue a sum-reduction of `n` floats at `data` on `stream`. `data` must + // be device-visible and is updated in place. Returns false when the + // payload is larger than the fabric was built for, so the caller can fall + // back rather than corrupt the step. Rows cross as bf16; `exact` keeps + // this payload in f32 on the wire and in the add, for values that must + // arrive bit for bit. + bool submit(float * data, size_t n, void * stream, bool exact = false); + + // Broadcast from `root` as a sum: every other rank zeroes its copy on the + // stream first, so the sum is the root's data everywhere (with `exact`, + // the root's data itself: x + 0 is x). + bool submit_from_root(float * data, size_t n, void * stream, int root, bool exact = false); + + // Hybrid exchange (LUCE_CLUSTER_HYBRID_EXCHANGE=1): a rank whose partial is + // split between a dGPU (hot experts plus shared expert, exported as one + // partial) and the iGPU that owns its cold experts. The dGPU exports its + // partial to pinned memory while it waits anyway; the iGPU adds its part to + // it in the order the dGPU used to, packs and publishes the rank's row, + // and adds the peer's when it lands; the dGPU imports the finished sum. + // The iGPU's partial never crosses the dGPU's PCIe link, and the dGPU + // crosses it once each way. `k` is the exchange's index in issue order; + // its export, combine and import may be enqueued in any host order. + bool hyb_available() const; + bool hyb_export(const float * data, size_t n, void * stream, uint64_t k); + bool hyb_combine(const float * data, size_t n, void * stream, uint64_t k); + bool hyb_import(float * data, size_t n, void * stream, uint64_t k); + + +private: + std::unique_ptr p_; +}; + +} // namespace luce::cluster diff --git a/server/src/cluster/fast_reduce_kernels.cu b/server/src/cluster/fast_reduce_kernels.cu new file mode 100644 index 000000000..93c5ecf47 --- /dev/null +++ b/server/src/cluster/fast_reduce_kernels.cu @@ -0,0 +1,128 @@ +// fast_reduce_kernels.cu - the device half of the lean decode-path reduction. +// +// The waits are stream memory operations (hipStreamWaitValue32, in +// fast_reduce.cpp); these kernels move and add the rows: +// +// pack write this rank's row into pinned memory (bf16, or f32 for +// an exact row) and raise its flag. +// add add the peer's row straight out of pinned memory. +// hyb_combine, the two-rank hybrid exchange (FastReduce::hyb_combine). +// hyb_final +// +// The flags live in pinned memory, not managed memory: gfx1151 reports XNACK +// disabled, so the GPU cannot take a page fault, and a managed page the CPU +// has touched faults the moment a kernel reads it. Pinned pages never migrate. +// +// A stream wait has no timeout of its own: the host's progress thread bounds +// it (LUCE_CLUSTER_FAST_REDUCE_STALL_S) and releases the GPU when a peer is +// gone. + +#include + +#include + +namespace luce::cluster { + +namespace { +constexpr int kBlock = 256; +constexpr int kMaxGrid = 128; + +int grid_for(int n) { + int grid = (n + kBlock - 1) / kBlock; + if (grid < 1) grid = 1; + if (grid > kMaxGrid) grid = kMaxGrid; + return grid; +} +} // namespace + +// bf16 rows round both operands, so every rank computes the same sum. +__device__ __forceinline__ uint16_t fr_f32_to_bf16(float f) { + uint32_t u = __float_as_uint(f); + u += 0x7FFFu + ((u >> 16) & 1u); + return (uint16_t) (u >> 16); +} +__device__ __forceinline__ float fr_bf16_to_f32(uint16_t h) { + return __uint_as_float(((uint32_t) h) << 16); +} + +// The grid that wrote an exchange row publishes its flag itself. Each block +// orders its writes at device scope and counts in; the last makes every +// block's writes visible system-wide with ONE release and stores the flag, +// which is cheaper than a hipStreamWriteValue32 after the kernel. counter is +// this device's, zeroed, and left zeroed for the next grid on the same stream. +__device__ __forceinline__ void fr_publish_last_block(uint32_t * flag, uint32_t value, uint32_t * counter) { + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) { + const uint32_t prev = atomicAdd(counter, 1u); + if (prev == gridDim.x - 1) { + *counter = 0; + __threadfence_system(); + __hip_atomic_store(flag, value, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM); + } + } +} + +__global__ void fast_reduce_pack_kernel(const float * __restrict__ src, void * dst, int n, int bf16, + uint32_t * flag, uint32_t value, uint32_t * counter) { + for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < n; k += blockDim.x * gridDim.x) { + if (bf16) ((uint16_t *) dst)[k] = fr_f32_to_bf16(src[k]); + else ((float *) dst)[k] = src[k]; + } + fr_publish_last_block(flag, value, counter); +} + +__global__ void fast_reduce_add_kernel(float * __restrict__ dst, const void * peer, int n, int bf16) { + for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < n; k += blockDim.x * gridDim.x) { + if (bf16) { + const float own = fr_bf16_to_f32(fr_f32_to_bf16(dst[k])); + dst[k] = own + fr_bf16_to_f32(((const uint16_t *) peer)[k]); + } else { + dst[k] = dst[k] + ((const float *) peer)[k]; + } + } +} + +// Hybrid exchange, iGPU side. own = cold + hot in f32, where hot already holds +// the dGPU's hot experts plus shared expert, then the bf16 wire row; the final +// sum is bf16(own) + peer exactly as fast_reduce_add_kernel forms it. +__global__ void fast_reduce_hyb_combine_kernel(const float * __restrict__ cold, const float * __restrict__ hot, + void * own_row, int n, + uint32_t * flag, uint32_t value, uint32_t * counter) { + for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < n; k += blockDim.x * gridDim.x) { + ((uint16_t *) own_row)[k] = fr_f32_to_bf16(cold[k] + hot[k]); + } + fr_publish_last_block(flag, value, counter); +} + +__global__ void fast_reduce_hyb_final_kernel(const void * own_row, const void * peer_row, float * __restrict__ dst, + int n, uint32_t * flag, uint32_t value, uint32_t * counter) { + for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < n; k += blockDim.x * gridDim.x) { + dst[k] = fr_bf16_to_f32(((const uint16_t *) own_row)[k]) + fr_bf16_to_f32(((const uint16_t *) peer_row)[k]); + } + fr_publish_last_block(flag, value, counter); +} + +void fast_reduce_launch_pack(const float * src, void * dst, int n, int bf16, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream) { + hipLaunchKernelGGL(fast_reduce_pack_kernel, dim3(grid_for(n)), dim3(kBlock), 0, stream, src, dst, n, bf16, + flag, value, counter); +} + +void fast_reduce_launch_add(float * dst, const void * peer, int n, int bf16, hipStream_t stream) { + hipLaunchKernelGGL(fast_reduce_add_kernel, dim3(grid_for(n)), dim3(kBlock), 0, stream, dst, peer, n, bf16); +} + +void fast_reduce_launch_hyb_combine(const float * cold, const float * hot, void * own_row, int n, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream) { + hipLaunchKernelGGL(fast_reduce_hyb_combine_kernel, dim3(grid_for(n)), dim3(kBlock), 0, stream, + cold, hot, own_row, n, flag, value, counter); +} + +void fast_reduce_launch_hyb_final(const void * own_row, const void * peer_row, float * dst, int n, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream) { + hipLaunchKernelGGL(fast_reduce_hyb_final_kernel, dim3(grid_for(n)), dim3(kBlock), 0, stream, + own_row, peer_row, dst, n, flag, value, counter); +} + +} // namespace luce::cluster diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index 0c5107f3b..0334901f8 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -16,6 +16,7 @@ #include "placement/remote_draft_config.h" #include "placement/remote_target_shard_config.h" #include "prefill_attention_mode.h" +#include "cluster/cluster_config.h" namespace luce::common { @@ -87,6 +88,8 @@ struct BackendArgs { // Expert ownership and routing files (see docs/DS41.md); empty = unused. std::string expert_placement; // --expert-placement (alias --ds4-expert-placement) std::string ds4_router_bias; + // Expert-parallel cluster rank (--cluster-*); size 0 = a single box. + cluster::ClusterConfig cluster; std::string ds4_protected_experts; // Attention and speculative-decode options. Individual backends consume diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 694d33cff..0f015b27c 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -409,6 +409,7 @@ std::unique_ptr construct_backend( cfg.expert_placement_path = execution.expert_placement; cfg.router_bias_path = execution.router_bias; cfg.protected_experts_path = execution.protected_experts; + cfg.cluster = execution.cluster; cfg.prefill_mode = execution.prefill_mode; cfg.paged_attention = cache.paged_attention; cfg.max_concurrency = execution.max_concurrency; diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index fb428c451..ff08ddb70 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -91,6 +91,7 @@ class BackendPlan final { std::string expert_placement; std::string router_bias; std::string protected_experts; + cluster::ClusterConfig cluster; }; BackendPlan(BackendPlan &&) noexcept = default; diff --git a/server/src/common/backend_plan.cpp b/server/src/common/backend_plan.cpp index 7819ba0b4..d5b8bd85a 100644 --- a/server/src/common/backend_plan.cpp +++ b/server/src/common/backend_plan.cpp @@ -177,6 +177,7 @@ BackendPreparation BackendPlanBuilder::resolve( plan.execution_.expert_placement = args.expert_placement; plan.execution_.router_bias = args.ds4_router_bias; plan.execution_.protected_experts = args.ds4_protected_experts; + plan.execution_.cluster = args.cluster; plan.warnings_ = std::move(warnings); return plan; diff --git a/server/src/common/cluster_participant.h b/server/src/common/cluster_participant.h new file mode 100644 index 000000000..2ac15389e --- /dev/null +++ b/server/src/common/cluster_participant.h @@ -0,0 +1,61 @@ +// common/cluster_participant.h - what a backend must offer to run in a cluster. +// +// The head decorator and the worker loop were written against +// DeepSeek4Backend because it was the only architecture that sharded. Nothing +// they do actually needs that type: they broadcast requests, replay them, and +// hand the backend a config and a communicator. This interface is the whole of +// what they need, so a second architecture can join without either of them +// learning its name. +// +// Attachment happens in two steps, in this order, because the placement must +// exist before the ranks can agree on it and the communicator does not exist +// until they have: +// +// cluster_attach(&cfg, nullptr) before init() -- build the rank's placement +// cluster_attach(&cfg, comm) after RCCL is up -- attach the transport +// +// Include convention: #include "common/cluster_participant.h" + +#pragma once + +#include + +namespace luce::cluster { +struct ClusterConfig; +class IClusterComm; +struct Ds4ClusterHooks; +} // namespace luce::cluster + +namespace luce::common { + +class IClusterParticipant { +public: + virtual ~IClusterParticipant() = default; + + // Returns false when the config is one this backend cannot honour. + virtual bool cluster_attach(const cluster::ClusterConfig * cfg, + cluster::IClusterComm * comm) = 0; + + // A digest of this rank's weight placement, compared across ranks during + // the handshake so a mismatched shard set fails at startup rather than as + // wrong output. Zero means "no placement built yet", which is an error at + // handshake time. + virtual uint64_t cluster_placement_hash() const = 0; + + // Where the rank-0-decides hooks are installed, and where they are removed + // again at shutdown (nullptr). + virtual void cluster_set_hooks(cluster::Ds4ClusterHooks * hooks) = 0; + + // Whether this rank can speculate. The head must not speculate while a + // worker cannot, so a mismatch is rejected rather than silently downgraded. + virtual bool cluster_spec_decode_ready() const { return false; } + + // For /props: bytes of routed experts resident on this rank, and whether + // the collective runs inside the graph rather than between per-layer + // computes. Defaults describe a rank that shards nothing and reduces + // between computes. + virtual uint64_t cluster_resident_expert_bytes() const { return 0; } + virtual bool cluster_ingraph_allreduce() const { return false; } +}; + +} // namespace luce::common diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index 6bb68b164..8a1571e31 100644 --- a/server/src/common/dspark_head.cpp +++ b/server/src/common/dspark_head.cpp @@ -1,4 +1,5 @@ #include "dspark_head.h" +#include "pinned_stage.h" #include "ggml-alloc.h" #include "ddtree.h" @@ -221,6 +222,8 @@ struct MarkovChainGraphCache { int n_candidates = -1; bool want_confidence = false; bool want_logit_margin = false; + // Pinned staging for the per-call uploads and readbacks. + PinnedStage stage; bool matches(const DraftWeights & dw, ggml_tensor * head, ggml_backend_t candidate_backend, int candidate_count, @@ -462,16 +465,21 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, MarkovChainGraph & g = cache.graph; // Candidate hidden states start at position 1 (position 0 is the seed). - ggml_backend_tensor_set(g.inp_hidden, local_hidden + (size_t)hdim, 0, - sizeof(float) * (size_t)hdim * (size_t)n_cand); + // Uploads and readbacks go through pinned slices, so the call costs one + // wait for the graph and its outputs instead of a round trip per tensor. + const size_t hidden_bytes = sizeof(float) * (size_t)hdim * (size_t)n_cand; + PinnedStage & stage = cache.stage; + stage.reserve(backend, 2 * hidden_bytes + 3 * sizeof(float) * (size_t)n_cand + 4096); + stage.reset(); + stage.set(backend, g.inp_hidden, local_hidden + (size_t)hdim, 0, hidden_bytes); if (want_confidence && g.inp_confidence_hidden) { const float * conf_src = confidence_hidden ? confidence_hidden : local_hidden; - ggml_backend_tensor_set(g.inp_confidence_hidden, conf_src + (size_t)hdim, 0, - sizeof(float) * (size_t)hdim * (size_t)n_cand); + stage.set(backend, g.inp_confidence_hidden, conf_src + (size_t)hdim, 0, hidden_bytes); } - ggml_backend_tensor_set(g.inp_seed, &last_tok, 0, sizeof(int32_t)); + stage.set(backend, g.inp_seed, &last_tok, 0, sizeof(int32_t)); - if (ggml_backend_graph_compute(backend, g.gf) != GGML_STATUS_SUCCESS) { + if (ggml_backend_graph_compute_async(backend, g.gf) != GGML_STATUS_SUCCESS) { + ggml_backend_synchronize(backend); std::fprintf(stderr, "dspark_fused: graph_compute failed\n"); cache.invalidate(); return false; @@ -479,23 +487,35 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, draft_tok.assign((size_t)q_len, 0); draft_tok[0] = last_tok; - // One synchronize instead of n_cand blocking readbacks. + // One synchronize for the graph and every readback. std::vector t_out((size_t)n_cand); std::vector c_out(want_confidence ? (size_t)n_cand : 0); std::vector m_out(want_logit_margin ? (size_t)n_cand : 0); + std::vector t_src((size_t)n_cand, nullptr); + std::vector c_src((size_t)n_cand, nullptr); + std::vector m_src((size_t)n_cand, nullptr); for (int i = 0; i < n_cand; ++i) { - ggml_backend_tensor_get_async(backend, g.toks[(size_t)i], &t_out[i], 0, sizeof(int32_t)); + t_src[(size_t)i] = stage.get(backend, g.toks[(size_t)i], 0, sizeof(int32_t)); if (want_confidence && g.confidence[(size_t)i]) { - ggml_backend_tensor_get_async( - backend, g.confidence[(size_t)i], &c_out[i], 0, sizeof(float)); + c_src[(size_t)i] = stage.get(backend, g.confidence[(size_t)i], 0, sizeof(float)); } if (want_logit_margin && g.logit_margin[(size_t)i]) { - ggml_backend_tensor_get_async( - backend, g.logit_margin[(size_t)i], &m_out[i], 0, - sizeof(float)); + m_src[(size_t)i] = stage.get(backend, g.logit_margin[(size_t)i], 0, sizeof(float)); } } ggml_backend_synchronize(backend); + for (int i = 0; i < n_cand; ++i) { + if (t_src[(size_t)i]) std::memcpy(&t_out[i], t_src[(size_t)i], sizeof(int32_t)); + else ggml_backend_tensor_get(g.toks[(size_t)i], &t_out[i], 0, sizeof(int32_t)); + if (want_confidence && g.confidence[(size_t)i]) { + if (c_src[(size_t)i]) std::memcpy(&c_out[i], c_src[(size_t)i], sizeof(float)); + else ggml_backend_tensor_get(g.confidence[(size_t)i], &c_out[i], 0, sizeof(float)); + } + if (want_logit_margin && g.logit_margin[(size_t)i]) { + if (m_src[(size_t)i]) std::memcpy(&m_out[i], m_src[(size_t)i], sizeof(float)); + else ggml_backend_tensor_get(g.logit_margin[(size_t)i], &m_out[i], 0, sizeof(float)); + } + } for (int i = 0; i < n_cand; ++i) { draft_tok[(size_t)i + 1] = t_out[i]; } @@ -509,6 +529,214 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, return true; } +namespace { + +struct SplitChainCache { + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + ggml_tensor * inp_hidden = nullptr; + ggml_tensor * inp_confidence_hidden = nullptr; + ggml_tensor * inp_seed = nullptr; + std::vector toks; + std::vector confidence; + std::vector arena; + ggml_gallocr_t allocator = nullptr; + ggml_backend_t alloc_backend = nullptr; + PinnedStage stage; + bool built = false; + uint64_t generation = 0; + const void * lm_head = nullptr; + const void * w1 = nullptr; + const void * w2 = nullptr; + const void * id = nullptr; + ggml_backend_t backend = nullptr; + int64_t row_offset = -1; + int64_t n_rows = -1; + int n_candidates = -1; + bool want_confidence = false; + + void invalidate() { + if (ctx) ggml_free(ctx); + ctx = nullptr; + gf = nullptr; + inp_hidden = inp_confidence_hidden = inp_seed = nullptr; + toks.clear(); + confidence.clear(); + built = false; + } +}; + +} // namespace + +bool dspark_markov_chain_vocab_split(const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const DsparkVocabSplit & split, + const float * local_hidden, + int n_candidates, + int32_t last_tok, + bool read, + std::vector & draft_tok, + std::vector * confidence_out, + const float * confidence_hidden) { + if (n_candidates <= 0 || !backend || !lm_head || !dw.dspark.enabled) return false; + if (!split.share || !split.pick || !split.markov_w1 || !split.markov_w2) return false; + if (split.n_rows <= 0 || split.row_offset < 0 || + split.row_offset + split.n_rows > lm_head->ne[1] || + split.markov_w2->ne[1] < split.n_rows) { + return false; + } + const int hdim = dw.n_embd; + if (hdim <= 0 || dw.dspark.markov_rank <= 0) return false; + const bool want_confidence = read && confidence_out != nullptr && + dw.dspark.confidence_w != nullptr && dw.dspark.confidence_b != nullptr && + (dw.dspark.confidence_dim == hdim || + dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank); + if (confidence_out) confidence_out->clear(); + + static thread_local SplitChainCache c; + const uint64_t generation = dspark_drafter_generation(); + const bool reuse = c.built && c.generation == generation && c.lm_head == lm_head && + c.w1 == split.markov_w1 && c.w2 == split.markov_w2 && c.id == split.id && + c.backend == backend && c.row_offset == split.row_offset && c.n_rows == split.n_rows && + c.n_candidates == n_candidates && c.want_confidence == want_confidence; + if (!reuse) { + c.invalidate(); + const size_t arena_size = ggml_tensor_overhead() * (size_t) (64 + 24 * n_candidates) + + ggml_graph_overhead_custom(512, false) + 2 * 1024 * 1024; + if (c.arena.size() < arena_size) c.arena.resize(arena_size); + ggml_init_params ip{}; + ip.mem_size = c.arena.size(); + ip.mem_buffer = c.arena.data(); + ip.no_alloc = true; + c.ctx = ggml_init(ip); + if (!c.ctx) return false; + c.gf = ggml_new_graph_custom(c.ctx, 512, false); + c.inp_hidden = ggml_new_tensor_2d(c.ctx, GGML_TYPE_F32, hdim, n_candidates); + ggml_set_input(c.inp_hidden); + c.inp_seed = ggml_new_tensor_1d(c.ctx, GGML_TYPE_I32, 1); + ggml_set_input(c.inp_seed); + if (want_confidence) { + c.inp_confidence_hidden = ggml_new_tensor_2d(c.ctx, GGML_TYPE_F32, hdim, n_candidates); + ggml_set_input(c.inp_confidence_hidden); + } + ggml_tensor * hidden = split.share(c.ctx, c.inp_hidden); + ggml_tensor * rows = hidden + ? ggml_view_2d(c.ctx, lm_head, lm_head->ne[0], split.n_rows, lm_head->nb[1], + (size_t) split.row_offset * lm_head->nb[1]) + : nullptr; + ggml_tensor * base = rows ? ggml_mul_mat(c.ctx, rows, hidden) : nullptr; // [n_rows, n_candidates] + ggml_tensor * w2 = split.markov_w2->ne[1] == split.n_rows ? split.markov_w2 + : ggml_view_2d(c.ctx, split.markov_w2, split.markov_w2->ne[0], split.n_rows, + split.markov_w2->nb[1], 0); + ggml_tensor * prev_ids = c.inp_seed; + for (int i = 0; base && i < n_candidates; ++i) { + ggml_tensor * prev_emb = ggml_get_rows(c.ctx, split.markov_w1, prev_ids); + ggml_tensor * bias = ggml_mul_mat(c.ctx, w2, prev_emb); + ggml_tensor * base_i = ggml_view_2d(c.ctx, base, split.n_rows, 1, base->nb[1], + (size_t) i * base->nb[1]); + ggml_tensor * tok = split.pick(c.ctx, ggml_add(c.ctx, base_i, bias)); + if (!tok) { base = nullptr; break; } + ggml_set_output(tok); + ggml_build_forward_expand(c.gf, tok); + c.toks.push_back(tok); + if (want_confidence) { + ggml_tensor * hidden_i = ggml_view_2d( + c.ctx, c.inp_confidence_hidden, hdim, 1, c.inp_confidence_hidden->nb[1], + (size_t) i * c.inp_confidence_hidden->nb[1]); + ggml_tensor * conf_in = hidden_i; + if (dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank) { + conf_in = ggml_concat(c.ctx, hidden_i, prev_emb, 0); + } + ggml_tensor * conf = ggml_mul_mat(c.ctx, dw.dspark.confidence_w, conf_in); + conf = ggml_add(c.ctx, conf, ggml_reshape_2d(c.ctx, dw.dspark.confidence_b, 1, 1)); + conf = ggml_sigmoid(c.ctx, conf); + ggml_set_output(conf); + ggml_build_forward_expand(c.gf, conf); + c.confidence.push_back(conf); + } + prev_ids = tok; + } + if (!base) { + std::fprintf(stderr, "dspark_split: graph build failed\n"); + c.invalidate(); + return false; + } + if (c.allocator && c.alloc_backend != backend) { + ggml_gallocr_free(c.allocator); + c.allocator = nullptr; + } + if (!c.allocator) { + c.allocator = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + c.alloc_backend = backend; + } + if (!c.allocator || !ggml_gallocr_alloc_graph(c.allocator, c.gf)) { + std::fprintf(stderr, "dspark_split: gallocr_alloc_graph failed\n"); + c.invalidate(); + return false; + } + c.generation = generation; + c.lm_head = lm_head; + c.w1 = split.markov_w1; + c.w2 = split.markov_w2; + c.id = split.id; + c.backend = backend; + c.row_offset = split.row_offset; + c.n_rows = split.n_rows; + c.n_candidates = n_candidates; + c.want_confidence = want_confidence; + c.built = true; + g_dspark_chain_graph_builds.fetch_add(1, std::memory_order_acq_rel); + } + + // Candidate hidden states start at slot 1 of the padded block (slot 0 is + // the seed). A rank without them leaves the input as it is: the share + // node replaces it with the root's columns. + const size_t hidden_bytes = sizeof(float) * (size_t) hdim * (size_t) n_candidates; + c.stage.reserve(backend, 2 * hidden_bytes + 2 * sizeof(float) * (size_t) n_candidates + 4096); + c.stage.reset(); + if (local_hidden) c.stage.set(backend, c.inp_hidden, local_hidden + (size_t) hdim, 0, hidden_bytes); + if (want_confidence) { + const float * conf_src = confidence_hidden ? confidence_hidden : local_hidden; + if (conf_src) { + c.stage.set(backend, c.inp_confidence_hidden, conf_src + (size_t) hdim, 0, hidden_bytes); + } + } + c.stage.set(backend, c.inp_seed, &last_tok, 0, sizeof(int32_t)); + if (ggml_backend_graph_compute_async(backend, c.gf) != GGML_STATUS_SUCCESS) { + ggml_backend_synchronize(backend); + std::fprintf(stderr, "dspark_split: graph_compute failed\n"); + c.invalidate(); + return false; + } + if (!read) return true; + + std::vector t_src((size_t) n_candidates, nullptr); + std::vector c_src((size_t) n_candidates, nullptr); + for (int i = 0; i < n_candidates; ++i) { + t_src[(size_t) i] = c.stage.get(backend, c.toks[(size_t) i], 0, sizeof(int32_t)); + if (want_confidence) { + c_src[(size_t) i] = c.stage.get(backend, c.confidence[(size_t) i], 0, sizeof(float)); + } + } + ggml_backend_synchronize(backend); + draft_tok.assign((size_t) n_candidates + 1, 0); + draft_tok[0] = last_tok; + std::vector conf_out(want_confidence ? (size_t) n_candidates : 0); + for (int i = 0; i < n_candidates; ++i) { + int32_t t = 0; + if (t_src[(size_t) i]) std::memcpy(&t, t_src[(size_t) i], sizeof(int32_t)); + else ggml_backend_tensor_get(c.toks[(size_t) i], &t, 0, sizeof(int32_t)); + draft_tok[(size_t) i + 1] = t; + if (want_confidence) { + if (c_src[(size_t) i]) std::memcpy(&conf_out[(size_t) i], c_src[(size_t) i], sizeof(float)); + else ggml_backend_tensor_get(c.confidence[(size_t) i], &conf_out[(size_t) i], 0, sizeof(float)); + } + } + if (want_confidence) *confidence_out = std::move(conf_out); + return true; +} + bool dspark_markov_project_topk(const DraftWeights & dw, ggml_backend_t backend, ggml_tensor * lm_head, diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h index c407800de..a2aa17376 100644 --- a/server/src/common/dspark_head.h +++ b/server/src/common/dspark_head.h @@ -4,6 +4,7 @@ #include "internal.h" #include +#include #include namespace luce::common { @@ -47,6 +48,42 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, const float * confidence_hidden = nullptr, std::vector * logit_margin_out = nullptr); +// The fused chain with the vocabulary split across ranks that run the same +// graph in lockstep (a cluster's head and workers). Each rank projects rows +// [row_offset, row_offset + n_rows) of lm_head and of the Markov bias, so the +// lm_head read is divided between the ranks' devices. The model supplies the +// collectives as graph adapters: +// share(ctx, hidden) -> the root's hidden columns on every rank, exactly; +// pick(ctx, logits) -> I32 [1], the argmax over every rank's slice +// (ties to the lowest index, as ggml_argmax). +// The tokens equal the unsplit chain's: each vocabulary row is the same product +// in either form. Every rank must call this the same number of times with the +// same n_candidates, in the same order relative to its other collectives. +// local_hidden (padded as in the fused chain) is read on the root only and may +// be null elsewhere. read = false enqueues the graph and returns at once: a rank +// that only takes part, whose stream orders the graph before its next work. +struct DsparkVocabSplit { + int64_t row_offset = 0; + int64_t n_rows = 0; + ggml_tensor * markov_w1 = nullptr; // [markov_rank, vocab], on `backend`'s device + ggml_tensor * markov_w2 = nullptr; // [markov_rank, >= n_rows]: its first n_rows rows are this rank's + std::function share; + std::function pick; + const void * id = nullptr; // identifies share/pick for the graph cache +}; + +bool dspark_markov_chain_vocab_split(const DraftWeights & dw, + ggml_backend_t backend, + ggml_tensor * lm_head, + const DsparkVocabSplit & split, + const float * local_hidden, + int n_candidates, + int32_t last_tok, + bool read, + std::vector & draft_tok, + std::vector * confidence_out = nullptr, + const float * confidence_hidden = nullptr); + // DDTree candidate generation with the Markov correction: base logits for // all n_tokens positions in ONE lm_head matmul; rows 1..n-1 get the low-rank // previous-token bias chained along the main (argmax) path; top-K extracted diff --git a/server/src/common/host_stream_flag.cpp b/server/src/common/host_stream_flag.cpp new file mode 100644 index 000000000..5aaf1dc7d --- /dev/null +++ b/server/src/common/host_stream_flag.cpp @@ -0,0 +1,53 @@ +#include "host_stream_flag.h" + +#if defined(LUCE_BACKEND_HIP) +#include +#endif + +namespace luce::common { + +#if defined(LUCE_BACKEND_HIP) + +HostStreamFlag::~HostStreamFlag() { + if (host_) (void) hipHostFree(host_); +} + +bool HostStreamFlag::init() { + if (host_) return true; + uint32_t * host = nullptr; + if (hipHostMalloc((void **) &host, 64, hipHostMallocDefault) != hipSuccess) return false; + *host = 0; + uint32_t * dev = nullptr; + if (hipHostGetDevicePointer((void **) &dev, host, 0) != hipSuccess) { + (void) hipHostFree(host); + return false; + } + host_ = host; + dev_ = dev; + return true; +} + +bool HostStreamFlag::enqueue_wait(void * stream, uint32_t value) const { + return dev_ && hipStreamWaitValue32((hipStream_t) stream, dev_, value, + hipStreamWaitValueEq, 0xffffffffu) == hipSuccess; +} + +void HostStreamFlag::signal(uint32_t value) const { + if (host_) __atomic_store_n(host_, value, __ATOMIC_RELEASE); +} + +bool HostStreamFlag::copy_to_device(void * stream, void * dst, const void * src, size_t bytes) { + return hipMemcpyAsync(dst, src, bytes, hipMemcpyHostToDevice, (hipStream_t) stream) == hipSuccess; +} + +#else + +HostStreamFlag::~HostStreamFlag() = default; +bool HostStreamFlag::init() { return false; } +bool HostStreamFlag::enqueue_wait(void *, uint32_t) const { return false; } +void HostStreamFlag::signal(uint32_t) const {} +bool HostStreamFlag::copy_to_device(void *, void *, const void *, size_t) { return false; } + +#endif + +} // namespace luce::common diff --git a/server/src/common/host_stream_flag.h b/server/src/common/host_stream_flag.h new file mode 100644 index 000000000..398a606d8 --- /dev/null +++ b/server/src/common/host_stream_flag.h @@ -0,0 +1,38 @@ +// host_stream_flag.h - a host-written word a GPU stream can wait on. +// +// Lets the host queue GPU work that depends on data a host thread is still +// producing: the stream waits (on the command processor, no kernel) until the +// word reaches a value, and the producing thread raises it when the data is +// in place. HIP only; elsewhere init() fails and callers keep their blocking +// path. +#pragma once + +#include +#include + +namespace luce::common { + +class HostStreamFlag { +public: + HostStreamFlag() = default; + HostStreamFlag(const HostStreamFlag &) = delete; + HostStreamFlag & operator=(const HostStreamFlag &) = delete; + ~HostStreamFlag(); + + bool init(); + bool ready() const { return host_ != nullptr; } + // Queue on `stream` a wait until the word equals `value`. + bool enqueue_wait(void * stream, uint32_t value) const; + // Raise the word to `value` (release: earlier host writes are visible + // to the GPU work released by it). + void signal(uint32_t value) const; + // Queue an asynchronous host-to-device copy on `stream` (the source must + // be pinned and stay valid until the stream ran the copy). + static bool copy_to_device(void * stream, void * dst, const void * src, size_t bytes); + +private: + uint32_t * host_ = nullptr; + uint32_t * dev_ = nullptr; +}; + +} // namespace luce::common diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index 17a3fa30e..1a4ab683f 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -49,9 +49,11 @@ inline ggml_tensor * swiglu_maybe_clamped(ggml_context * ctx, ggml_tensor * up, float clamp) { if (clamp > 1.0e-6f) { - // DeepSeek V4 clamps only the upper side of the gate, and both sides of up. - gate = ggml_clamp(ctx, gate, -INFINITY, clamp); - up = ggml_clamp(ctx, up, -clamp, clamp); + // DeepSeek V4 clamps only the upper side of the gate, and both sides + // of up. One launch: SWIGLU_DS4 applies exactly these two clamps + // before silu(gate) * up (the same fminf/fmaxf bounds and the same + // silu expression as clamp + clamp + swiglu, which cost three). + return ggml_swiglu_ds4_split(ctx, gate, up, clamp); } return ggml_swiglu_split(ctx, gate, up); } @@ -766,7 +768,9 @@ static bool build_batched_routed_graph( std::vector * backend_nodes = nullptr, bool allow_fused_combine = false, bool force_fused_combine = false, - bool defer_route_reduction = false) + bool defer_route_reduction = false, + ggml_tensor * combine_valid = nullptr, + ggml_tensor * combine_ids = nullptr) { const auto track = [&](ggml_tensor * t) -> ggml_tensor * { if (backend_nodes && t) backend_nodes->push_back(t); @@ -891,6 +895,14 @@ static bool build_batched_routed_graph( ggml_tensor * experts = track(apply_scale2(ctx, mixed_mmq(ggml_mul_mat_id(ctx, down_tensor, gu, sel), mixed_mmq_policy), down_scale)); + // The owner's route mask, weights and route sum in one launch (wts are the + // unmasked router weights then; see build_moe_owner_remap). + if (combine_valid && combine_ids) { + GGML_ASSERT(!defer_route_reduction && !tokenwise); + *out_routed = track(ggml_moe_combine_masked(ctx, experts, wts, combine_valid, combine_ids)); + return *out_routed != nullptr; + } + // Weight and sum over experts: [n_embd, n_used, n_tokens] * [1, n_used, n_tokens] if (!defer_route_reduction && allow_fused_combine && (force_fused_combine || moe_hybrid_graph_policy().fused_combine)) { @@ -932,6 +944,10 @@ struct MoeOwnerGraphSpec { std::vector * branch_nodes = nullptr; ggml_tensor * local_ids = nullptr; ggml_tensor * masked_weights = nullptr; + // Set when the route mask is applied inside the owner's combine + // (ggml_moe_combine_masked): masked_weights are then the router weights. + ggml_tensor * combine_valid = nullptr; + ggml_tensor * combine_ids = nullptr; ggml_tensor * output = nullptr; bool available() const { @@ -947,7 +963,8 @@ static bool build_moe_owner_remap( int n_tokens, int dynamic_main_slots_x4, bool main_owner, - MoeOwnerGraphSpec & owner) { + MoeOwnerGraphSpec & owner, + bool mask_in_combine = false) { if (!owner.local_by_global || (int) owner.local_by_global->size() != cfg.n_expert || !owner.local_lut || !owner.valid_lut) { @@ -999,7 +1016,14 @@ static bool build_moe_owner_remap( ctx, *owner.local_lut, global_ids)); mapped = track(ggml_reshape_2d( ctx, mapped, cfg.n_expert_used, n_tokens)); - owner.local_ids = track(ggml_cont(ctx, mapped)); + // The gathered ids are already dense: a copy would be one more launch. + owner.local_ids = ggml_is_contiguous(mapped) ? mapped : track(ggml_cont(ctx, mapped)); + if (mask_in_combine) { + owner.masked_weights = router_weights; + owner.combine_valid = *owner.valid_lut; + owner.combine_ids = global_ids; + return true; + } ggml_tensor * valid = track(ggml_get_rows( ctx, *owner.valid_lut, global_ids)); valid = track(ggml_reshape_2d( @@ -1016,10 +1040,40 @@ static bool prepare_moe_owner_branch( int n_tokens, int dynamic_main_slots_x4, bool main_owner, - MoeOwnerGraphSpec & owner) { + MoeOwnerGraphSpec & owner, + bool mask_in_combine = false) { return !owner.available() || build_moe_owner_remap( ctx, cfg, global_ids, router_weights, n_tokens, - dynamic_main_slots_x4, main_owner, owner); + dynamic_main_slots_x4, main_owner, owner, mask_in_combine); +} + +// Whether build_moe_owner_branch runs this owner token by token. +static bool moe_owner_branch_tokenwise(const MoeOwnerGraphSpec & owner, int n_tokens) { + const ggml_tensor * dispatch_weights = owner.gate_up ? owner.gate_up : owner.gate; + return dispatch_weights && dispatch_weights->type == GGML_TYPE_Q2_0_ROCMFP2 && + !(n_tokens > 1 && moe_hybrid_graph_policy().grouped_mmvq); +} + +// Whether build_batched_routed_graph hands this owner to a coarse DS4 owner +// op (which weights and sums its routes itself). +static bool moe_owner_takes_coarse_op(const MoeOwnerGraphSpec & owner, const MoeLayerDesc & desc) { + const MoeHybridGraphPolicy & policy = moe_hybrid_graph_policy(); + if (!policy.coarse_owner || !owner.down) return false; + const bool down_fpx = owner.down->type == GGML_TYPE_Q3_0_ROCMFPX; + const bool split = policy.coarse_owner_split && owner.gate && owner.up && + owner.gate->type == GGML_TYPE_Q2_0_ROCMFP2 && owner.up->type == GGML_TYPE_Q2_0_ROCMFP2 && down_fpx; + const bool fused = owner.gate_up && desc.ffn_gate_up_exps_s == 1.0f && + owner.gate_up->type == GGML_TYPE_Q2_0_ROCMFP2 && down_fpx; + return split || fused; +} + +// LUCE_MOE_Q8_HANDOFF=1: see build_moe_hybrid_ffn_graph. +static bool moe_q8_handoff_enabled() { + static const bool on = [] { + const char * v = std::getenv("LUCE_MOE_Q8_HANDOFF"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + return on; } static void align_moe_owner_routes( @@ -1062,7 +1116,8 @@ static bool build_moe_owner_branch( cfg.n_embd, cfg.n_ff_exp, cfg.n_expert_used, n_tokens, cfg.swiglu_clamp, cfg.mixed_mmq_policy, &owner.output, tokenwise, owner.branch_nodes, allow_fused_combine, - /*force_fused_combine=*/false, canonical_route_join); + /*force_fused_combine=*/false, canonical_route_join, + owner.combine_valid, owner.combine_ids); } ggml_tensor * build_moe_routed_experts( @@ -1177,6 +1232,11 @@ static ggml_tensor * build_moe_owner_join( if (primary) { ggml_build_forward_expand(schedule_graph, primary); } + if (out.external_join && primary) { + out.main_output = primary; + out.peer_raw = secondary; + return primary; + } ggml_tensor * secondary_ready = ggml_ds4_deferred_peer_copy(ctx, secondary); // Diagnostic host-copy paths may prefill this tensor before its split @@ -1238,6 +1298,7 @@ bool build_moe_hybrid_ffn_graph( out.output = nullptr; out.main_output = nullptr; out.peer_output = nullptr; + out.peer_raw = 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 || @@ -1377,12 +1438,20 @@ bool build_moe_hybrid_ffn_graph( // Keep graph construction order stable: both remaps, then both optional ID // alignments, then both expert branches. + // The route mask rides in the owner's combine on the GPU owners (where + // the fused combine is allowed) when the branch reduces its own routes. + const auto mask_in_combine = [&](const MoeOwnerGraphSpec & owner) { + return allow_fused_combine && + !canonical_route_join && dynamic_main_slots_x4 == 0 && + !moe_owner_branch_tokenwise(owner, n_tokens) && + !moe_owner_takes_coarse_op(owner, desc); + }; if (!prepare_moe_owner_branch( ctx, cfg, global_ids, router_weights, n_tokens, - dynamic_main_slots_x4, true, primary_owner) || + dynamic_main_slots_x4, true, primary_owner, mask_in_combine(primary_owner)) || !prepare_moe_owner_branch( ctx, cfg, global_ids, router_weights, n_tokens, - dynamic_main_slots_x4, false, secondary_owner) || + dynamic_main_slots_x4, false, secondary_owner, mask_in_combine(secondary_owner)) || !prepare_moe_owner_branch( ctx, cfg, global_ids, router_weights, n_tokens, /*dynamic_main_slots_x4=*/0, false, streamed_owner)) { @@ -1391,11 +1460,29 @@ bool build_moe_hybrid_ffn_graph( align_moe_owner_routes(ctx, n_tokens, primary_owner); align_moe_owner_routes(ctx, n_tokens, secondary_owner); align_moe_owner_routes(ctx, n_tokens, streamed_owner); + // LUCE_MOE_Q8_HANDOFF=1 (secondary owner on a GPU): the secondary owner gets + // its input as the q8_1 rows its quantized matvecs read anyway, quantized + // once on the main owner. 5,120 floats become 160 blocks of 36 bytes per + // token, 3.6x fewer bytes across the link, and the secondary owner + // skips its own quantization. Same kernel, same bytes: the result is unchanged. + // Up to 8 tokens: every quantized type takes MMVQ there (decode, verify). + ggml_tensor * inp_secondary = inp; + const ggml_tensor * secondary_w = secondary_owner.gate_up ? secondary_owner.gate_up : secondary_owner.gate; + if (moe_q8_handoff_enabled() && schedule_graph && secondary_owner.available() && secondary_w && + ggml_is_quantized(secondary_w->type) && inp->type == GGML_TYPE_F32 && + inp->ne[0] % 512 == 0 && inp->ne[1] == n_tokens && n_tokens <= 8) { + ggml_tensor * src = ggml_is_contiguous(inp) ? inp : ggml_cont(ctx, inp); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, inp->ne[0], n_tokens); + inp_secondary = ggml_cpy(ctx, src, rows); + if (src != inp) out.hot_remap_nodes.push_back(src); + out.hot_remap_nodes.push_back(inp_secondary); + ggml_build_forward_expand(schedule_graph, inp_secondary); + } if (!build_moe_owner_branch( ctx, cfg, desc, inp, n_tokens, canonical_route_join, allow_fused_combine, primary_owner) || !build_moe_owner_branch( - ctx, cfg, desc, inp, n_tokens, canonical_route_join, + ctx, cfg, desc, inp_secondary, n_tokens, canonical_route_join, allow_fused_combine, secondary_owner) || !build_moe_owner_branch( ctx, cfg, desc, inp, n_tokens, canonical_route_join, @@ -1784,7 +1871,8 @@ bool eval_moe_hybrid_ffn_single( const int32_t gid = selected_ids[i]; // Cold owner None: routes masked to -1 by the cluster runtime are // evaluated elsewhere and contribute zero here. - if (gid < 0 && storage.cold_backend_kind == MoeHybridColdBackend::None) { + if (gid < 0 && (storage.cold_backend_kind == MoeHybridColdBackend::None || + storage.foreign_routes)) { continue; } if (gid < 0 || gid >= (int32_t)storage.hot_local_by_global.size()) { @@ -4393,7 +4481,8 @@ bool eval_moe_hybrid_ffn_gpu_resident( for (int i = 0; i < n_selected; ++i) { const int32_t gid = selected_ids[i]; - if (gid < 0 && storage.cold_backend_kind == MoeHybridColdBackend::None) continue; + if (gid < 0 && (storage.cold_backend_kind == MoeHybridColdBackend::None || + storage.foreign_routes)) continue; if (gid < 0 || gid >= (int32_t)storage.hot_local_by_global.size()) return false; const int32_t hot_local = storage.hot_local_by_global[(size_t)gid]; if (hot_local >= 0) { diff --git a/server/src/common/moe_hybrid_ffn_eval.h b/server/src/common/moe_hybrid_ffn_eval.h index 097e9acec..64bfcac46 100644 --- a/server/src/common/moe_hybrid_ffn_eval.h +++ b/server/src/common/moe_hybrid_ffn_eval.h @@ -169,6 +169,12 @@ struct MoeHybridGraphInputs { // neither tensor changes expert placement or routing semantics. ggml_tensor * main_output = nullptr; ggml_tensor * peer_output = nullptr; + // Input, set by the caller before building: return the owners' partials + // unjoined (main_output on the main backend, peer_raw on the secondary + // owner's backend, output = main_output) so the caller can combine them + // itself -- the cluster hybrid exchange adds them on the secondary owner. + bool external_join = false; + ggml_tensor * peer_raw = nullptr; // Backend-affinity hints consumed after the multi-backend scheduler is // created. Keeping every intermediate of a routed branch on its weight // backend avoids gate/up -> activation -> down ping-pong copies. diff --git a/server/src/common/moe_hybrid_placement.cpp b/server/src/common/moe_hybrid_placement.cpp index efb42a47c..b577a8486 100644 --- a/server/src/common/moe_hybrid_placement.cpp +++ b/server/src/common/moe_hybrid_placement.cpp @@ -695,10 +695,15 @@ bool MoeExpertOwnership::load_json(const std::string & path, int layers, int exp tmp.owner[i] = Secondary; } else if (name == "stream") { tmp.owner[i] = Stream; + } else if (name == "remote") { + // Another expert-parallel cluster rank evaluates this expert. + tmp.owner[i] = Stream; + if (tmp.remote.empty()) tmp.remote.assign(tmp.owner.size(), 0); + tmp.remote[i] = 1; } else { if (err) { *err = "expert placement entry " + std::to_string(i) + - " must be \"primary\", \"secondary\" or \"stream\""; + " must be \"primary\", \"secondary\", \"stream\" or \"remote\""; } return false; } diff --git a/server/src/common/moe_hybrid_placement.h b/server/src/common/moe_hybrid_placement.h index 159399b2b..709f4f872 100644 --- a/server/src/common/moe_hybrid_placement.h +++ b/server/src/common/moe_hybrid_placement.h @@ -121,6 +121,9 @@ struct MoeExpertOwnership { int n_expert = 0; std::vector owner; // [n_layer * n_expert] Owner std::vector pinned; // [n_layer * n_expert] 1 = keep on the primary + // [n_layer * n_expert] 1 = "remote": another cluster rank evaluates it. + // Held as Stream here; the cluster runtime keeps it out of every tier. + std::vector remote; bool init(int n_layer, int n_expert, Owner fill); static bool load_json(const std::string & path, int n_layer, int n_expert, diff --git a/server/src/common/moe_hybrid_storage.cpp b/server/src/common/moe_hybrid_storage.cpp index d46eac8cd..798556cac 100644 --- a/server/src/common/moe_hybrid_storage.cpp +++ b/server/src/common/moe_hybrid_storage.cpp @@ -337,7 +337,29 @@ static bool assign_cold_experts(const MoeHybridConfig & cfg, int il, } int owned = 0; for (uint8_t hot : is_hot) owned += hot ? 1 : 0; - dst.n_streamed = cfg.n_expert - owned - (int)dst.cold_expert_ids.size(); + int foreign = 0; + if (!cfg.foreign_expert_ids.empty()) { + if (cfg.foreign_expert_ids.size() != (size_t)cfg.n_layer) { + if (err) *err = "foreign expert lists need one list per layer"; + return false; + } + dst.foreign_by_global.assign((size_t)cfg.n_expert, 0); + dst.foreign_routes = true; + for (int32_t expert : cfg.foreign_expert_ids[(size_t)il]) { + if (expert < 0 || expert >= cfg.n_expert || is_hot[(size_t)expert] || + dst.cold_local_by_global[(size_t)expert] >= 0 || + dst.foreign_by_global[(size_t)expert]) { + if (err) { + *err = "foreign expert " + std::to_string(expert) + " in layer " + + std::to_string(il) + " is out of range, resident, or listed twice"; + } + return false; + } + dst.foreign_by_global[(size_t)expert] = 1; + ++foreign; + } + } + dst.n_streamed = cfg.n_expert - owned - (int)dst.cold_expert_ids.size() - foreign; return true; } diff --git a/server/src/common/moe_hybrid_storage.h b/server/src/common/moe_hybrid_storage.h index 805311378..fd7896bbe 100644 --- a/server/src/common/moe_hybrid_storage.h +++ b/server/src/common/moe_hybrid_storage.h @@ -106,11 +106,19 @@ struct MoeHybridLayerStorage { // the whole cold set when the cold stack is not materialized, otherwise // what an explicit cold-stack list leaves out. int n_streamed = 0; + // Experts another cluster rank evaluates (MoeHybridConfig:: + // foreign_expert_ids): 1 per such expert. Their routes are masked to -1 + // before the owners run, which the evaluators then skip. + std::vector foreign_by_global; + bool foreign_routes = false; bool is_streamed(int global_expert) const { if (global_expert < 0 || (size_t) global_expert >= hot_local_by_global.size() || hot_local_by_global[(size_t) global_expert] >= 0) { return false; } + if (!foreign_by_global.empty() && foreign_by_global[(size_t) global_expert]) { + return false; + } const bool cold_stack = down_cold || gate_up_cold; return cold_stack ? cold_local_by_global[(size_t) global_expert] < 0 : n_streamed > 0 && cold_backend_kind != MoeHybridColdBackend::None; diff --git a/server/src/common/moe_hybrid_types.h b/server/src/common/moe_hybrid_types.h index 7f28882bc..0cb4dade1 100644 --- a/server/src/common/moe_hybrid_types.h +++ b/server/src/common/moe_hybrid_types.h @@ -64,6 +64,10 @@ struct MoeHybridConfig { // cold-stack experts, and the experts in neither stack are streamed from // the model file on demand. std::vector> cold_expert_ids; + // Experts another process evaluates (an expert-parallel cluster rank's + // foreign shard, server/src/cluster): never resident here, never + // streamed; their routes arrive masked to id -1 and contribute zero. + std::vector> foreign_expert_ids; // Cold owner None has no cold experts, so nothing to materialize whatever // materialize_cold_experts says. diff --git a/server/src/common/pinned_stage.h b/server/src/common/pinned_stage.h new file mode 100644 index 000000000..5fce857b2 --- /dev/null +++ b/server/src/common/pinned_stage.h @@ -0,0 +1,130 @@ +// pinned_stage.h - pinned host staging for a graph launch's small transfers. +// +// A tensor_set or tensor_get from pageable memory is a full round trip on +// HIP: the call returns only after the copy ran, and a speculative step does a +// few dozen of them (inputs, readbacks), each ~30 us with the GPU idle. Copies +// from pinned memory are truly asynchronous, so a step pays one wait for all +// of them instead of one per tensor. +// +// Usage per launch: reset(), take() slices for each transfer, issue the async +// copies on the backend that owns the tensors, synchronize that backend once, +// then read the readback slices. A slice stays valid until the next reset(), +// and must not be reused before the copies that read or write it completed. +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include + +namespace luce::common { + +class PinnedStage { +public: + PinnedStage() = default; + PinnedStage(const PinnedStage &) = delete; + PinnedStage & operator=(const PinnedStage &) = delete; + // Moving hands the block over; the caller has no copies in flight. + PinnedStage(PinnedStage && other) noexcept { take_from(other); } + PinnedStage & operator=(PinnedStage && other) noexcept { + if (this != &other) { + release(); + take_from(other); + } + return *this; + } + ~PinnedStage() { release(); } + + // Room for `bytes` in one launch. Growing frees the old block, so the + // caller must not have copies in flight from it. + bool reserve(ggml_backend_t backend, size_t bytes) { + if (buf_ && size_ >= bytes && backend_ == backend) return true; + release(); + ggml_backend_dev_t dev = backend ? ggml_backend_get_device(backend) : nullptr; + ggml_backend_buffer_type_t buft = dev ? ggml_backend_dev_host_buffer_type(dev) : nullptr; + if (!buft) return false; + const size_t want = std::max(bytes, (size_t) 1 << 20); + buf_ = ggml_backend_buft_alloc_buffer(buft, want); + if (!buf_) return false; + base_ = static_cast(ggml_backend_buffer_get_base(buf_)); + if (!base_) { release(); return false; } + size_ = want; + backend_ = backend; + return true; + } + + void reset() { used_ = 0; } + + // A 64-byte aligned slice, or nullptr when the stage is missing or full + // (the caller then takes its synchronous path). + void * take(size_t bytes) { + const size_t off = (used_ + 63) & ~(size_t) 63; + if (!base_ || off + bytes > size_) return nullptr; + used_ = off + bytes; + return base_ + off; + } + + // Asynchronous upload of `data` into `t` on `backend` through a slice. + // Falls back to the blocking set when no slice is left or the tensor does + // not live in `backend`'s device memory. + void set(ggml_backend_t backend, ggml_tensor * t, const void * data, + size_t offset, size_t bytes) { + if (!t || !t->buffer || bytes == 0) return; + void * slice = owned_by(backend, t) ? take(bytes) : nullptr; + if (!slice) { + ggml_backend_tensor_set(t, data, offset, bytes); + return; + } + std::memcpy(slice, data, bytes); + ggml_backend_tensor_set_async(backend, t, slice, offset, bytes); + } + + // Asynchronous readback of `t` into a slice; the bytes are valid after + // the backend was synchronized. nullptr: read it with the blocking get. + const void * get(ggml_backend_t backend, const ggml_tensor * t, + size_t offset, size_t bytes) { + if (!t || !t->buffer || bytes == 0 || !owned_by(backend, t)) return nullptr; + void * slice = take(bytes); + if (!slice) return nullptr; + ggml_backend_tensor_get_async(backend, t, slice, offset, bytes); + return slice; + } + +private: + static bool owned_by(ggml_backend_t backend, const ggml_tensor * t) { + const ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; + return backend && buf && !ggml_backend_buffer_is_host(buf) && + ggml_backend_buffer_get_type(buf) == ggml_backend_get_default_buffer_type(backend); + } + + void take_from(PinnedStage & other) { + buf_ = other.buf_; + base_ = other.base_; + size_ = other.size_; + used_ = other.used_; + backend_ = other.backend_; + other.buf_ = nullptr; + other.base_ = nullptr; + other.size_ = other.used_ = 0; + other.backend_ = nullptr; + } + + void release() { + if (buf_) ggml_backend_buffer_free(buf_); + buf_ = nullptr; + base_ = nullptr; + size_ = used_ = 0; + backend_ = nullptr; + } + + ggml_backend_buffer_t buf_ = nullptr; + uint8_t * base_ = nullptr; + size_t size_ = 0; + size_t used_ = 0; + ggml_backend_t backend_ = nullptr; +}; + +} // namespace luce::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index ac4bf5ea6..f64a69cf9 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -2,6 +2,8 @@ #include "deepseek4_roctx.h" #include "deepseek4_backend.h" +#include "deepseek4_cluster.h" +#include "cluster/cluster_decision_hooks.h" #include "thinking_budget.h" #include "deepseek4_internal.h" #include "deepseek4_image_spans.h" @@ -1140,6 +1142,222 @@ void log_deepseek4_step_telemetry(const char * phase, DeepSeek4Backend::DeepSeek4Backend(DeepSeek4BackendConfig cfg) : cfg_(std::move(cfg)) {} +// ─── Expert-parallel cluster ─────────────────────────────────────────── + +bool DeepSeek4Backend::cluster_active() const { + return hooks_ != nullptr && hooks_->is_cluster(); +} + +bool DeepSeek4Backend::cluster_worker() const { + return cluster_active() && !hooks_->is_head(); +} + +uint64_t DeepSeek4Backend::cluster_placement_hash() const { + return cluster_ ? cluster_->placement.hash() : 0; +} + +uint64_t DeepSeek4Backend::cluster_resident_expert_bytes() const { + return cluster_ ? cluster_->resident_expert_bytes : 0; +} + +bool DeepSeek4Backend::cluster_ingraph_allreduce() const { + return cluster_ && ds4_cluster_fused_graph_available(cluster_.get()); +} + +bool DeepSeek4Backend::set_cluster(const cluster::ClusterConfig * cfg, + cluster::IClusterComm * comm) { + if (!cfg || !cfg->enabled()) { + if (cluster_ && !w_.layers.empty()) { + std::fprintf(stderr, + "[deepseek4-cluster] set_cluster(nullptr) after the sharded model " + "was loaded; keeping the runtime (reload is not supported)\n"); + return false; + } + cluster_.reset(); + cache_.cluster_rt = nullptr; + return true; + } + if (comm && (comm->size() != cfg->size || comm->rank() != cfg->rank)) { + std::fprintf(stderr, + "[deepseek4-cluster] communicator %d/%d does not match config %d/%d\n", + comm->rank(), comm->size(), cfg->rank, cfg->size); + return false; + } + if (cluster_) { + // Second call (head/worker after RCCL is up): the placement-defining + // fields must not change; the communicator is attached. + const cluster::ClusterConfig & have = cluster_->cfg_storage; + if (have.rank != cfg->rank || have.size != cfg->size || + have.placement_source != cfg->placement_source || + have.placement_file != cfg->placement_file || + have.replicate_hot != cfg->replicate_hot || + have.shared_expert != cfg->shared_expert) { + std::fprintf(stderr, + "[deepseek4-cluster] set_cluster(): placement-defining config " + "changed after the expert shard was built; refusing\n"); + return false; + } + cluster_->cfg_storage = *cfg; + cluster_->cfg = &cluster_->cfg_storage; + if (comm) cluster_->comm = comm; + cluster_->trace = cluster::cluster_env_trace(); + cache_.cluster_rt = cluster_.get(); + if (comm && cfg->size > 1 && !cluster_->fast) { + // Every rank must reduce the same way or the collectives desync: + // keep the fast path only if all ranks brought it up. Every rank + // takes part in the agreement, even one with the switch off. + const char * fe = std::getenv("LUCE_CLUSTER_FAST_REDUCE"); + cluster::FastReduce::Config fc; + std::unique_ptr fast; + std::string ferr = "LUCE_CLUSTER_FAST_REDUCE is off"; + bool up = false; + if (fe && std::atoi(fe) != 0) { + fc.rank = cfg->rank; + fc.size = cfg->size; + fc.hca = cfg->ib_hca.substr(0, cfg->ib_hca.find(':')); + fc.gid_index = cfg->gid_index > 0 ? cfg->gid_index : 1; + fc.bootstrap_host = cfg->head_host; + fc.bootstrap_port = cfg->head_port + 100; + fc.max_elems = 81920; // q=6 verify: 6 x 5120 reduce, 6 x 12288 Engram broadcast + fc.slots = 256; // ~6 verify steps of lookahead + if (const char * e = std::getenv("LUCE_CLUSTER_FAST_REDUCE_ELEMS")) fc.max_elems = std::max(1024, std::atoi(e)); + if (const char * e = std::getenv("LUCE_CLUSTER_FAST_REDUCE_SLOTS")) fc.slots = std::max(8, std::atoi(e)); + fast = std::make_unique(); + ferr.clear(); + up = fast->init(fc, &ferr); + } + std::vector ups; + std::string aerr; + if (!ds4_cluster_allgather_i32(*cluster_, backend_, up ? 1 : 0, ups, &aerr)) { + std::fprintf(stderr, "[deepseek4-cluster] fast reduce agreement failed: %s\n", aerr.c_str()); + return false; + } + const bool all_up = std::all_of(ups.begin(), ups.end(), [](int32_t v) { return v != 0; }); + if (all_up) { + cluster_->fast = std::move(fast); + std::fprintf(stderr, "[deepseek4-cluster] rank %d fast reduce up on %s gid %d " + "(%d floats x %d slots), bootstrap :%d\n", cfg->rank, fc.hca.c_str(), + fc.gid_index, fc.max_elems, fc.slots, fc.bootstrap_port); + } else if (fast) { + std::fprintf(stderr, "[deepseek4-cluster] rank %d fast reduce unavailable, keeping " + "RCCL: %s\n", cfg->rank, up ? "a peer could not bring it up" : ferr.c_str()); + } + } + if (comm && (!cluster_agree_prefill_chunk() || !cluster_agree_drafter())) return false; + return true; + } + if (backend_ && !w_.layers.empty()) { + std::fprintf(stderr, "[deepseek4-cluster] set_cluster() must run before init()\n"); + return false; + } + cluster_ = std::make_unique(); + cluster_->cfg_storage = *cfg; + cluster_->cfg = &cluster_->cfg_storage; + cluster_->comm = comm; + cluster_->trace = cluster::cluster_env_trace(); + return true; +} + +// Drafting changes the shape of every verify step, and a split drafter +// all-reduces inside its own graph. Rank 0 decides whether a request drafts +// (no drafter there: every rank decodes AR), so every other rank must draft, +// and split, the way rank 0 does, or the collectives desync on the first +// speculative request. +bool DeepSeek4Backend::cluster_agree_drafter() { + const int32_t mine = !spec_enabled_ || !spec_drafter_ ? 0 : spec_drafter_->split_size > 1 ? 2 : 1; + std::vector all; + std::string err; + if (!ds4_cluster_allgather_i32(*cluster_, backend_, mine, all, &err)) { + std::fprintf(stderr, "[deepseek4-cluster] drafter agreement failed: %s\n", err.c_str()); + return false; + } + static const char * const kMode[] = {"no drafter", "a whole drafter", "a split drafter"}; + const int32_t head = all.empty() ? 0 : all[0]; + for (size_t r = 1; head != 0 && r < all.size(); ++r) { + if (all[r] == head) continue; + std::fprintf(stderr, "[deepseek4-cluster] rank 0 has %s, rank %zu has %s: start every rank " + "with the same drafter and draft device\n", kMode[std::clamp(head, 0, 2)], + r, kMode[std::clamp(all[r], 0, 2)]); + return false; + } + return true; +} + +// The batched prefill chunk is fitted to each rank's free memory at load, +// and every chunk is a collective: all ranks must use the same one. +bool DeepSeek4Backend::cluster_agree_prefill_chunk() { + std::string err; + std::vector caps; + if (!ds4_cluster_allgather_i32(*cluster_, backend_, hybrid_prefill_chunk_cap_, caps, &err)) { + std::fprintf(stderr, "[deepseek4-cluster] prefill chunk agreement failed: %s\n", err.c_str()); + return false; + } + int agreed = 0; // 0 = uncapped + for (int32_t c : caps) { + if (c > 0) agreed = agreed > 0 ? std::min(agreed, (int) c) : (int) c; + } + std::vector long_chunks; + if (!ds4_cluster_allgather_i32(*cluster_, backend_, hybrid_long_context_chunk_, long_chunks, &err)) { + std::fprintf(stderr, "[deepseek4-cluster] prefill chunk agreement failed: %s\n", err.c_str()); + return false; + } + int long_chunk = hybrid_long_context_chunk_; + for (int32_t c : long_chunks) long_chunk = std::min(long_chunk, (int) c); + if (agreed != hybrid_prefill_chunk_cap_ || long_chunk != hybrid_long_context_chunk_) { + std::fprintf(stderr, "[deepseek4-cluster] rank %d prefill chunk cap %d -> %d, long-context " + "chunk %d -> %d (agreed across ranks)\n", cluster_->rank(), + hybrid_prefill_chunk_cap_, agreed, hybrid_long_context_chunk_, long_chunk); + } + // Every prefill chunk reads all of a rank's experts once, so the smallest + // rank's chunk sets the whole cluster's prompt speed: name it. + for (size_t r = 0; agreed > 0 && cluster_->rank() == 0 && r < caps.size(); ++r) { + if (caps[r] == agreed && agreed < kDs4ClusterPipelineMinBand) { + std::fprintf(stderr, "[deepseek4-cluster] rank %zu holds every prefill chunk to %d rows; a " + "smaller LUCE_EXPERT_BUDGET_MB on that rank leaves room for longer chunks\n", + r, agreed); + break; + } + } + hybrid_prefill_chunk_cap_ = agreed; + hybrid_long_context_chunk_ = long_chunk; + // The drafter swap's long-prompt chunk is agreed the same way. Every rank + // swaps for the same requests (the decision depends on the prompt alone), + // so the band sizes stay collective. A rank without the swap reports its + // ordinary cap, which keeps the agreed chunk within its memory. + const int32_t without = draft_swap_.host ? draft_swap_.cap_without : agreed; + std::vector withouts; + if (!ds4_cluster_allgather_i32(*cluster_, backend_, without, withouts, &err)) { + std::fprintf(stderr, "[deepseek4-cluster] prefill chunk agreement failed: %s\n", err.c_str()); + return false; + } + int agreed_without = 0; // 0 = uncapped + bool all_swap = true; + for (int32_t c : withouts) { + if (c > 0) agreed_without = agreed_without > 0 ? std::min(agreed_without, (int) c) : (int) c; + } + std::vector swaps; + if (!ds4_cluster_allgather_i32(*cluster_, backend_, draft_swap_.host ? 1 : 0, swaps, &err)) { + std::fprintf(stderr, "[deepseek4-cluster] prefill chunk agreement failed: %s\n", err.c_str()); + return false; + } + for (int32_t v : swaps) all_swap = all_swap && v != 0; + if (draft_swap_.host) { + if (!all_swap) { + // One rank cannot swap: none does, or the bands would differ. + draft_swap_release(); + std::fprintf(stderr, "[deepseek4-cluster] rank %d draft swap off (not every rank " + "can swap)\n", cluster_->rank()); + } else { + draft_swap_.cap_with = agreed; + draft_swap_.cap_without = agreed_without; + std::fprintf(stderr, "[deepseek4-cluster] rank %d draft swap: prefill chunk cap %d with " + "the drafter, %d with it swapped out (agreed across ranks)\n", + cluster_->rank(), agreed, agreed_without); + } + } + return true; +} + DeepSeek4Backend::~DeepSeek4Backend() { shutdown(); } @@ -1970,8 +2188,36 @@ bool DeepSeek4Backend::apply_expert_ownership(bool secondary_owner, int secondar } } } + // Cluster rank: the rank map (--cluster-expert-placement) decides which + // rank evaluates each expert; this rank's tiers hold only its own share. + // Every other expert stays out of all tiers (never resident, never + // streamed) and its routes are masked before the owners run. + std::vector> foreign((size_t) w_.n_layer); + if (cluster_) { + for (int il = 0; il < w_.n_layer; ++il) { + for (int e = 0; e < w_.n_expert; ++e) { + const size_t i = (size_t) il * (size_t) w_.n_expert + (size_t) e; + const bool here = ds4_cluster_resident(*cluster_, il, e); + const bool marked_remote = !own.remote.empty() && own.remote[i]; + if (here && marked_remote) { + std::fprintf(stderr, "[deepseek4-cluster] %s marks layer %d expert %d remote, " + "but the rank map gives it to rank %d\n", + cfg_.expert_placement_path.c_str(), il, e, cluster_->rank()); + return false; + } + if (here) { + if (own.at(il, e) == MoeExpertOwnership::Stream) own.set(il, e, MoeExpertOwnership::Secondary); + } else { + own.set(il, e, MoeExpertOwnership::Stream); + foreign[(size_t) il].push_back(e); + } + } + } + } for (size_t i = 0; i < w_.protected_experts.size(); ++i) { - if (w_.protected_experts[i]) own.pin_primary((int) (i / (size_t) w_.n_expert), (int) (i % (size_t) w_.n_expert)); + const int il = (int) (i / (size_t) w_.n_expert), e = (int) (i % (size_t) w_.n_expert); + if (cluster_ && !ds4_cluster_resident(*cluster_, il, e)) continue; + if (w_.protected_experts[i]) own.pin_primary(il, e); } MoeHybridRoutingStats usage; const char * usage_path = ds4_usage_profile_path(); @@ -1981,6 +2227,25 @@ bool DeepSeek4Backend::apply_expert_ownership(bool secondary_owner, int secondar return false; } + if (cluster_) { + // Nothing of this rank's share may stream: the streamed tier would + // also serve the other rank's experts (and a rank's sparse model copy + // holds none of them). + int demoted = 0; + for (int il = 0; il < w_.n_layer; ++il) { + for (int e = 0; e < w_.n_expert; ++e) { + if (ds4_cluster_resident(*cluster_, il, e) && + own.at(il, e) == MoeExpertOwnership::Stream) ++demoted; + } + } + if (demoted > 0) { + std::fprintf(stderr, "[deepseek4-cluster] %d of this rank's experts do not fit its " + "primary and secondary devices; lower LUCE_EXPERT_BUDGET_MB or " + "--max-ctx, or give the rank fewer experts\n", demoted); + return false; + } + hybrid_cfg.foreign_expert_ids = foreign; + } moe_placement_.hot_expert_ids = own.expert_ids(MoeExpertOwnership::Primary); moe_placement_.total_hot = 0; for (int il = 0; il < w_.n_layer; ++il) { @@ -2260,8 +2525,24 @@ bool DeepSeek4Backend::load_spec_drafter() { } auto drafter = std::make_unique(); + // LUCE_CLUSTER_DRAFT_SPLIT=1 (two ranks, drafter beside the target): each + // rank loads its half of the routed experts and drafts with its half of + // the heads; the partial sums meet in the drafter graph's all-reduces. + DSparkLoadSplit draft_split{}; + if (env_flag_enabled("LUCE_CLUSTER_DRAFT_SPLIT") && cluster_ && cluster_->cfg && + cluster_->cfg->size == 2) { + if (draft_backend != backend_) { + std::fprintf(stderr, + "[deepseek4] LUCE_CLUSTER_DRAFT_SPLIT needs the drafter on the " + "target GPU (no --draft-device); drafting whole\n"); + } else { + draft_split.rank = cluster_->cfg->rank; + draft_split.size = cluster_->cfg->size; + } + } if (!load_deepseek4_dspark_drafter( - spec_draft_path_, draft_backend, *drafter)) { + spec_draft_path_, draft_backend, *drafter, + draft_split.size > 1 ? &draft_split : nullptr)) { std::fprintf(stderr, "[deepseek4] DSpark drafter load FAILED: %s\n", deepseek4_dspark_last_error()); if (spec_backend_) { @@ -2280,6 +2561,14 @@ bool DeepSeek4Backend::load_spec_drafter() { return false; } + if (drafter->split_size > 1) { + drafter->cluster_rt = cluster_.get(); + std::fprintf(stderr, + "[deepseek4] DSpark split: rank %d of %d drafts with experts %d..%d of %d " + "and half the heads\n", + drafter->split_rank, drafter->split_size, drafter->expert_first, + drafter->expert_first + drafter->expert_local, drafter->core.n_expert); + } const DSparkDrafter & d = *drafter; // A drafter is trained against one target: its hyper-connection rules // (the staggered pre-mix of V4.1) must match the target's. @@ -2547,6 +2836,7 @@ bool DeepSeek4Backend::supports_batched_spec_feature_capture( } bool DeepSeek4Backend::init() { + if (cfg_.cluster.enabled() && !set_cluster(&cfg_.cluster, nullptr)) return false; // --draft selects DSpark directly; LUCE_DS4_SPEC + LUCE_DS4_DRAFT remain // the environment spelling of the same request. spec_requested_ = !cfg_.draft_path.empty() || env_flag_enabled("LUCE_DS4_SPEC"); @@ -2795,7 +3085,17 @@ bool DeepSeek4Backend::init() { } if (!init_streamed_expert_tier() || !check_device_headroom()) return false; if (!size_hybrid_prefill_chunk()) return false; + // A cluster rank's prefill chunks, the swapped-out one included, are + // agreed with the other ranks once the communicator is up. if (!setup_draft_swap()) return false; + if (cluster_) { + if (!moe_hybrid_ || !expert_backend_) { + std::fprintf(stderr, "[deepseek4-cluster] a cluster rank needs the hybrid expert " + "tiers (--expert-device for the secondary owner)\n"); + return false; + } + cache_.cluster_rt = cluster_.get(); + } image_capable_ = vision_ != nullptr; return true; } @@ -3334,6 +3634,15 @@ bool DeepSeek4Backend::init_hybrid_model() { TargetLoadPlan plan; plan.skip_expert_tensors = true; plan.load_ds4_image_bias = !cfg_.mmproj_path.empty(); + // A cluster rank loads only the dense slices its head split and + // shared-expert shard read, which leaves room for more hot experts. + if (cluster_ && cluster_->cfg && cluster_->cfg->size > 1) { + plan.slice_rank = cluster_->cfg->rank; + plan.slice_ranks = cluster_->cfg->size; + plan.slice_attention = ds4_cluster_attention_parallel_enabled(); + plan.slice_shared_expert = + cluster_->cfg->shared_expert == cluster::SharedExpertMode::Shard; + } if (!load_deepseek4_gguf_partial(cfg_.model_path, backend_, plan, w_)) { std::fprintf(stderr, "[deepseek4] failed to partially load model for hybrid mode: %s (%s)\n", cfg_.model_path.c_str(), luce_last_error()); @@ -3479,6 +3788,21 @@ bool DeepSeek4Backend::init_hybrid_model() { hybrid_cfg.cold_expert_backend = MoeHybridColdBackend::Gpu; } const int host_spill_gpu = same_runtime_tp ? tp.secondary_gpu : -1; + if (cluster_) { + if (!inprocess_tp) { + std::fprintf(stderr, "[deepseek4-cluster] a cluster rank needs the secondary " + "expert owner (--expert-device)\n"); + return fail_hybrid_init(); + } + const char * hotness_env = std::getenv("LUCE_DS4_HOTNESS_CSV"); + const std::string hotness_csv = hotness_env ? hotness_env : ""; + if (!ds4_cluster_build_placement(*cluster_->cfg, w_, + hotness_csv.empty() ? nullptr : &hotness_csv, + *cluster_, &err)) { + std::fprintf(stderr, "[deepseek4-cluster] placement failed: %s\n", err.c_str()); + return fail_hybrid_init(); + } + } if (!apply_expert_ownership(inprocess_tp, host_spill_gpu, hybrid_cfg)) { if (host_spill_gpu >= 0) ggml_backend_cuda_set_host_spill(host_spill_gpu, 0, 0); return fail_hybrid_init(); @@ -4501,8 +4825,15 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, } } + // Cluster lockstep: rank 0 samples and + // broadcasts every token with its stop flags; workers adopt them, so all + // ranks leave the loop on the same step for the same reason. + const bool cluster = cluster_active(); + const bool worker = cluster_worker(); + const uint64_t cluster_req = hooks_ ? hooks_->current_request() : 0; + for (int generated = 0; generated < n_gen; generated++) { - if (io.is_cancelled()) break; + if (!cluster && io.is_cancelled()) break; // Get last logits and sample std::vector logits; @@ -4558,8 +4889,25 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, } int32_t next_token = 0; + uint8_t decision_flags = 0; const auto sample_t0 = Clock::now(); - if (process_logits) { + if (worker) { + // Workers never sample: the head's token arrives in DecisionMsg. + std::string err; + if (!hooks_->decide_next_token(cluster_req, (uint32_t) generated, + logits.data(), w_.n_vocab, sampler_, + /*is_eos_candidate_unused=*/false, + next_token, decision_flags, &err)) { + std::fprintf(stderr, "[deepseek4] cluster decision missing at step %d: %s\n", + generated, err.c_str()); + return false; + } + if (next_token < 0 || next_token >= w_.n_vocab) { + std::fprintf(stderr, "[deepseek4] cluster decision token %d out of range at step %d\n", + next_token, generated); + return false; + } + } else if (process_logits) { next_token = sample_logits(logits.data(), w_.n_vocab, sampler_, history, sampler_rng_); } else { @@ -4574,9 +4922,11 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, if (timing) tel_acc.sample_us += elapsed_us(sample_t0, Clock::now()); // Account only output tokens, before recording their emitted values - // in penalty history. The controller also covers the first seed. - if (io.is_cancelled()) break; - next_token = budget.apply(next_token).token; + // in penalty history. The controller also covers the first seed. In a + // cluster the head alone checks cancel (it travels in the decision + // flags) and applies the budget: a worker's token already carries both. + if (!cluster && io.is_cancelled()) break; + if (!worker) next_token = budget.apply(next_token).token; if (process_logits) { history.push_back(next_token); @@ -4593,7 +4943,33 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, io.emit(next_token); if (timing) tel_acc.emit_us += elapsed_us(emit_t0, Clock::now()); - if (deepseek4_is_eos_tok(next_token, w_)) { + const bool eos = deepseek4_is_eos_tok(next_token, w_); + if (cluster && !worker) { + // Head: publish the decision. Cancel is only ever evaluated here + // (workers have no client) and travels in the flags. + const bool cancel = io.is_cancelled(); + const bool is_final = eos || cancel || generated + 1 >= n_gen; + decision_flags = cluster::make_decision_flags( + eos, /*stop=*/false, cancel, budget.forced_close(), is_final); + std::string err; + if (!hooks_->decide_next_token(cluster_req, (uint32_t) generated, + logits.data(), w_.n_vocab, sampler_, + /*is_eos_candidate_unused=*/eos, + next_token, decision_flags, &err)) { + std::fprintf(stderr, "[deepseek4] cluster decision broadcast failed at step %d: %s\n", + generated, err.c_str()); + return false; + } + } + if (cluster) { + // Every rank leaves the loop on the same step for the same reason. + if (decision_flags & (cluster::kDecisionEos | cluster::kDecisionStop | + cluster::kDecisionCancel)) { + break; + } + continue; + } + if (eos) { break; } } @@ -4612,6 +4988,8 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, GenerateResult DeepSeek4Backend::generate_from_state( const GenerateRequest & req, const DaemonIO & io, int kv_offset) { GenerateResult result; + if (cluster_) cache_.cluster_rt = cluster_.get(); + const bool cluster = cluster_active(); DaemonIO out_io = io.with_token_callback(req.on_token); auto t0 = Clock::now(); sampler_ = req.sampler; @@ -4717,9 +5095,10 @@ GenerateResult DeepSeek4Backend::generate_from_state( // once from the request's starting state with every disposable arena // retired and the pipeline off (one band in flight, the same bands, so // the same numerics), instead of failing the request. Other prefills - // fail as before. + // fail as before. A cluster rank does not retry on its own: the ranks + // run every prefill chunk together. if (committed < 0 && moe_hybrid_ && deepseek4_prefill_pipeline_bands() > 0 && - !cache_.pipeline_off && !out_io.is_cancelled() && !images && + !cache_.pipeline_off && !out_io.is_cancelled() && !images && !cluster_ && (kv_offset == 0 || prefill_retry_slot_ >= 0)) { deepseek4_release_retry_scratch(cache_, moe_hybrid_.get()); cache_.pipeline_off = true; @@ -4738,11 +5117,16 @@ GenerateResult DeepSeek4Backend::generate_from_state( if (draft_swap_.out) { // Keep any sticky reduction of the swapped-out chunk, then bring the // drafter back before decode. A failed swap-in decodes this request - // without the drafter; the next request tries again. + // without the drafter and the next request tries again; a cluster rank + // fails the request instead, since every rank must draft alike. draft_swap_.cap_without = hybrid_prefill_chunk_cap_; hybrid_prefill_chunk_cap_ = draft_swap_.cap_with; deepseek4_release_prefill_scratch(cache_, moe_hybrid_.get()); if (!draft_swap_in()) { + if (cluster_) { + result.fail(GenerateErrorCode::PrefillFailed, "DSpark drafter swap-in failed"); + return result; + } std::fprintf(stderr, "[deepseek4] draft swap: this request decodes without DSpark\n"); } } @@ -4821,12 +5205,40 @@ GenerateResult DeepSeek4Backend::generate_from_state( float mv = last_logits_[0]; for (int i = 1; i < w_.n_vocab; i++) if (last_logits_[i] > mv) { mv = last_logits_[i]; seed = i; } } - if (out_io.is_cancelled()) { + if (!cluster && out_io.is_cancelled()) { result.succeed(); maybe_save_routing_stats(); return result; } - seed = budget.apply(seed).token; + if (!cluster_worker()) seed = budget.apply(seed).token; + uint8_t seed_flags = 0; + if (cluster) { + // The seed is this request's first generated token, and every rank + // must emit the same one: rank 0 decides and the workers adopt its + // token. The flags also carry the head's cancel state, because a + // worker has no client to ask. + const bool seed_eos = deepseek4_is_eos_tok(seed, w_); + int32_t seed_tok = seed; + if (!cluster_worker()) { + seed_flags = cluster::make_decision_flags( + seed_eos, /*stop=*/false, out_io.is_cancelled(), budget.forced_close(), + /*is_final=*/seed_eos || req.n_gen <= 1); + } + std::string err; + if (!hooks_->decide_next_token(hooks_->current_request(), /*step=*/0, + last_logits_.data(), w_.n_vocab, sampler_, + seed_eos, seed_tok, seed_flags, &err)) { + result.fail(GenerateErrorCode::DecodeFailed, + "cluster: DSpark seed exchange failed: " + err); + return result; + } + if (seed_tok < 0 || seed_tok >= w_.n_vocab) { + result.fail(GenerateErrorCode::DecodeFailed, + "cluster: DSpark seed token out of range"); + return result; + } + seed = seed_tok; + } if (spec_sampler) spec_sampler->history.push_back(seed); if (env_flag_enabled("LUCE_DS4_TIMING")) { size_t nonfinite_logits = 0; @@ -4844,7 +5256,12 @@ GenerateResult DeepSeek4Backend::generate_from_state( out_io.emit(seed); float accept_rate = 0.0f; bool spec_ran = false; - if (!out_io.is_cancelled() && !deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { + // Cluster: the head's seed decision says whether the request already + // ends here; a worker must not consult its own (never cancelled) IO. + const bool seed_ends_request = cluster + ? cluster::decision_terminates(seed_flags) + : (out_io.is_cancelled() || deepseek4_is_eos_tok(seed, w_)); + if (!seed_ends_request && req.n_gen > 1) { const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; const int win_len = feat_row > 0 ? (int) (spec_feat_window_.size() / feat_row) : 0; std::vector spec_toks; @@ -4867,7 +5284,11 @@ GenerateResult DeepSeek4Backend::generate_from_state( expert_runtime_.compute ? &expert_runtime_ : nullptr, routing_stats_.get(), spec_sampler ? &*spec_sampler : nullptr, - [&out_io]() { return out_io.is_cancelled(); })) { + // A cluster head stops through its draft broadcast, never + // on a local check the workers cannot see. + cluster ? std::function{} + : std::function([&out_io]() { return out_io.is_cancelled(); }), + cluster ? hooks_ : nullptr)) { result.fail(GenerateErrorCode::DecodeFailed, "DSpark speculative decode failed"); return result; @@ -5373,6 +5794,9 @@ void DeepSeek4Backend::maybe_save_routing_stats() { } void DeepSeek4Backend::shutdown() { + cache_.cluster_rt = nullptr; + hooks_ = nullptr; + if (cluster_) cluster_->free_scratch(); // device scratch belongs to backend_ release_vision(); maybe_save_routing_stats(); free_drafter(); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index e042ae786..f4a4a3b23 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -7,6 +7,7 @@ #pragma once #include "common/model_backend.h" +#include "common/cluster_participant.h" #include "common/moe_expert_compute.h" #include "common/sampler.h" #include "../common/moe_hybrid_expert_cache.h" @@ -39,6 +40,7 @@ namespace luce::common { class DeepSeek4ImagePrompt; +struct Ds4ClusterRuntime; // deepseek4_cluster.h // Bounds the sparse heterogeneous prefill arena once accumulated attention // context dominates its memory footprint. Decode batching is unaffected. @@ -73,7 +75,7 @@ bool deepseek4_mix_mmq_prefill_default( ggml_mixed_mmq_policy deepseek4_mix_mmq_prefill_policy( PrefillAttentionMode mode, const char * gcn_arch, const char * explicit_value); -class DeepSeek4Backend : public ModelBackend { +class DeepSeek4Backend : public ModelBackend, public IClusterParticipant { public: explicit DeepSeek4Backend(DeepSeek4BackendConfig cfg); ~DeepSeek4Backend() override; @@ -133,7 +135,34 @@ class DeepSeek4Backend : public ModelBackend { return routing_stats_.get(); } + // ── Expert-parallel cluster ── + // set_cluster(&cfg, nullptr) before init() builds this rank's shard + // (init() does it from cfg_.cluster); set_cluster(&cfg, comm) once RCCL + // is up only attaches the communicator. + bool set_cluster(const cluster::ClusterConfig * cfg, cluster::IClusterComm * comm); + bool cluster_attach(const cluster::ClusterConfig * cfg, + cluster::IClusterComm * comm) override { + return set_cluster(cfg, comm); + } + uint64_t cluster_placement_hash() const override; + void cluster_set_hooks(cluster::Ds4ClusterHooks * hooks) override { hooks_ = hooks; } + bool cluster_spec_decode_ready() const override { + return spec_enabled_ && spec_drafter_ != nullptr; + } + uint64_t cluster_resident_expert_bytes() const override; + bool cluster_ingraph_allreduce() const override; + private: + // Rank-0-decides hooks; nullptr outside a cluster. + cluster::Ds4ClusterHooks * hooks_ = nullptr; + bool cluster_active() const; + bool cluster_worker() const; + // Load-time agreements every rank must share: the prefill chunk size and + // the drafter (none, whole or split). + bool cluster_agree_prefill_chunk(); + bool cluster_agree_drafter(); + std::unique_ptr cluster_; + DeepSeek4BackendConfig cfg_; ggml_backend_t backend_ = nullptr; ggml_backend_t snap_backend_ = nullptr; diff --git a/server/src/deepseek4/deepseek4_cluster.cpp b/server/src/deepseek4/deepseek4_cluster.cpp new file mode 100644 index 000000000..3c1b89697 --- /dev/null +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -0,0 +1,806 @@ +// deepseek4_cluster.cpp - DeepSeek4 expert-parallel cluster runtime. +// +// See deepseek4_cluster.h for the model. This file has no RCCL dependency: +// collectives go through IClusterComm, the stream comes from +// ggml_backend_cuda_get_stream. + +#include "deepseek4_cluster.h" +#include "deepseek4_internal.h" + +#include "common/moe_hybrid_routing_stats.h" + +#include "ggml-alloc.h" +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include + +namespace luce::common { + +namespace { + +using ClusterClock = std::chrono::steady_clock; + +uint64_t us_since(ClusterClock::time_point t0) { + return (uint64_t) std::chrono::duration_cast( + ClusterClock::now() - t0).count(); +} + +// Slots cached from ClusterExpertPlacement::slot_validity. Covers decode +// (1) and every DSpark verify width (<= 5, q5 opt-in) plus headroom; larger +// prefill batches use placement.slot_owner() directly. +constexpr int kSlotValidityCacheSlots = 8; + +void set_err(std::string * err, const std::string & msg) { + if (err) *err = msg; +} + +// Mixed ROCmFP qtypes keep their codebooks out of band, so the shard has to +// register a decode table per resident tensor (ds4_cluster_init_experts). +// Reports the first mixed expert tensor found. +bool ds4_cluster_has_mix_experts(const DeepSeek4Weights & w, std::string * what) { + for (const auto & L : w.layers) { + const struct { ggml_tensor * t; int qtype; const char * name; } mix[] = { + { 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" }, + { L.ffn_down_exps, GGML_TYPE_Q2_1_ROCMFP2_MIX, "qtype-106 (mixed ROCmFP2) down" }, + }; + for (const auto & m : mix) { + if (m.t && m.t->type == m.qtype) { + if (what) *what = m.name; + return true; + } + } + } + return false; +} + +const char * placement_source_label(cluster::PlacementSource s) { + return cluster::placement_source_name(s); +} + +// Enqueue the collective for n floats at dev_ptr on the backend stream. +// force_wait: block the host until the reduction completed (host-resident +// consumers need this); otherwise stay stream-ordered unless rt.trace. +bool allreduce_device(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + void * dev_ptr, + size_t n, + int layer, + bool force_wait, + DeepSeek4StepTelemetry * telemetry, + std::string * err) { + if (!rt.cfg) { + set_err(err, "cluster runtime has no config"); + return false; + } + // Before the communicator is attached (load-time forwards) a rank runs + // its own shard alone; nothing reads those outputs. + if (!rt.comm || rt.comm->size() <= 1 || n == 0) { + return true; // single rank: the partial already is the sum + } + if (!dev_ptr) { + set_err(err, "cluster all-reduce: null device pointer"); + return false; + } + + // ggml_backend_cuda_get_stream returns the stream ggml_backend_graph_compute + // enqueues on (cuda_ctx->stream()), so the collective is ordered after the + // producing kernels. + cluster::DeviceStream stream = ggml_backend_cuda_get_stream(backend); + + const uint64_t bytes_f32 = (uint64_t) n * sizeof(float); + const cluster::AllreduceDType dtype = rt.cfg->allreduce_dtype; + // `auto` resolves to F32 for every size. The bf16 + // compressed collective rounds each per-layer partial to 8 mantissa + // bits; over 43 layers of a multi-token prefill (>256 KiB payload, which + // `auto` would have sent as bf16) that visibly changes the first + // generated token. bf16 stays available as an explicit + // --cluster-allreduce-dtype bf16. `auto` can use the bf16 threshold once + // a bf16 prefill is token-identical to f32 on the qualification prompts. + const bool use_bf16 = dtype == cluster::AllreduceDType::BF16; + if (dtype == cluster::AllreduceDType::Auto && bytes_f32 > cluster::kAutoBf16ThresholdBytes) { + static bool logged_auto_f32 = false; + if (!logged_auto_f32) { + logged_auto_f32 = true; + std::fprintf(stderr, + "[deepseek4-cluster] allreduce dtype auto: using f32 for %llu-byte " + "partials (bf16 auto-switch disabled until verified)\n", + (unsigned long long) bytes_f32); + } + } + + bool ok = false; + uint64_t payload = 0; + if (use_bf16) { + if (!rt.ensure_scratch(backend, n, err)) return false; + ok = rt.comm->allreduce_sum_bf16_compressed(dev_ptr, n, rt.scratch_bf16->data, stream, err); + payload = (uint64_t) n * 2u; + } else { + ok = rt.comm->allreduce_sum_f32(dev_ptr, n, stream, err); + payload = bytes_f32; + } + if (!ok) { + if (err && err->empty()) *err = "cluster all-reduce enqueue failed"; + return false; + } + rt.telemetry.allreduce_calls += 1; + rt.telemetry.allreduce_bytes += payload; + if (telemetry) telemetry->cluster_allreduce_bytes += payload; + + if (force_wait || rt.trace) { + const auto wait_t0 = ClusterClock::now(); + if (!rt.comm->wait_stream(stream, rt.cfg->timeout_ms, err)) { + if (err && err->empty()) *err = "cluster all-reduce wait timed out"; + return false; + } + const uint64_t wait_us = us_since(wait_t0); + rt.telemetry.allreduce_wait_us += wait_us; + if (rt.trace) { + std::fprintf(stderr, + "[deepseek4-cluster] rank %d layer %d allreduce n=%zu %s wait=%llu us\n", + rt.rank(), layer, n, use_bf16 ? "bf16" : "f32", + (unsigned long long) wait_us); + } + } + return true; +} + +} // namespace + +// ─── Ds4ClusterRuntime ───────────────────────────────────────────────── + +Ds4ClusterRuntime::~Ds4ClusterRuntime() { + free_scratch(); +} + +bool Ds4ClusterRuntime::evaluates(int layer, int expert, int slot) const { + if (slot >= 0 && slot < slot_validity_n_slots && !slot_validity_cache.empty() && + layer >= 0 && layer < placement.n_layer && expert >= 0 && expert < placement.n_expert) { + const size_t idx = + ((size_t) layer * (size_t) slot_validity_n_slots + (size_t) slot) * + (size_t) placement.n_expert + (size_t) expert; + if (idx < slot_validity_cache.size()) return slot_validity_cache[idx] != 0; + } + return placement.slot_owner(layer, expert, slot) == rank(); +} + +bool Ds4ClusterRuntime::ensure_scratch(ggml_backend_t backend, size_t n_elems, std::string * err) { + if (scratch_buf && scratch_backend == backend && scratch_elems >= n_elems) { + return true; + } + free_scratch(); + ggml_init_params ip{}; + ip.mem_size = 4 * ggml_tensor_overhead(); + ip.mem_buffer = nullptr; + ip.no_alloc = true; + scratch_ctx = ggml_init(ip); + if (!scratch_ctx) { + set_err(err, "cluster scratch: ggml_init failed"); + return false; + } + scratch_f32 = ggml_new_tensor_1d(scratch_ctx, GGML_TYPE_F32, (int64_t) n_elems); + ggml_set_name(scratch_f32, "ds4_cluster_scratch_f32"); + // bf16 payload for the compressed collective: 2 bytes per element. + scratch_bf16 = ggml_new_tensor_1d(scratch_ctx, GGML_TYPE_I16, (int64_t) n_elems); + ggml_set_name(scratch_bf16, "ds4_cluster_scratch_bf16"); + scratch_buf = ggml_backend_alloc_ctx_tensors(scratch_ctx, backend); + if (!scratch_buf) { + free_scratch(); + set_err(err, "cluster scratch: device allocation failed"); + return false; + } + scratch_backend = backend; + scratch_elems = n_elems; + return true; +} + +void Ds4ClusterRuntime::free_scratch() { + if (scratch_buf) { ggml_backend_buffer_free(scratch_buf); scratch_buf = nullptr; } + if (scratch_ctx) { ggml_free(scratch_ctx); scratch_ctx = nullptr; } + scratch_f32 = nullptr; + scratch_bf16 = nullptr; + scratch_elems = 0; + scratch_backend = nullptr; +} + +// ─── Config / placement ──────────────────────────────────────────────── + +MoeHybridConfig ds4_cluster_moe_config(const DeepSeek4Weights & w) { + MoeHybridConfig cfg; + cfg.n_embd = w.n_embd; + cfg.n_expert = w.n_expert; + cfg.n_expert_used = w.n_expert_used; + cfg.n_ff_exp = w.n_ff_exp; + cfg.n_ff_shexp = w.n_ff_exp; + cfg.n_layer = w.n_layer; + cfg.first_moe_layer = 0; + cfg.swiglu_clamp = w.swiglu_clamp_exp; + cfg.cold_expert_backend = MoeHybridColdBackend::None; + cfg.materialize_hot_experts = true; + cfg.materialize_cold_experts = false; + // Same as the single-node DS4 hybrid config: reduced hot stacks stay on + // the q<=4 MMVQ sub-batch path on gfx1151 (mmq_safe_full_batch=false). + return cfg; +} + +bool ds4_cluster_build_placement(const cluster::ClusterConfig & cfg, + const DeepSeek4Weights & w, + const std::string * hotness_csv_path, + Ds4ClusterRuntime & rt, + std::string * err) { + if (!cfg.enabled() || cfg.rank < 0 || cfg.rank >= cfg.size) { + set_err(err, "cluster config is not enabled or rank is out of range"); + return false; + } + + const MoeHybridConfig mcfg = ds4_cluster_moe_config(w); + MoeHybridRoutingStats stats; + bool have_stats = false; + const bool need_stats = + cfg.placement_source == cluster::PlacementSource::Balanced || + (cfg.placement_source == cluster::PlacementSource::Uniform && cfg.replicate_hot > 0); + if (need_stats) { + if (!hotness_csv_path || hotness_csv_path->empty()) { + set_err(err, std::string(placement_source_label(cfg.placement_source)) + + " placement with replicate_hot/balanced requires LUCE_DS4_HOTNESS_CSV"); + return false; + } + if (!MoeHybridRoutingStats::load_csv(*hotness_csv_path, stats, err)) return false; + if (!stats.matches(mcfg)) { + set_err(err, "routing hotness CSV shape does not match the DeepSeek V4 target " + "(n_layer/n_expert/n_expert_used)"); + return false; + } + have_stats = true; + } + + switch (cfg.placement_source) { + case cluster::PlacementSource::Uniform: + if (!cluster::ClusterExpertPlacement::build_uniform( + cfg.size, mcfg, cfg.replicate_hot, have_stats ? &stats : nullptr, + rt.placement, err)) { + return false; + } + break; + case cluster::PlacementSource::Balanced: + if (!cluster::ClusterExpertPlacement::build_balanced( + cfg.size, stats, cfg.replicate_hot, rt.placement, err)) { + return false; + } + break; + case cluster::PlacementSource::File: + if (!cluster::ClusterExpertPlacement::load_json(cfg.placement_file, rt.placement, err)) { + return false; + } + if (rt.placement.n_ranks != cfg.size) { + set_err(err, "placement file was built for " + std::to_string(rt.placement.n_ranks) + + " ranks, cluster size is " + std::to_string(cfg.size)); + return false; + } + if (!rt.placement.matches(mcfg)) { + set_err(err, "placement file dimensions do not match the DeepSeek V4 target"); + return false; + } + break; + } + + if (!rt.placement.to_rank_placement(cfg.rank, rt.rank_placement, err)) return false; + rt.placement.slot_validity(cfg.rank, kSlotValidityCacheSlots, rt.slot_validity_cache); + rt.slot_validity_n_slots = kSlotValidityCacheSlots; + if (&cfg != &rt.cfg_storage) rt.cfg_storage = cfg; + rt.cfg = &rt.cfg_storage; + rt.trace = cluster::cluster_env_trace(); + + // Shared-expert sharding: this rank's slice of the intermediate axis. The + // down projection contracts over that axis, so the slice is a partial sum + // and has to be added to the routed partial BEFORE the all-reduce. + rt.shexp_ff_begin = 0; + rt.shexp_ff_count = 0; + if (cfg.shared_expert == cluster::SharedExpertMode::Shard && cfg.size > 1) { + if (w.n_ff_exp > 0 && w.n_ff_exp % cfg.size == 0) { + rt.shexp_ff_count = w.n_ff_exp / cfg.size; + rt.shexp_ff_begin = cfg.rank * rt.shexp_ff_count; + if (w.cluster_shexp_ff_local > 0) { + if (w.cluster_shexp_ff_local != rt.shexp_ff_count) { + set_err(err, "loaded shared-expert slice does not match the shard"); + return false; + } + rt.shexp_ff_begin = 0; // the loaded tensors are this slice + } + std::fprintf(stderr, + "[deepseek4-cluster] rank %d/%d shared expert intermediate %d..%d " + "of %d; its partial rides in the routed all-reduce\n", + cfg.rank, cfg.size, rt.shexp_ff_begin, + rt.shexp_ff_begin + rt.shexp_ff_count - 1, w.n_ff_exp); + } else { + std::fprintf(stderr, + "[deepseek4-cluster] shared expert stays replicated: %d " + "intermediate units do not split %d ways\n", + w.n_ff_exp, cfg.size); + } + } + + // Attention head parallelism. The split has to be in whole output groups, + // because the grouped output projection's first stage is indexed by group; + // a rank count that does not divide the groups (or the heads) keeps + // attention replicated instead of splitting it unevenly. + rt.attn_head_begin = 0; + rt.attn_head_count = 0; + if (ds4_cluster_attention_parallel_enabled() && cfg.size > 1 && + w.n_out_group > 0 && w.n_head > 0 && + w.n_out_group % cfg.size == 0 && w.n_head % cfg.size == 0) { + rt.attn_head_count = w.n_head / cfg.size; + rt.attn_head_begin = cfg.rank * rt.attn_head_count; + if (w.cluster_heads_local > 0) { + if (w.cluster_heads_local != rt.attn_head_count) { + set_err(err, "loaded attention slice does not match the head split"); + return false; + } + rt.attn_head_begin = 0; // the loaded tensors are this slice + } + std::fprintf(stderr, + "[deepseek4-cluster] rank %d/%d attention heads %d..%d of %d " + "(%d of %d output groups); the attention output is a partial sum " + "and is all-reduced with the routed experts\n", + cfg.rank, cfg.size, rt.attn_head_begin, + rt.attn_head_begin + rt.attn_head_count - 1, w.n_head, + w.n_out_group / cfg.size, w.n_out_group); + } else if (cfg.size > 1) { + std::fprintf(stderr, + "[deepseek4-cluster] rank %d/%d attention stays replicated " + "(%d heads / %d output groups do not split %d ways, or the " + "kill-switch is set)\n", + cfg.rank, cfg.size, w.n_head, w.n_out_group, cfg.size); + } + + std::fprintf(stderr, "[deepseek4-cluster] rank %d/%d %s source=%s resident_experts=%d/%d\n", + cfg.rank, cfg.size, + rt.placement.describe(have_stats ? &stats : nullptr).c_str(), + placement_source_label(cfg.placement_source), + rt.rank_placement.total_hot, w.n_layer * w.n_expert); + return true; +} + +bool ds4_cluster_init_experts(const std::string & model_path, + ggml_backend_t backend, + const DeepSeek4Weights & w, + Ds4ClusterRuntime & rt, + MoeHybridStorage & storage, + std::string * err) { + std::string mix_what; + const bool has_mix = ds4_cluster_has_mix_experts(w, &mix_what); + const MoeHybridConfig mcfg = ds4_cluster_moe_config(w); + if (!rt.rank_placement.matches(mcfg)) { + set_err(err, "rank placement does not match the loaded model dimensions"); + return false; + } + + // The non-mmap builder materializes exactly the hot (resident) slices of + // every expert tensor from the file and unmaps afterwards; with cold owner + // None nothing else is read, so no streaming engine or retained mmap is + // needed. (build_deepseek4_moe_hybrid_storage_from_file_with_mmap would + // keep the whole file mapped for cold-expert streaming.) + if (!build_deepseek4_moe_hybrid_storage_from_file( + model_path, backend, w, rt.rank_placement, &mcfg, storage, err)) { + return false; + } + if (storage.cold_backend_kind != MoeHybridColdBackend::None || + storage.materialized_cold_experts) { + set_err(err, "cluster expert storage did not come back with cold owner None"); + return false; + } + + uint64_t bytes = 0; + int resident = 0; + int layers_with_experts = 0; + for (const MoeHybridLayerStorage & layer : storage.layers) { + if (!layer.hot_buf) continue; + bytes += (uint64_t) ggml_backend_buffer_get_size(layer.hot_buf); + resident += (int) layer.hot_expert_ids.size(); + ++layers_with_experts; + if (!layer.cold_expert_ids.empty() || layer.cold_buf) { + set_err(err, "cluster expert storage unexpectedly holds cold experts"); + return false; + } + } + // Adaptive (mixed ROCmFPX) experts keep their codebooks out of band and + // need a decode table registered per resident tensor. The shard's resident + // set is exactly its hot experts, and the registrar treats an absent cold + // owner as "nothing to register", so hot-only storage is enough. + if (has_mix) { + std::string why; + if (!register_deepseek4_moe_hybrid_mix_tables(model_path, w, storage, &why)) { + set_err(err, "cluster shard could not register decode tables for " + + mix_what + " experts" + (why.empty() ? "" : ": " + why)); + return false; + } + } + if (has_mix) { + std::fprintf(stderr, + "[deepseek4-cluster] rank %d/%d registered %s decode tables for " + "its resident experts\n", rt.rank(), rt.size(), mix_what.c_str()); + } + + rt.resident_expert_bytes = bytes; + const int total_experts = w.n_layer * w.n_expert; + std::fprintf(stderr, + "[deepseek4-cluster] rank %d/%d resident routed experts: %d/%d (%.1f%%) " + "%.2f GiB (%.1f MiB/layer over %d layers), owned=%d replicated=%d\n", + rt.rank(), rt.size(), resident, total_experts, + total_experts > 0 ? 100.0 * (double) resident / (double) total_experts : 0.0, + (double) bytes / (1024.0 * 1024.0 * 1024.0), + layers_with_experts > 0 + ? (double) bytes / (1024.0 * 1024.0) / (double) layers_with_experts : 0.0, + layers_with_experts, + rt.placement.total_owned(rt.rank()), + resident - rt.placement.total_owned(rt.rank())); + return true; +} + +bool ds4_cluster_ingraph_allreduce_enabled() { + static const bool enabled = [] { + const char * v = std::getenv("LUCE_CLUSTER_NO_INGRAPH_ALLREDUCE"); + const bool off = v && v[0] && std::strcmp(v, "0") != 0; + if (off) { + std::fprintf(stderr, + "[deepseek4-cluster] LUCE_CLUSTER_NO_INGRAPH_ALLREDUCE=1: " + "per-layer host all-reduce, no fused graph\n"); + } + return !off; + }(); + return enabled; +} + +bool ds4_cluster_prefill_pipeline_enabled(const std::vector & bands) { + if (bands.empty()) return false; + static const int forced = [] { // -1 unset, 0 off, 1 on + const char * v = std::getenv("LUCE_CLUSTER_PREFILL_PIPELINE"); + return v && *v ? (std::strcmp(v, "0") != 0 ? 1 : 0) : -1; + }(); + if (forced >= 0) return forced == 1; + return std::all_of(bands.begin(), bands.end(), + [](int rows) { return rows >= kDs4ClusterPipelineMinBand; }); +} + +bool ds4_cluster_attention_parallel_enabled() { + // Off unless set; cluster mode sets it (cluster_launch_env), and =0 keeps + // attention replicated. + static const bool enabled = [] { + const char * v = std::getenv("LUCE_CLUSTER_ATTENTION_PARALLEL"); + return v && v[0] && std::strcmp(v, "0") != 0; + }(); + return enabled; +} + +bool ds4_cluster_fused_graph_available(const Ds4ClusterRuntime * rt) { + if (!rt || !rt->comm || rt->comm->size() <= 1) return false; + if (!ds4_cluster_ingraph_allreduce_enabled()) return false; + for (int32_t owner : rt->placement.owner) { + if (owner == cluster::kReplicated) { + static bool warned = false; + if (!warned) { + warned = true; + std::fprintf(stderr, + "[deepseek4-cluster] fused graph unavailable: a " + "replicated expert is owned per token slot, which the " + "fused graph's per-expert owner table cannot express; " + "using the per-layer host all-reduce\n"); + } + return false; + } + } + static bool logged = false; + if (!logged) { + logged = true; + std::fprintf(stderr, + "[deepseek4-cluster] fused whole-model graph with an " + "in-graph all-reduce per MoE layer\n"); + } + return true; +} + +// ─── Route masking ────────────────────────────────────────────────────── + +void ds4_cluster_mask_routes(const Ds4ClusterRuntime & rt, + int layer, + int32_t * selected, + float * weights, + int route_width, + int n_tokens) { + if (!selected || !weights || route_width <= 0 || n_tokens <= 0) return; + for (int t = 0; t < n_tokens; ++t) { + int32_t * ids = selected + (size_t) t * (size_t) route_width; + float * wts = weights + (size_t) t * (size_t) route_width; + for (int j = 0; j < route_width; ++j) { + const int32_t e = ids[j]; + if (e < 0) continue; + if (!rt.evaluates(layer, (int) e, t)) { + ids[j] = -1; + wts[j] = 0.0f; + } + } + } +} + +// ─── All-reduce ───────────────────────────────────────────────────────── + +// Partials at least this large are prefill-sized and may take the compressed +// RCCL paths below; decode-sized partials stay far below it. +constexpr size_t kPrefillPartialMinElems = 262144; + +// Invoked by GGML_MOE_FUSED_CLUSTER_ALLREDUCE while the graph runs. +// It enqueues the collective on the stream ggml handed us, which is the same +// stream the surrounding kernels use, so ordering needs no host synchronize. +static void ds4_cluster_allreduce_graph_callback(void * user, void * data, + size_t n, void * stream) { + Ds4ClusterRuntime * rt = static_cast(user); + if (!rt || !rt->comm || rt->comm->size() <= 1 || n == 0 || !data) return; + // LUCE_CLUSTER_PREFILL_BF16=1: prefill-sized partials cross the link as + // bf16 (RCCL sums in bf16; every rank gets the same result). Decode-sized + // partials never reach this branch (they stay on fast_reduce). + static const bool prefill_bf16 = [] { + const char * v = std::getenv("LUCE_CLUSTER_PREFILL_BF16"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + // LUCE_CLUSTER_PREFILL_Q8=1: prefill-sized partials cross the link as + // 8-bit blocks (32 values, one scale), about 0.56x the bytes of bf16; every + // rank sums the same dequantized values in rank order. The f32 scratch + // holds both ranks' blocks. + static const bool prefill_q8 = [] { + const char * v = std::getenv("LUCE_CLUSTER_PREFILL_Q8"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + if (prefill_q8 && n >= kPrefillPartialMinElems && n % cluster::kQ8ExchangeBlock == 0 && rt->scratch_f32 && + rt->scratch_elems * sizeof(float) >= (size_t) rt->comm->size() * cluster::q8_exchange_slot_bytes(n)) { + std::string errq; + if (!rt->comm->allreduce_sum_q8_exchange(data, n, rt->scratch_f32->data, + (cluster::DeviceStream) stream, &errq)) { + if (rt->node_error.empty()) rt->node_error = errq.empty() ? "in-graph q8 exchange failed" : errq; + return; + } + rt->telemetry.allreduce_calls += 1; + rt->telemetry.allreduce_bytes += (uint64_t) cluster::q8_exchange_slot_bytes(n); + return; + } + if (prefill_bf16 && n >= kPrefillPartialMinElems && rt->scratch_bf16 && rt->scratch_elems >= n) { + std::string err16; + if (!rt->comm->allreduce_sum_bf16_compressed(data, n, rt->scratch_bf16->data, + (cluster::DeviceStream) stream, &err16)) { + if (rt->node_error.empty()) rt->node_error = err16.empty() ? "in-graph bf16 all-reduce failed" : err16; + return; + } + rt->telemetry.allreduce_calls += 1; + rt->telemetry.allreduce_bytes += (uint64_t) n * 2u; + return; + } + if (rt->fast && rt->fast->ok() && rt->fast->submit((float *) data, n, stream)) { + rt->telemetry.allreduce_calls += 1; + rt->telemetry.allreduce_bytes += (uint64_t) n * sizeof(float); + return; + } + std::string err; + if (!rt->comm->allreduce_sum_f32(data, n, (cluster::DeviceStream) stream, &err)) { + // Keep the first failure: the rest of the graph still runs, and the + // caller turns this into a failed forward once the compute returns. + if (rt->node_error.empty()) { + rt->node_error = err.empty() ? "in-graph all-reduce failed" : err; + } + return; + } + rt->telemetry.allreduce_calls += 1; + rt->telemetry.allreduce_bytes += (uint64_t) n * sizeof(float); +} + +// Hybrid exchange callbacks. Each kind counts its own calls; the k-th call +// of every kind belongs to exchange k, so the host order of the dGPU and +// iGPU splits does not matter. A failure is fatal for the step: the peer is +// already waiting on this rank's row. +static void ds4_hyb_fail(Ds4ClusterRuntime * rt, const char * what) { + if (rt->node_error.empty()) rt->node_error = std::string("hybrid exchange: ") + what + " failed"; +} +static void ds4_hyb_export_cb(void * user, void * data, size_t n, void * stream) { + auto * rt = static_cast(user); + if (!rt || !rt->fast || !rt->fast->hyb_export((const float *) data, n, stream, rt->hyb_export++)) ds4_hyb_fail(rt, "export"); +} +static void ds4_hyb_combine_cb(void * user, void * data, size_t n, void * stream) { + auto * rt = static_cast(user); + if (!rt || !rt->fast || !rt->fast->hyb_combine((const float *) data, n, stream, rt->hyb_combine++)) ds4_hyb_fail(rt, "combine"); + else { rt->telemetry.allreduce_calls += 1; rt->telemetry.allreduce_bytes += (uint64_t) n * sizeof(float); } +} +static void ds4_hyb_import_cb(void * user, void * data, size_t n, void * stream) { + auto * rt = static_cast(user); + if (!rt || !rt->fast || !rt->fast->hyb_import((float *) data, n, stream, rt->hyb_import++)) ds4_hyb_fail(rt, "import"); +} + +bool ds4_cluster_hybrid_exchange_enabled(const Ds4ClusterRuntime & rt) { + static const bool env = [] { + const char * v = std::getenv("LUCE_CLUSTER_HYBRID_EXCHANGE"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + return env && rt.comm && rt.comm->size() == 2 && rt.fast && rt.fast->hyb_available(); +} + +ggml_tensor * ds4_cluster_hybrid_exchange_nodes(ggml_context * ctx, + ggml_tensor * hot, + ggml_tensor * shared, + ggml_tensor * cold, + Ds4ClusterRuntime & rt, + ggml_tensor ** combine_node, + ggml_tensor ** export_node) { + if (!ctx || !hot || !shared || !cold || !combine_node || !export_node) return nullptr; + for (ggml_tensor * t : {hot, shared, cold}) { + if (t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || ggml_nelements(t) != ggml_nelements(hot)) return nullptr; + } + // In place throughout: each node passes its own tensor to its callback. + // The dGPU exports hot + shared as one partial (one crossing of its link), + // the iGPU adds cold + (hot + shared) and exchanges, and the import + // overwrites the exported partial with the reduced sum. + ggml_tensor * ex = ggml_cluster_allreduce_inplace(ctx, ggml_add(ctx, hot, shared), &ds4_hyb_export_cb, &rt); + ggml_tensor * comb = ggml_cluster_allreduce_inplace(ctx, cold, &ds4_hyb_combine_cb, &rt); + ggml_tensor * imp = ggml_cluster_allreduce_inplace(ctx, ex, &ds4_hyb_import_cb, &rt); + *combine_node = comb; + *export_node = ex; + return imp; +} + +ggml_tensor * ds4_cluster_allreduce_node(ggml_context * ctx, + ggml_tensor * partial, + Ds4ClusterRuntime & rt, + bool inplace) { + if (!ctx || !partial) return nullptr; + if (!rt.comm || rt.comm->size() <= 1) return partial; + if (partial->type != GGML_TYPE_F32 || !ggml_is_contiguous(partial)) { + std::fprintf(stderr, + "[deepseek4-cluster] in-graph all-reduce needs a contiguous " + "F32 partial\n"); + return nullptr; + } + return inplace + ? ggml_cluster_allreduce_inplace(ctx, partial, + &ds4_cluster_allreduce_graph_callback, &rt) + : ggml_cluster_allreduce(ctx, partial, + &ds4_cluster_allreduce_graph_callback, &rt); +} + +bool ds4_cluster_allreduce_layer_host(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + float * partial_host, + int n_embd, + int n_tokens, + int layer, + DeepSeek4StepTelemetry * telemetry, + std::string * err) { + if (!rt.comm || rt.comm->size() <= 1) return true; + if (!partial_host || n_embd <= 0 || n_tokens <= 0) { + set_err(err, "cluster all-reduce: invalid host partial"); + return false; + } + const auto t0 = ClusterClock::now(); + const size_t n = (size_t) n_embd * (size_t) n_tokens; + if (!rt.ensure_scratch(backend, n, err)) return false; + // ggml_backend_tensor_set completes the H2D copy before returning + // (ggml-cuda copies on cudaStreamPerThread and synchronizes it), so the + // collective enqueued next on the backend stream reads final data. + ggml_backend_tensor_set(rt.scratch_f32, partial_host, 0, n * sizeof(float)); + if (!allreduce_device(rt, backend, rt.scratch_f32->data, n, layer, + /*force_wait=*/true, telemetry, err)) { + return false; + } + ggml_backend_tensor_get(rt.scratch_f32, partial_host, 0, n * sizeof(float)); + if (telemetry) telemetry->cluster_allreduce_us += us_since(t0); + return true; +} + +// ─── Lucebox additions ──────────────────────────────────────────────── + +namespace { + +// A broadcast from rank 0 as an in-graph collective. `exact` keeps rank 0's +// values bit for bit (f32 on the wire even where the fast reduce rounds its +// payloads to bf16) and skips the upload hook, which belongs to the Engram +// keys of the plain broadcast. +void broadcast_in_graph(void * user, void * data, size_t n, void * stream, bool exact) { + Ds4ClusterRuntime * rt = static_cast(user); + if (!rt || !rt->comm || rt->comm->size() <= 1 || n == 0 || !data) return; + if (!exact && rt->before_broadcast && !rt->before_broadcast(data, n, stream)) { + if (rt->node_error.empty()) rt->node_error = "upload before the in-graph broadcast failed"; + } + if (rt->fast && rt->fast->ok() && + rt->fast->submit_from_root((float *) data, n, stream, /*root=*/0, exact)) { + return; + } + std::string err; + if (!rt->comm->broadcast_bytes(data, n * sizeof(float), /*root=*/0, + (cluster::DeviceStream) stream, &err)) { + if (rt->node_error.empty()) { + rt->node_error = err.empty() ? "in-graph broadcast failed" : err; + } + } +} + +void broadcast_graph_callback(void * user, void * data, size_t n, void * stream) { + broadcast_in_graph(user, data, n, stream, /*exact=*/false); +} + +void broadcast_exact_graph_callback(void * user, void * data, size_t n, void * stream) { + broadcast_in_graph(user, data, n, stream, /*exact=*/true); +} + +// The in-graph collective node runs its callback in place on a copy of +// `value`. +ggml_tensor * broadcast_node(ggml_context * ctx, ggml_tensor * value, Ds4ClusterRuntime & rt, + ggml_cluster_allreduce_fn callback) { + if (!ctx || !value) return nullptr; + if (rt.size() <= 1) return value; + if (value->type != GGML_TYPE_F32) { + std::fprintf(stderr, "[deepseek4-cluster] broadcast needs an F32 tensor\n"); + return nullptr; + } + if (!ggml_is_contiguous(value)) value = ggml_cont(ctx, value); + return ggml_cluster_allreduce(ctx, value, callback, &rt); +} + +} // namespace + +ggml_tensor * ds4_cluster_broadcast_exact_node(ggml_context * ctx, + ggml_tensor * value, + Ds4ClusterRuntime & rt) { + return broadcast_node(ctx, value, rt, &broadcast_exact_graph_callback); +} + +ggml_tensor * ds4_cluster_broadcast_node(ggml_context * ctx, + ggml_tensor * value, + Ds4ClusterRuntime & rt) { + return broadcast_node(ctx, value, rt, &broadcast_graph_callback); +} + +bool ds4_cluster_broadcast_host(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + float * data, + size_t n, + std::string * err) { + if (!rt.comm || rt.comm->size() <= 1 || n == 0) return true; + if (!data) { + set_err(err, "cluster broadcast: null host buffer"); + return false; + } + if (!rt.ensure_scratch(backend, n, err)) return false; + cluster::DeviceStream stream = ggml_backend_cuda_get_stream(backend); + ggml_backend_tensor_set(rt.scratch_f32, data, 0, n * sizeof(float)); + if (!rt.comm->broadcast_bytes(rt.scratch_f32->data, n * sizeof(float), 0, stream, err)) { + return false; + } + if (!rt.comm->wait_stream(stream, rt.cfg ? rt.cfg->timeout_ms : 30000u, err)) return false; + ggml_backend_tensor_get(rt.scratch_f32, data, 0, n * sizeof(float)); + return true; +} + +bool ds4_cluster_allgather_i32(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + int32_t value, + std::vector & out, + std::string * err) { + const int n = std::max(1, rt.size()); + std::vector slots((size_t) n, 0.0f); + slots[(size_t) std::max(0, rt.rank())] = (float) value; // exact below 2^24 + if (rt.comm && rt.comm->size() > 1 && + !ds4_cluster_allreduce_layer_host(rt, backend, slots.data(), n, 1, -1, nullptr, err)) { + return false; + } + out.resize((size_t) n); + for (int i = 0; i < n; ++i) out[(size_t) i] = (int32_t) slots[(size_t) i]; + return true; +} + +bool ds4_cluster_resident(const Ds4ClusterRuntime & rt, int layer, int expert) { + return rt.rank_placement.is_hot(layer, expert); +} + +} // namespace luce::common diff --git a/server/src/deepseek4/deepseek4_cluster.h b/server/src/deepseek4/deepseek4_cluster.h new file mode 100644 index 000000000..f64b5f479 --- /dev/null +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -0,0 +1,259 @@ +// deepseek4_cluster.h - DeepSeek4 expert-parallel cluster runtime, after +// maikzz32's lucebox-halo-cluster. +// +// Every rank loads only the dense slices it computes (its attention heads +// and its shared-expert slice) and only its shard of routed experts (owned + replicated, see cluster_expert_placement.h). Per MoE +// layer the rank evaluates its resident experts as an owner partial sum +// through the existing hybrid machinery with cold owner +// MoeHybridColdBackend::None, then ONE all-reduce (IClusterComm:: +// allreduce_sum_f32 on the backend stream) makes the layer's output identical +// on every rank before HC-post runs. +// +// By default the all-reduce is a node of the fused whole-model verify/decode +// graph (ggml_cluster_allreduce, one per MoE layer). With the fused graph +// unavailable (see ds4_cluster_fused_graph_available) it is a host-driven +// call between per-layer graph computes inside +// deepseek4_step_layer_range / eval_ds4_layer_range_hybrid_ffn. +// +// Shared expert: with SharedExpertMode::Shard (the default) each rank owns a +// slice of its intermediate axis, so its down projection is a partial sum +// that rides in the routed all-reduce. With Replicate every rank computes it +// from the same normalized input and adds it after the all-reduce; the +// routed partial is then evaluated with the shared-expert tensors removed +// from the MoeLayerDesc (the per-layer equivalent of include_shared=false). +// +// Replicated experts: expert e at batch slot t is evaluated only by rank +// (e + t) % N. Slot = token index inside the current batch (0..n_tokens-1), +// identical on every rank because routing is replicated; for decode +// (n_tokens == 1) this pins replicated expert e to rank e % N. + +#pragma once + +#include "deepseek4_internal.h" + +#include "cluster/cluster_comm.h" +#include "cluster/fast_reduce.h" +#include "cluster/cluster_config.h" +#include "cluster/cluster_expert_placement.h" +#include "cluster/cluster_telemetry.h" +#include "common/moe_hybrid_placement.h" +#include "common/moe_hybrid_storage.h" +#include "common/moe_hybrid_types.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include +#include +#include + +namespace luce::common { + +struct Ds4ClusterRuntime { + // cfg points at cfg_storage (a copy taken by DeepSeek4Backend::set_cluster) + // so the caller's ClusterConfig need not outlive the backend. comm is + // borrowed and may be attached after the model was loaded; it is required + // by the first forward. + cluster::ClusterConfig cfg_storage; + const cluster::ClusterConfig * cfg = nullptr; + cluster::IClusterComm * comm = nullptr; + // Decode-path all-reduces over raw RDMA (LUCE_CLUSTER_FAST_REDUCE=1); RCCL keeps + // whatever it declines (prefill-sized payloads). + std::unique_ptr fast; + // Hybrid exchange call counters: the k-th export/combine/import call of + // each kind belongs to exchange k. Like FastReduce's sequence numbers they + // assume every rank runs every exchange for the life of the process: a + // step that fails on one rank leaves its peer waiting, and the stall + // watchdog then ends both ranks. + uint64_t hyb_export = 0, hyb_combine = 0, hyb_import = 0; + cluster::ClusterExpertPlacement placement; + // Rank-local hot set (owned + replicated experts) for the hybrid storage. + MoeHybridPlacement rank_placement; + // [n_layer][slot_validity_n_slots][n_expert] 0/1 from + // ClusterExpertPlacement::slot_validity for this rank; consulted for + // slots below slot_validity_n_slots, larger slots fall back to + // placement.slot_owner(). + std::vector slot_validity_cache; + int slot_validity_n_slots = 0; + cluster::ClusterStepTelemetry telemetry; + bool trace = false; + // A graph node cannot return an error, so the in-graph + // all-reduce records the first failure here and the forward fails after + // the compute. Cleared before every graph that contains such a node. + std::string node_error; + // Attention head parallelism: this rank builds only heads + // [attn_head_begin, attn_head_begin + attn_head_count) and the ranks + // all-reduce the attention output, which is a partial sum over the + // grouped output projection. count 0 = attention stays replicated. The + // range indexes the loaded tensors: 0 when they hold only this slice. + int attn_head_begin = 0; + int attn_head_count = 0; + // Shared-expert sharding (--cluster-shared-expert shard, the default): + // this rank's slice of the intermediate axis. count 0 = the shared expert is + // replicated and added once after the reduction. + int shexp_ff_begin = 0; + int shexp_ff_count = 0; + + // Device staging for host-resident partials (per-layer path): one F32 [n] and + // one bf16 [n] (stored as I16) tensor in a single backend buffer, grown + // on demand. Owned here, freed in the destructor / free_scratch(). + ggml_backend_t scratch_backend = nullptr; + ggml_context * scratch_ctx = nullptr; + ggml_backend_buffer_t scratch_buf = nullptr; + ggml_tensor * scratch_f32 = nullptr; + ggml_tensor * scratch_bf16 = nullptr; + size_t scratch_elems = 0; + + // Resident routed-expert bytes on this rank (sum of hot buffers), for logs. + uint64_t resident_expert_bytes = 0; + + // Set by a forward for the duration of one graph: called from the + // in-graph broadcast's callback with its buffer and stream before the + // broadcast is queued, so data a host thread is still producing can be + // uploaded right before its consumer (the verify's Engram keys). False + // fails the forward. + std::function before_broadcast; + + Ds4ClusterRuntime() = default; + ~Ds4ClusterRuntime(); + Ds4ClusterRuntime(const Ds4ClusterRuntime &) = delete; + Ds4ClusterRuntime & operator=(const Ds4ClusterRuntime &) = delete; + + int rank() const { return cfg ? cfg->rank : 0; } + int size() const { return cfg ? cfg->size : 1; } + + // True when this rank evaluates (layer, expert) for batch slot `slot`. + bool evaluates(int layer, int expert, int slot) const; + + bool ensure_scratch(ggml_backend_t backend, size_t n_elems, std::string * err); + void free_scratch(); +}; + +// MoeHybridConfig for cluster expert storage: cold owner None, cold experts +// never materialized, hot experts materialized on the local GPU. +MoeHybridConfig ds4_cluster_moe_config(const DeepSeek4Weights & w); + +// Build the N-rank placement from cfg (uniform | balanced | file) and derive +// this rank's MoeHybridPlacement and slot-validity cache. hotness_csv_path +// (from LUCE_DS4_HOTNESS_CSV) is required for `balanced` and used by +// `uniform` when replicate_hot > 0; may be nullptr otherwise. +bool ds4_cluster_build_placement(const cluster::ClusterConfig & cfg, + const DeepSeek4Weights & w, + const std::string * hotness_csv_path, + Ds4ClusterRuntime & rt, + std::string * err); + +// Build this rank's expert storage: hot set = rt.rank_placement, cold owner +// None. `w` must already be loaded with TargetLoadPlan.skip_expert_tensors +// (expert tensors present as metadata only). Logs resident expert bytes. +bool ds4_cluster_init_experts(const std::string & model_path, + ggml_backend_t backend, + const DeepSeek4Weights & w, + Ds4ClusterRuntime & rt, + MoeHybridStorage & storage, + std::string * err); + +// Zero out every route this rank does not evaluate: selected id -> -1, +// weight -> 0. selected/weights are [route_width, n_tokens] row-major +// (token-major), slot = token index. Routing statistics must be observed +// before this call. +void ds4_cluster_mask_routes(const Ds4ClusterRuntime & rt, + int layer, + int32_t * selected, + float * weights, + int route_width, + int n_tokens); + +// An all-reduce as an ordinary graph node. `partial` is summed +// across the ranks on the executing backend's stream, so no host round trip +// separates the kernels that produce it from those that consume the sum. +// Returns `partial` unchanged for a single-rank runtime and nullptr when the +// tensor cannot be reduced (non-contiguous or not F32). Errors during the +// compute land in rt.node_error. `inplace` sums into `partial` itself (no +// device copy first); only for a partial that no other node reads. +ggml_tensor * ds4_cluster_allreduce_node(ggml_context * ctx, + ggml_tensor * partial, + Ds4ClusterRuntime & rt, + bool inplace = false); + +// LUCE_CLUSTER_HYBRID_EXCHANGE=1: the decode FFN reduction of a rank whose +// experts are split between its dGPU (hot + shared) and iGPU (cold) runs as +// export (dGPU) -> combine + exchange (iGPU) -> import (dGPU); see +// cluster::FastReduce::hyb_export. ds4_cluster_hybrid_exchange_nodes returns +// the dGPU tensor holding the reduced sum (null when unavailable) and the iGPU +// combine node, which the caller must keep in the graph and on the iGPU. +bool ds4_cluster_hybrid_exchange_enabled(const Ds4ClusterRuntime & rt); +ggml_tensor * ds4_cluster_hybrid_exchange_nodes(ggml_context * ctx, + ggml_tensor * hot, + ggml_tensor * shared, + ggml_tensor * cold, + Ds4ClusterRuntime & rt, + ggml_tensor ** combine_node, + ggml_tensor ** export_node); + +// LUCE_CLUSTER_ATTENTION_PARALLEL: split the attention heads across the ranks +// (on in cluster mode; =0 keeps attention replicated). Cached after the first +// call. +bool ds4_cluster_attention_parallel_enabled(); +// Band-pipelined prefill on a cluster rank: one part's cold FFN runs beside +// the next part's attention. On when every band of the pass has at least +// kDs4ClusterPipelineMinBand rows; smaller bands split into parts too short +// to pay for the overlap. LUCE_CLUSTER_PREFILL_PIPELINE=1 forces it on, =0 +// off. A pure function of the bands, which every rank agrees on. +constexpr int kDs4ClusterPipelineMinBand = 1536; +bool ds4_cluster_prefill_pipeline_enabled(const std::vector & bands); + +// True when this runtime may use the fused whole-model graph: the +// opt-in is set, a real multi-rank communicator is attached, the shared +// expert is replicated, and expert ownership is a function of the expert id +// alone. Replicated experts are owned per token slot, which the fused graph's +// per-expert owner LUT cannot express, so they fall back to the per-layer path. +bool ds4_cluster_fused_graph_available(const Ds4ClusterRuntime * rt); + +// In-graph all-reduce (default). LUCE_CLUSTER_NO_INGRAPH_ALLREDUCE=1 falls back +// to the host-enqueued per-layer all-reduce, which also gives up the +// fused whole-model graph. Cached after the first call. +bool ds4_cluster_ingraph_allreduce_enabled(); + +// Host-resident variant used by the per-layer path: uploads to the runtime's +// device scratch, all-reduces on the backend stream, waits (deadline +// cfg->timeout_ms) and downloads the sum back into partial_host. +bool ds4_cluster_allreduce_layer_host(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + float * partial_host, + int n_embd, + int n_tokens, + int layer, + DeepSeek4StepTelemetry * telemetry, + std::string * err); + +// ─── Lucebox additions ──────────────────────────────────────────────── +// Engram tables live on rank 0 only (a rank's sparse model copy leaves them +// out). Rank 0's decoded rows reach the other ranks by broadcast: as a graph +// node in the fused verify/decode graph, host-staged in the per-layer path. +ggml_tensor * ds4_cluster_broadcast_node(ggml_context * ctx, + ggml_tensor * value, + Ds4ClusterRuntime & rt); +// The same broadcast, exact on a bf16 fast-reduce wire (f32 payload). +ggml_tensor * ds4_cluster_broadcast_exact_node(ggml_context * ctx, + ggml_tensor * value, + Ds4ClusterRuntime & rt); +bool ds4_cluster_broadcast_host(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + float * data, + size_t n, + std::string * err); +// Every rank's value in rank order: an all-gather built on the all-reduce +// (one writer per slot). Used to agree on load-time sizes that would +// otherwise depend on each rank's free memory, such as the prefill chunk. +bool ds4_cluster_allgather_i32(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + int32_t value, + std::vector & out, + std::string * err); +// True when this rank evaluates `expert` of `layer` for some token slot +// (owned or replicated): the experts its tiers hold. +bool ds4_cluster_resident(const Ds4ClusterRuntime & rt, int layer, int expert); + +} // namespace luce::common diff --git a/server/src/deepseek4/deepseek4_dspark.cpp b/server/src/deepseek4/deepseek4_dspark.cpp index ad990cf20..594aa8522 100644 --- a/server/src/deepseek4/deepseek4_dspark.cpp +++ b/server/src/deepseek4/deepseek4_dspark.cpp @@ -136,7 +136,8 @@ const char * deepseek4_dspark_last_error() { return g_dspark_err.c_str(); } bool load_deepseek4_dspark_drafter(const std::string & path, ggml_backend_t backend, - DSparkDrafter & out) { + DSparkDrafter & out, + const DSparkLoadSplit * split) { g_dspark_err.clear(); ggml_context * meta = nullptr; gguf_init_params gip{}; @@ -257,6 +258,26 @@ bool load_deepseek4_dspark_drafter(const std::string & path, w.layers.resize(n_layer); w.backend = backend; + // A cluster split keeps only this rank's routed experts: a contiguous + // range of each expert tensor's expert axis, read straight from the file. + const bool expert_split = split && split->size > 1 && split->rank >= 0 && + split->rank < split->size && w.n_expert % split->size == 0; + const int64_t n_expert_local = expert_split ? w.n_expert / split->size : w.n_expert; + const int64_t expert_first = expert_split ? (int64_t) split->rank * n_expert_local : 0; + const auto is_routed_expert = [](const char * name) { + const std::string s(name); + const auto ends = [&s](const char * suffix) { + const size_t n = std::strlen(suffix); + return s.size() >= n && s.compare(s.size() - n, n, suffix) == 0; + }; + return ends("ffn_gate_exps.weight") || ends("ffn_up_exps.weight") || + ends("ffn_down_exps.weight"); + }; + out.split_rank = expert_split ? split->rank : 0; + out.split_size = expert_split ? split->size : 1; + out.expert_first = (int) expert_first; + out.expert_local = expert_split ? (int) n_expert_local : 0; + // Validate the complete tensor table before allocating device memory or // forming an absolute file offset. The drafter contract intentionally has // no embedding/lm-head tensors; rejecting unknown tensors prevents a stray @@ -305,6 +326,23 @@ bool load_deepseek4_dspark_drafter(const std::string & path, return false; } + if (expert_split) { + for (int64_t ti = 0; ti < n_tensors; ++ti) { + const char * tname = gguf_get_tensor_name(g, ti); + if (!is_routed_expert(tname)) continue; + ggml_tensor * t = ggml_get_tensor(meta, tname); + if (!t || t->ne[2] != w.n_expert || t->ne[3] != 1) { + set_err(std::string("cannot split the routed experts of ") + tname); + gguf_free(g); + ggml_free(meta); + out = DSparkDrafter{}; + return false; + } + t->ne[2] = n_expert_local; + t->nb[3] = t->nb[2] * (size_t) n_expert_local; + } + } + // ── Allocate validated tensors into one backend buffer ────────────── ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(meta, backend); if (!buf) { @@ -322,8 +360,12 @@ bool load_deepseek4_dspark_drafter(const std::string & path, const char * tname = gguf_get_tensor_name(g, ti); ggml_tensor * t = ggml_get_tensor(meta, tname); const size_t tensor_off = gguf_get_tensor_offset(g, ti); - const size_t off = data_off + tensor_off; // preflight proved no overflow - const size_t size = gguf_get_tensor_size(g, ti); + size_t off = data_off + tensor_off; // preflight proved no overflow + size_t size = gguf_get_tensor_size(g, ti); + if (expert_split && is_routed_expert(tname)) { + off += (size_t) expert_first * t->nb[2]; + size = ggml_nbytes(t); + } staging.resize(size); std::memcpy(staging.data(), file_bytes + off, size); ggml_backend_tensor_set(t, staging.data(), 0, size); @@ -482,11 +524,11 @@ bool load_deepseek4_dspark_drafter(const std::string & path, expect_shape(L.ffn_exp_probs_b, p + "exp_probs_b.bias", {w.n_expert}); } expect_shape(L.ffn_gate_exps, p + "ffn_gate_exps.weight", - {n_embd, w.n_ff_exp, w.n_expert}); + {n_embd, w.n_ff_exp, n_expert_local}); expect_shape(L.ffn_up_exps, p + "ffn_up_exps.weight", - {n_embd, w.n_ff_exp, w.n_expert}); + {n_embd, w.n_ff_exp, n_expert_local}); expect_shape(L.ffn_down_exps, p + "ffn_down_exps.weight", - {w.n_ff_exp, n_embd, w.n_expert}); + {w.n_ff_exp, n_embd, n_expert_local}); expect_shape(L.ffn_gate_shexp, p + "ffn_gate_shexp.weight", {n_embd, w.n_ff_exp}); expect_shape(L.ffn_up_shexp, p + "ffn_up_shexp.weight", {n_embd, w.n_ff_exp}); expect_shape(L.ffn_down_shexp, p + "ffn_down_shexp.weight", {w.n_ff_exp, n_embd}); diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index bf64afd18..6015695aa 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -41,6 +41,10 @@ #include #include +namespace luce::cluster { +struct Ds4ClusterHooks; // cluster/cluster_decision_hooks.h +} + namespace luce::common { // The drafter weights. `core` reuses DeepSeek4Weights for the n_layer decoder @@ -80,13 +84,31 @@ struct DSparkDrafter { // architecture makes the target GPU's next GEMM of that name fail. bool flash_attention = false; std::vector capture_layer_ids; // [40,41,42] + + // Cluster split (LUCE_CLUSTER_DRAFT_SPLIT=1, two ranks): this rank drafts + // with its half of the heads and holds routed experts [expert_first, + // expert_first + expert_local); the partial sums meet in two in-graph + // all-reduces a layer. split_size 1 = the whole drafter on one rank. + int split_rank = 0; + int split_size = 1; + int expert_first = 0; + int expert_local = 0; + void * cluster_rt = nullptr; // the target's Ds4ClusterRuntime, set after load +}; + +// A cluster rank's share of the drafter at load: routed experts +// [rank, rank + 1) * n_expert / size of every expert tensor. +struct DSparkLoadSplit { + int rank = 0; + int size = 1; }; // Load a "deepseek4-dflash-draft" GGUF into `out`. Returns false on error; // deepseek4_dspark_last_error() has the message. bool load_deepseek4_dspark_drafter(const std::string & path, ggml_backend_t backend, - DSparkDrafter & out); + DSparkDrafter & out, + const DSparkLoadSplit * split = nullptr); // Copy only the DSpark Markov/confidence heads to `backend` and repoint the // drafter fields. Decoder-block tensors remain on `core.backend`. @@ -315,6 +337,9 @@ bool run_deepseek4_dspark_spec_decode( // (deepseek4_spec_sampling.h), so tokens follow the request's sampler. DSparkSpecSampling * sampling = nullptr, // Polled before applying the budget or emitting another token. - const std::function & is_cancelled = {}); + const std::function & is_cancelled = {}, + // Cluster lockstep hooks: head/worker + // hooks broadcast/receive DraftMsg and AcceptMsg; nullptr = one box. + luce::cluster::Ds4ClusterHooks * cluster_hooks = nullptr); } // namespace luce::common diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 787eb3918..1250d9680 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -24,11 +24,15 @@ #include "deepseek4_dspark.h" #include "thinking_budget.h" #include "deepseek4_internal.h" +#include "cluster/cluster_decision_hooks.h" +#include "cluster/fast_reduce.h" +#include "deepseek4_cluster.h" #include "deepseek4_snapshot.h" #include "deepseek4_roctx.h" #include "internal.h" #include "common/adaptive_spec_width.h" #include "common/dspark_head.h" +#include "common/cuda_graph_overrides.h" #include "ggml.h" #include "ggml-backend.h" @@ -43,6 +47,7 @@ #include #include #include +#include #include #include @@ -240,9 +245,19 @@ class DeepSeek4DFlashTarget : public DFlashTarget { } } bool copied = false; + // One launch for every layer's pair on a CUDA/HIP backend (a blit per + // tensor costs a dispatch each, four per layer); the pairs are + // disjoint, so their order does not matter. + std::vector descs; + const bool batch = ggml_backend_is_cuda(backend_); auto copy_pair = [&](ggml_tensor * src, ggml_tensor * dst) { if (!src) return; - ggml_backend_tensor_copy_async(backend_, backend_, src, dst); + if (batch && ggml_is_contiguous(src) && ggml_is_contiguous(dst) && + ggml_nbytes(src) == ggml_nbytes(dst)) { + descs.push_back({src->data, dst->data, ggml_nbytes(src)}); + } else { + ggml_backend_tensor_copy_async(backend_, backend_, src, dst); + } copied = true; }; for (const DeepSeek4SpecBoundaryCheckpointLayer & layer : @@ -252,6 +267,9 @@ class DeepSeek4DFlashTarget : public DFlashTarget { copy_pair(layer.index_kv_src, layer.index_kv_dst); copy_pair(layer.index_score_src, layer.index_score_dst); } + if (!descs.empty()) { + ggml_backend_cuda_copy_batch_async(backend_, descs.data(), (int) descs.size()); + } return copied; } void clear_snapshot() { free_deepseek4_snapshot(snap_); } @@ -316,9 +334,11 @@ bool spec_env_default_on(const char * name) { // fixed drafter/head work (DSpark computes its whole proposal block whatever // the verify width). Only the ratios matter: q3 costs almost as much as q4, // so high-acceptance text settles on q4/q5 while low-acceptance prose drops -// to q2/q3. observe() refines the table online, so other backends start from -// this curve instead of importing a DS4 policy. The fused verify cache is -// sized so every width stays resident across the ratio-4 phases. +// to q2/q3. The controller takes the curve as a shape: it scales it to the +// widths measured on this machine and measures a wider width before trusting +// the curve's price for it, so a machine with a flatter curve (a two-box +// cluster) widens on its own. The fused verify cache is sized so every width +// stays resident across the ratio-4 phases. constexpr int kDs4AdaptiveMinWidth = 2; // Depths whose confidence-head score decides the width. The head was // calibrated on q<=4 traffic (three candidates); the fourth candidate of a q5 @@ -431,7 +451,20 @@ bool device_rollback_enabled() { bool rollback_tensor_on(const ggml_tensor * t, ggml_backend_buffer_type_t buft) { if (!t) return true; const ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; - return buf && ggml_backend_buffer_get_type(buf) == buft; + const bool on = buf && ggml_backend_buffer_get_type(buf) == buft; + if (!on) { + // Device staging falls back to one copy per saved row; say which + // tensor forced that, once. + static bool reported = false; + if (!reported) { + reported = true; + std::fprintf(stderr, + "[ds4-spec] device rollback staging off: %s is in %s, not %s\n", + t->name, buf ? ggml_backend_buft_name(ggml_backend_buffer_get_type(buf)) : "(no buffer)", + ggml_backend_buft_name(buft)); + } + } + return on; } // Device staging needs every saved tensor in the backend's own device memory, @@ -756,8 +789,13 @@ void spec_rollback_apply(const DeepSeek4SpecRollback & rb, const DeepSeek4Weight // and a window is at most two tokens; its remaining rows are rewritten by // the next tokens before the window is pooled. void spec_restore_window_rows(const Ds4VerifyWindowRows & rows, const DeepSeek4Weights & w, - DeepSeek4Cache & cache, int verify_pos, int commit_pos) { - for (size_t il = 0; il < cache.layers.size() && il < rows.kv.size(); ++il) { + DeepSeek4Cache & cache, int verify_pos, int commit_pos, + ggml_backend_t backend) { + // Rows the fused verify left on the device go back in one batched copy on + // the compute stream, ordered before the next verify reads the state. + std::vector copies; + for (size_t il = 0; il < cache.layers.size() && + (il < rows.kv.size() || il < rows.kv_dev.size()); ++il) { DeepSeek4LayerCache & lc = cache.layers[il]; const int ratio = il < w.compress_ratios.size() ? (int) w.compress_ratios[il] : 0; if (!deepseek4_is_window_state(lc.attn_compressor, ratio)) continue; @@ -765,6 +803,21 @@ void spec_restore_window_rows(const Ds4VerifyWindowRows & rows, const DeepSeek4W const int window_start = commit_pos / ratio * ratio; for (int p = std::max(window_start, verify_pos); p < commit_pos; ++p) { const size_t t = (size_t) (p - verify_pos); + if (il < rows.kv_dev.size() && t < rows.kv_dev[il].size() && + t < rows.score_dev[il].size()) { + const ggml_tensor * kv = rows.kv_dev[il][t]; + const ggml_tensor * score = rows.score_dev[il][t]; + if (kv && score && ggml_nbytes(kv) >= row_bytes && ggml_nbytes(score) >= row_bytes) { + const size_t offset = (size_t) (p % ratio) * row_bytes; + copies.push_back({kv->data, (char *) lc.attn_compressor.state_kv->data + offset, + row_bytes}); + copies.push_back({score->data, + (char *) lc.attn_compressor.state_score->data + offset, + row_bytes}); + } + continue; + } + if (il >= rows.kv.size()) continue; if (rows.kv[il].size() < (t + 1) * row_bytes || rows.score[il].size() < (t + 1) * row_bytes) { continue; @@ -776,6 +829,76 @@ void spec_restore_window_rows(const Ds4VerifyWindowRows & rows, const DeepSeek4W rows.score[il].data() + t * row_bytes, offset, row_bytes); } } + if (!copies.empty()) { + ggml_backend_cuda_copy_batch_async(backend, copies.data(), (int) copies.size()); + } +} + +// The split draft head's Markov weights on the target's device, for a rank +// whose drafter sits elsewhere (--draft-device) or whose rows of markov_w2 do +// not start at row 0: markov_w1 whole (the previous token's row is any row) +// and the rank's rows of markov_w2, copied once per drafter. +struct SplitMarkovCopies { + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + ggml_tensor * w1 = nullptr; + ggml_tensor * w2 = nullptr; + uint64_t generation = 0; + ggml_backend_t backend = nullptr; + int64_t row_offset = -1; + int64_t n_rows = -1; + + void release() { + if (buf) ggml_backend_buffer_free(buf); + if (ctx) ggml_free(ctx); + buf = nullptr; + ctx = nullptr; + w1 = w2 = nullptr; + } +}; + +bool spec_tensor_on(const ggml_tensor * t, ggml_backend_t backend) { + return t && t->buffer && + ggml_backend_buffer_get_type(t->buffer) == ggml_backend_get_default_buffer_type(backend); +} + +bool split_markov_copies(SplitMarkovCopies & m, const DraftWeights & dw, ggml_backend_t backend, + int64_t row_offset, int64_t n_rows) { + const uint64_t generation = dspark_drafter_generation(); + if (m.w1 && m.generation == generation && m.backend == backend && + m.row_offset == row_offset && m.n_rows == n_rows) { + return true; + } + m.release(); + const ggml_tensor * s1 = dw.dspark.markov_w1; + const ggml_tensor * s2 = dw.dspark.markov_w2; + if (!s1 || !s2 || !ggml_is_contiguous(s1) || !ggml_is_contiguous(s2) || + row_offset + n_rows > s2->ne[1]) { + return false; + } + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * 4; + ip.no_alloc = true; + m.ctx = ggml_init(ip); + if (!m.ctx) return false; + m.w1 = ggml_new_tensor_2d(m.ctx, s1->type, s1->ne[0], s1->ne[1]); + m.w2 = ggml_new_tensor_2d(m.ctx, s2->type, s2->ne[0], n_rows); + m.buf = ggml_backend_alloc_ctx_tensors(m.ctx, backend); + if (!m.buf) { + m.release(); + return false; + } + std::vector rows(ggml_nbytes(s1)); + ggml_backend_tensor_get(s1, rows.data(), 0, rows.size()); + ggml_backend_tensor_set(m.w1, rows.data(), 0, rows.size()); + rows.resize((size_t) n_rows * s2->nb[1]); + ggml_backend_tensor_get(s2, rows.data(), (size_t) row_offset * s2->nb[1], rows.size()); + ggml_backend_tensor_set(m.w2, rows.data(), 0, rows.size()); + m.generation = generation; + m.backend = backend; + m.row_offset = row_offset; + m.n_rows = n_rows; + return true; } using SpecClock = std::chrono::steady_clock; @@ -939,8 +1062,32 @@ bool run_deepseek4_dspark_spec_decode( MoeExpertComputeRuntime * expert_runtime, MoeHybridRoutingStats * routing_stats, DSparkSpecSampling * sampling, - const std::function & is_cancelled) { + const std::function & is_cancelled, + luce::cluster::Ds4ClusterHooks * cluster_hooks) { const int n_embd = target_w.n_embd; + // Cluster lockstep. Only a real + // head/worker hook changes behaviour. + const bool cluster_spec = cluster_hooks && cluster_hooks->is_cluster(); + const bool cluster_spec_head = cluster_spec && cluster_hooks->is_head(); + // Workers verify the head's draft block, never their own, so with + // LUCE_CLUSTER_WORKER_SKIP_DRAFT=1 they skip the drafter forward and the + // head. The drafter can then sit off the worker's dGPU (--draft-device), + // which frees that memory for hot experts. + // LUCE_CLUSTER_DRAFT_SPLIT: the drafter itself is split across the ranks + // (half the heads and experts each, all-reduces in its graph), so every + // rank runs it every step, in the same place in the collective order. + const bool draft_split = cluster_spec && drafter.split_size > 1; + const bool worker_skips_draft = cluster_spec && !cluster_spec_head && !draft_split && + spec_env_flag("LUCE_CLUSTER_WORKER_SKIP_DRAFT"); + const uint64_t cluster_req = cluster_hooks ? cluster_hooks->current_request() : 0; + // A cluster request that drafts carries no thinking budget (the head + // decodes budgeted requests AR) and the head alone stops early (an empty + // draft block tells the workers): every rank must keep the same tokens. + if (cluster_spec && (!budget.settled() || is_cancelled)) { + std::fprintf(stderr, "[ds4-spec] a cluster rank drafts with a thinking budget or a " + "local cancel check; refusing\n"); + return false; + } const int n_tgt = drafter.n_target_layers; const int block = drafter.block_size; const int n_swa = target_w.n_swa; @@ -1054,7 +1201,18 @@ bool run_deepseek4_dspark_spec_decode( width <= std::min(q_cap, DS4_Q5_VERIFY_TOKENS); ++width) { width_cost_ms[(size_t) width] = kDs4VerifyWidthCostMs[width]; } - width_controller.set_relative_costs(width_cost_ms); + // The table is one machine's curve: a shape that prices the widths this + // machine has not measured yet, never mixed with its measurements. + width_controller.set_relative_costs(width_cost_ms, AdaptiveSpecWidth::CostSeed::kShape); + // Step costs belong to the machine and the context length, not to the + // text: every request starts from the costs earlier ones measured at a + // similar context (their warmup hold re-arms). + static std::mutex width_costs_mu; + static SpecWidthCostMemory width_costs; + { + std::lock_guard lock(width_costs_mu); + width_costs.load(width_controller, committed); + } if (timing && width_controller.enabled()) { std::fprintf(stderr, "[ds4-spec] adaptive width policy=%s\n", @@ -1144,6 +1302,81 @@ bool run_deepseek4_dspark_spec_decode( if (sampling) target.set_keep_logits(true); // verify returns every row's logits const SpecClock::time_point run_t0 = SpecClock::now(); + // LUCE_CLUSTER_DRAFT_VOCAB_SPLIT=1 (two ranks): the draft head reads the + // target's whole lm_head every step on the head's dGPU while the worker's + // waits for the block. Split it like the verifier's + // head (LUCE_CLUSTER_VOCAB_SPLIT): each rank projects its half of the + // vocabulary and of the Markov bias, the head's hidden states reach the + // worker exactly, and the ranks agree on every depth's token in the graph. + // Same tokens as the unsplit head. Both ranks run it once per step, the + // worker right before it waits for the head's block. + Ds4ClusterRuntime * split_rt = target_cache.cluster_rt; + const bool draft_vocab_split = cluster_spec && split_rt && split_rt->size() == 2 && + split_rt->fast && split_rt->fast->ok() && q_cap >= 2 && + spec_env_flag("LUCE_CLUSTER_DRAFT_VOCAB_SPLIT"); + DsparkVocabSplit vocab_split; + if (draft_vocab_split) { + const int split_rank = split_rt->rank(); + ggml_tensor * head = target.lm_head_tensor(); + const int64_t V = head->ne[1]; + const int64_t half = ((V / 2) + 255) / 256 * 256; + vocab_split.row_offset = split_rank == 0 ? 0 : half; + vocab_split.n_rows = split_rank == 0 ? half : V - half; + vocab_split.markov_w1 = dw.dspark.markov_w1; + vocab_split.markov_w2 = dw.dspark.markov_w2; + if (vocab_split.row_offset != 0 || !spec_tensor_on(dw.dspark.markov_w1, backend) || + !spec_tensor_on(dw.dspark.markov_w2, backend)) { + static SplitMarkovCopies split_copies; + if (!split_markov_copies(split_copies, dw, backend, vocab_split.row_offset, + vocab_split.n_rows)) { + std::fprintf(stderr, "[ds4-spec] draft vocab split: Markov weight copy failed\n"); + return false; + } + vocab_split.markov_w1 = split_copies.w1; + vocab_split.markov_w2 = split_copies.w2; + } + const int64_t r0 = vocab_split.row_offset; + // A split drafter leaves the same hidden states on every rank. + if (draft_split) { + vocab_split.share = [](ggml_context *, ggml_tensor * t) { return t; }; + } else { + vocab_split.share = [split_rt](ggml_context * ctx, ggml_tensor * t) { + return ds4_cluster_broadcast_exact_node(ctx, t, *split_rt); + }; + } + vocab_split.pick = [split_rt, r0, split_rank](ggml_context * ctx, ggml_tensor * logits) -> ggml_tensor * { + ggml_tensor * pair = ggml_ds4_argmax_pair(ctx, logits, (int) r0, split_rank, 2); + pair = ds4_cluster_allreduce_node(ctx, pair, *split_rt, /*inplace=*/true); + return pair ? ggml_ds4_argmax_pick(ctx, pair, 2) : nullptr; + }; + vocab_split.id = split_rt; + static bool split_reported = false; + if (!split_reported) { + split_reported = true; + std::fprintf(stderr, "[ds4-spec] draft vocab split: rank %d rows %lld..%lld of %lld (%s Markov weights)\n", + split_rank, (long long) r0, (long long) (r0 + vocab_split.n_rows), (long long) V, + vocab_split.markov_w1 == dw.dspark.markov_w1 ? "drafter's" : "copied"); + } + } + // One split head per step on every rank, in the same place in the + // collective order: the head's own, or one whose tokens nobody reads. + const auto run_split_head = [&](const float * hidden, bool read, std::vector & toks, + std::vector * conf, const float * conf_hidden) { + const ScopedCudaGraphOverrides head_products( + /*disable_graphs=*/false, /*mmvq_max_ncols=*/0, + /*skip_property_check=*/false, /*ds4_mix_mmv_max_tokens=*/0, + /*mmvq_batch_invariant=*/true); + split_rt->node_error.clear(); + const bool ran = dspark_markov_chain_vocab_split( + dw, backend, target.lm_head_tensor(), vocab_split, hidden, q_cap - 1, lt, read, + toks, conf, conf_hidden); + if (ran && read && !split_rt->node_error.empty()) { + std::fprintf(stderr, "[ds4-spec] draft vocab split: %s\n", split_rt->node_error.c_str()); + return false; + } + return ran; + }; + while (n_generated < n_gen) { if (is_cancelled && is_cancelled()) break; const SpecClock::time_point step_t0 = SpecClock::now(); @@ -1154,7 +1387,7 @@ bool run_deepseek4_dspark_spec_decode( // Noise block = [seed] + [MASK]*(block-1). SpecClock::time_point t0 = SpecClock::now(); - if (q_cap >= 2 && !forcing) { + if (q_cap >= 2 && (draft_split || (!forcing && !worker_skips_draft))) { noise_ids[0] = lt; for (int i = 1; i < block; i++) { noise_ids[i] = drafter.mask_token_id; @@ -1202,6 +1435,7 @@ bool run_deepseek4_dspark_spec_decode( draft_tok.clear(); draft_confidence.clear(); bool ds_ok = false; + bool split_head_ran = false; // Batched-verify exactness: the batch must not cross a ratio-4 // boundary except at its last token (state rows stay distinct and the // comp emission matches AR). Boundaries sit at p % 4 == 3. The @@ -1239,7 +1473,7 @@ bool run_deepseek4_dspark_spec_decode( draft_tok.push_back(budget.pending_close_token((size_t) i)); } ds_ok = true; - } else if (q_step_cap >= 2) { + } else if (q_step_cap >= 2 && !worker_skips_draft) { std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), sizeof(float) * (size_t) n_embd * block); if (use_confidence_width) { @@ -1247,13 +1481,39 @@ bool run_deepseek4_dspark_spec_decode( confidence_hidden.data(), sizeof(float) * (size_t) n_embd * block); } - ds_ok = dspark_markov_correct_greedy_chain_fused( + if (draft_vocab_split) { + // Every depth the head always offers, so the worker's graph + // matches; the controller's narrower cap keeps the prefix. + split_head_ran = true; + if (!run_split_head(padded_hidden.data(), /*read=*/true, draft_tok, + use_confidence_width ? &draft_confidence : nullptr, + use_confidence_width ? padded_confidence_hidden.data() : nullptr)) { + std::fprintf(stderr, "[ds4-spec] split draft head failed\n"); + ok = false; + break; + } + ds_ok = true; + if ((int) draft_tok.size() > q_step_cap) draft_tok.resize((size_t) q_step_cap); + if ((int) draft_confidence.size() > q_step_cap - 1) { + draft_confidence.resize((size_t) std::max(0, q_step_cap - 1)); + } + } else { + // The head projects every candidate column through the + // vocabulary rows; batch-invariant products read those rows + // once for all columns on the single-column kernel (the RDNA4 + // Q8_0 path). + const ScopedCudaGraphOverrides head_products( + /*disable_graphs=*/false, /*mmvq_max_ncols=*/0, + /*skip_property_check=*/false, /*ds4_mix_mmv_max_tokens=*/0, + /*mmvq_batch_invariant=*/true); + ds_ok = dspark_markov_correct_greedy_chain_fused( dw, backend, target.lm_head_tensor(), padded_hidden.data(), q_step_cap, lt, draft_tok, use_confidence_width ? &draft_confidence : nullptr, use_confidence_width ? padded_confidence_hidden.data() : nullptr, nullptr); + } if (!ds_ok) { ds_ok = dspark_markov_correct_greedy_chain(dw, backend, target, padded_hidden.data(), q_step_cap, lt, 0.0f, draft_tok); @@ -1305,6 +1565,51 @@ bool run_deepseek4_dspark_spec_decode( // tokens (drafted or forced close tokens) that are never returned. const int remaining_out = std::max(1, n_gen - n_generated); if ((int) draft_tok.size() > remaining_out) draft_tok.resize((size_t) remaining_out); + // The split draft head runs on every rank every step: here on a worker + // (it then waits for the head's block below with its share already + // queued) and on a head step that did not project (forced tokens). + if (draft_vocab_split && !split_head_ran) { + std::vector unused_tok; + if (!run_split_head((cluster_spec_head || draft_split) ? padded_hidden.data() : nullptr, + /*read=*/cluster_spec_head, unused_tok, nullptr, nullptr)) { + std::fprintf(stderr, "[ds4-spec] split draft head failed\n"); + ok = false; + break; + } + } + // Cluster lockstep: the head's draft block is authoritative. Every + // rank drafts (the drafter forwards run concurrently on all ranks, so + // the redundant work costs no wall-clock and keeps every feature + // window warm), but only rank 0's block and adaptive width are + // verified, so the verify batch is identical everywhere. An EMPTY + // block is the head's end marker for a head-only stop (cancel). + if (cluster_spec_head) { + std::string err; + if (!cluster_hooks->decide_draft(cluster_req, (uint32_t) steps, pos, draft_tok, &err)) { + std::fprintf(stderr, "[ds4-spec] cluster draft broadcast failed: %s\n", err.c_str()); + ok = false; + break; + } + } else if (cluster_spec) { + std::string err; + const bool got_draft = + cluster_hooks->decide_draft(cluster_req, (uint32_t) steps, pos, draft_tok, &err); + if (!got_draft) { + std::fprintf(stderr, "[ds4-spec] cluster draft missing: %s\n", err.c_str()); + ok = false; + break; + } + if (draft_tok.empty()) { + stop_requested = true; + break; + } + if (draft_tok[0] != lt) { + std::fprintf(stderr, "[ds4-spec] cluster draft seed %d != local seed %d\n", + draft_tok[0], lt); + ok = false; + break; + } + } const int q = (int) draft_tok.size(); // seed + candidates if (q >= 0 && q <= q_cap) width_steps[(size_t) q]++; tm_head += spec_ms_since(t0); @@ -1463,6 +1768,32 @@ bool run_deepseek4_dspark_spec_decode( } bonus = tgt_am[accept - 1]; // target's token at the accept point } + // Cluster: the head's accept/bonus are authoritative. Workers compute + // their own (identical when the expert-parallel graph is exact) but + // adopt the head's, so a numerically different argmax on one rank + // cannot desynchronize the rollback below. + if (cluster_spec) { + int32_t accept_msg = accept; + int32_t bonus_msg = bonus; + uint8_t accept_flags = 0; + std::string err; + if (!cluster_hooks->decide_accept(cluster_req, (uint32_t) steps, + accept_msg, bonus_msg, accept_flags, &err)) { + std::fprintf(stderr, "[ds4-spec] cluster accept exchange failed: %s\n", err.c_str()); + ok = false; + break; + } + if (!cluster_spec_head) { + if (accept_msg < 1 || accept_msg > q) { + std::fprintf(stderr, "[ds4-spec] cluster accept=%d outside [1,%d]\n", + accept_msg, q); + ok = false; + break; + } + accept = accept_msg; + bonus = bonus_msg; + } + } // Apply the shared controller and emit the accepted drafts/bonus in // order, without copying a block or advancing over a stopped suffix. // Any cut (including a forced close equal to the draft) defers its @@ -1498,8 +1829,13 @@ bool run_deepseek4_dspark_spec_decode( if (on_token && !on_token(t)) stop_requested = true; hit_eos = target.is_eos(t); if (stop_requested || hit_eos) { - accept = i; - bonus = t; + // A cluster head's local stop must not change what the ranks + // keep: the workers keep the whole accepted block, so the head + // does too (an EOS is seen by every rank at the same token). + if (!(cluster_spec && stop_requested && !hit_eos)) { + accept = i; + bonus = t; + } break; } } @@ -1541,7 +1877,7 @@ bool run_deepseek4_dspark_spec_decode( // compressors keep no ratio-4 halves, only pooled windows. spec_rollback_apply(rollback, target_w, target_cache, commit_pos, false); spec_restore_window_rows(target.last_window_rows(), target_w, target_cache, - pos, commit_pos); + pos, commit_pos, backend); } else if (!full_snap && accept < q && q > DS4_CONSERVATIVE_VERIFY_MAX_TOKENS) { // Rejected wide (q5) verify over positions [pos, pos + 4]. @@ -1639,7 +1975,49 @@ bool run_deepseek4_dspark_spec_decode( tm_verify / steps, tm_probe_submit / steps, tm_probe_wait / steps, tm_apply / steps, tm_feat / steps); } - if (hit_eos || stop_requested) break; + if (hit_eos || stop_requested) { + // Head-local stop while the workers would keep going: send the + // empty end marker for the step they are about to ask for. When + // the loop ends on every rank anyway (EOS, or the token budget is + // exhausted) no marker is sent. + if (cluster_spec_head && stop_requested && !hit_eos && n_generated < n_gen) { + std::vector end_marker; + std::string err; + // A split drafter runs on every rank as well: the workers + // are in its all-reduces first, so run this rank's share. + if (draft_split) { + const int stop_ctx_len = feat_count < n_swa ? feat_count : n_swa; + noise_ids[0] = lt; + for (int i = 1; i < block; i++) noise_ids[i] = drafter.mask_token_id; + if (!target.embed_tokens(noise_ids.data(), block, noise_embed.data()) || + !deepseek4_dspark_draft_forward( + drafter_backend, drafter, noise_embed.data(), + stop_ctx_len > 0 ? feat_win.data() : nullptr, stop_ctx_len, pos, + local_hidden, use_confidence_width ? &confidence_hidden : nullptr)) { + std::fprintf(stderr, "[ds4-spec] split drafter before the end marker failed\n"); + ok = false; + } else { + std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), + sizeof(float) * (size_t) n_embd * block); + } + } + // The workers queue their split head before they read the + // marker: release it with the head's share of that step. + std::vector unused_tok; + if (draft_vocab_split && + !run_split_head(padded_hidden.data(), /*read=*/true, unused_tok, nullptr, nullptr)) { + std::fprintf(stderr, "[ds4-spec] split draft head before the end marker failed\n"); + ok = false; + } + if (!cluster_hooks->decide_draft(cluster_req, (uint32_t) steps, pos, + end_marker, &err)) { + std::fprintf(stderr, "[ds4-spec] cluster end marker failed: %s\n", + err.c_str()); + ok = false; + } + } + break; + } } const double total_ms = spec_ms_since(run_t0); @@ -1664,6 +2042,19 @@ bool run_deepseek4_dspark_spec_decode( q_cap >= 3 ? width_steps[3] : 0, q_cap >= 4 ? width_steps[4] : 0, q_cap >= 5 ? width_steps[5] : 0); + { + std::lock_guard lock(width_costs_mu); + width_costs.store(width_controller, committed); + } + std::string costs = "[ds4-spec] measured step ms"; + for (int w = kDs4AdaptiveMinWidth; w <= q_cap; ++w) { + const float c = width_controller.measured_costs()[(size_t) w]; + char buf[32]; + if (std::isfinite(c)) std::snprintf(buf, sizeof(buf), " q%d=%.1f", w, c); + else std::snprintf(buf, sizeof(buf), " q%d=-", w); + costs += buf; + } + std::fprintf(stderr, "%s\n", costs.c_str()); if (use_confidence_width) { // Recent predicted against observed acceptance per depth and // the scale the controller applies to the head's scores. diff --git a/server/src/deepseek4/deepseek4_engram.h b/server/src/deepseek4/deepseek4_engram.h index 9c7d02b2d..bee231e30 100644 --- a/server/src/deepseek4/deepseek4_engram.h +++ b/server/src/deepseek4/deepseek4_engram.h @@ -225,7 +225,8 @@ ggml_tensor * deepseek4_build_engram_apply(ggml_context * ctx, const DeepSeek4Layer & L, int n_embd, int n_hc, float rms_eps, - ggml_tensor ** gate_out = nullptr); + ggml_tensor ** gate_out = nullptr, + ggml_tensor * kv_pre = nullptr); // engram_wkv @ keys, computed ahead // Runs the apply on residual copies kept in host memory (the host // hyper-connection paths): uploads them with their keys, computes on diff --git a/server/src/deepseek4/deepseek4_engram_apply.cpp b/server/src/deepseek4/deepseek4_engram_apply.cpp index 22ffc0498..aca977700 100644 --- a/server/src/deepseek4/deepseek4_engram_apply.cpp +++ b/server/src/deepseek4/deepseek4_engram_apply.cpp @@ -58,7 +58,8 @@ ggml_tensor * deepseek4_build_engram_apply(ggml_context * ctx, const DeepSeek4Layer & L, int n_embd, int n_hc, float rms_eps, - ggml_tensor ** gate_out) { + ggml_tensor ** gate_out, + ggml_tensor * kv_pre) { if (!ctx || !h || !keys || !L.engram_wkv || !L.engram_q || !L.engram_k) return nullptr; const int64_t n_tokens = keys->ne[1]; GGML_ASSERT(h->ne[0] == n_embd && h->ne[1] == n_hc && h->ne[2] == n_tokens); @@ -66,10 +67,14 @@ ggml_tensor * deepseek4_build_engram_apply(ggml_context * ctx, GGML_ASSERT(L.engram_wkv->ne[1] == (int64_t) n_embd * (n_hc + 1)); // [n_embd * (n_hc + 1), n_tokens]: the n_hc keys first, the value last. - ggml_tensor * kv = ggml_mul_mat(ctx, L.engram_wkv, keys); - // The F16 projection must accumulate in F32: with F16 accumulation (the - // BLAS default on gfx1151) the update drifts past 2e-3 of its RMS. - ggml_mul_mat_set_prec(kv, GGML_PREC_F32); + ggml_tensor * kv = kv_pre; + if (!kv) { + kv = ggml_mul_mat(ctx, L.engram_wkv, keys); + // The F16 projection must accumulate in F32: with F16 accumulation (the + // BLAS default on gfx1151) the update drifts past 2e-3 of its RMS. + ggml_mul_mat_set_prec(kv, GGML_PREC_F32); + } + GGML_ASSERT(kv->ne[0] == L.engram_wkv->ne[1] && kv->ne[1] == n_tokens); ggml_tensor * key = ggml_view_3d(ctx, kv, n_embd, n_hc, n_tokens, (size_t) n_embd * ggml_element_size(kv), kv->nb[1], 0); ggml_tensor * value = ggml_view_3d(ctx, kv, n_embd, 1, n_tokens, diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 353f184fb..676a3f986 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -188,9 +188,15 @@ static int32_t ds4_fused_invalid_route() { static void ds4_fused_verify_refresh_hybrid_luts( const DeepSeek4Weights & w, MoeHybridStorage * hybrid, - DeepSeek4FusedDecodeGraph & fg) { + DeepSeek4FusedDecodeGraph & fg, + bool force_masked_routes = false) { if (!hybrid) return; - const int32_t invalid_route = ds4_fused_invalid_route(); + // A route this owner does not hold is either sent to local slot 0 and + // multiplied by a zero weight (correct, but it pays the full expert + // matmul) or given id -1, which the MMVQ kernels retire early with an + // exact zero. A cluster rank owns only its shard, so masking is where the + // sharding actually saves work: it is not optional there. + const int32_t invalid_route = force_masked_routes ? -1 : ds4_fused_invalid_route(); for (size_t hi = 0; hi < fg.hybrid_inputs.size(); ++hi) { const int il = (int) hi; const MoeHybridLayerStorage & layer = hybrid->layers[(size_t) il]; @@ -449,6 +455,9 @@ static bool ds4_build_fused_verify_graph( bool have_token_ids, const std::vector & capture_ids, MoeHybridStorage * hybrid, + // Cluster greedy verify: the head covers this rank's vocabulary half + // and the ranks agree on the argmax (part of the slot's shape key). + bool vocab_split, std::vector && shape_key, DeepSeek4PagedCache * paged_cache = nullptr, const std::vector> * paged_rows = nullptr, @@ -515,6 +524,7 @@ static bool ds4_build_fused_verify_graph( step_graph_free(fg.sg); fg.reset_nodes(); ex.reset(); + ex.vocab_split = vocab_split; fg.hash_ids.assign((size_t) w.n_layer, nullptr); ex.boundary_checkpoint.layers.resize((size_t) w.n_layer); if (hybrid) { @@ -526,6 +536,20 @@ static bool ds4_build_fused_verify_graph( ? backend_pair_capabilities(backend, hybrid->cold_backend) : BackendPairCapabilities{true, true}; const bool same_gpu_runtime = pair_capabilities.same_runtime; + // Expert-parallel cluster: the + // hybrid owners hold this rank's shard of the routed experts (hot on the + // primary GPU, cold on the secondary), routes to other ranks' experts are + // masked in the owner LUTs, the joined routed partial is summed across the + // ranks by an in-graph all-reduce on the primary's stream, and the + // shared expert is added to that sum. + Ds4ClusterRuntime * cluster_rt = cache ? cache->cluster_rt : nullptr; + const bool cluster_hybrid = hybrid && cluster_rt && cluster_rt->size() > 1; + std::vector cluster_nodes; + // Attention head parallelism: this rank's query heads, in + // whole output groups; 0 = attention replicated. + const int attn_head_begin = cluster_hybrid ? cluster_rt->attn_head_begin : 0; + const int attn_head_count = cluster_hybrid ? cluster_rt->attn_head_count : 0; + const int attn_n_head = attn_head_count > 0 ? attn_head_count : w.n_head; // The 43-layer q<=5 verifier currently consumes at most ~6.5 MiB of // metadata. Keeping 32 MiB leaves nearly 5x headroom while avoiding the @@ -556,12 +580,23 @@ static bool ds4_build_fused_verify_graph( engram->n_layers()); ggml_set_input(ex.engram_keys); } + // Engram tables live on rank 0; its decoded rows reach the other ranks. + ggml_tensor * engram_keys_all = ex.engram_keys; + if (cluster_hybrid && engram_keys_all) { + engram_keys_all = ds4_cluster_broadcast_node(ctx, ex.engram_keys, *cluster_rt); + if (!engram_keys_all) return false; + cluster_nodes.push_back(engram_keys_all); + } ex.q = q; ex.pos_q = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.pos_q); ex.neg_q = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.neg_q); ex.rawrows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.rawrows); ex.saved_rawrows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.saved_rawrows); + if (!paged_mode && q > 1) { + ex.pos_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); + ggml_set_input(ex.pos_rows); + } for (uint32_t ratio : w.compress_ratios) { if (ratio == 0 || ex.rows_for((int) ratio)) continue; auto & rows = ex.ratio_rows.emplace_back(); @@ -686,6 +721,41 @@ static bool ds4_build_fused_verify_graph( // The latest index source's top-k of each attention lane (paged: of each // row), for the layers after it (w.shared_index_topk). std::vector lane_selection((size_t) (paged_mode ? q : n_lanes), nullptr); + // Lane batch: the batched ops' selections, shared across layers, and the + // latest index source's selection for all lanes (lane_selection's twin). + Ds4LaneBatchSelections lane_batch_selections; + ggml_tensor * lane_batch_carrier = nullptr; + + // LUCE_DS4_ENGRAM_PRECOMPUTE= (hybrid exchange): an Engram layer's + // projection engram_wkv @ keys depends only on the step's tokens, so it is + // computed in row chunks of that many rows, one after each of the + // preceding layers' FFN exports, where the dGPU waits for the iGPU's sum; + // the Engram layer computes whatever rows remain and joins them. Each row + // is the same MMVF product as in the whole projection. + static const int engram_chunk_rows = [] { + const char * v = std::getenv("LUCE_DS4_ENGRAM_PRECOMPUTE"); + return v && *v ? std::max(0, std::atoi(v)) : 0; + }(); + std::vector> engram_kv_parts(engram ? (size_t) engram->n_layers() : 0u); + std::vector engram_kv_done(engram_kv_parts.size(), 0); + // The chunk (Engram layer index, first row) queued after layer il's export. + auto engram_chunk_for = [&](int il, int & e_out, int64_t & r0_out) -> bool { + if (!engram || engram_chunk_rows <= 0 || !engram_keys_all) return false; + for (int e = 0; e < engram->n_layers(); ++e) { + const int il_e = engram->layer_id(e); + const int64_t rows = w.layers[(size_t) il_e].engram_wkv + ? w.layers[(size_t) il_e].engram_wkv->ne[1] : 0; + const int n_chunks = (int) ((rows + engram_chunk_rows - 1) / engram_chunk_rows); + // Chunks fill the layers just before il_e, the last chunk last. + const int first = std::max(0, il_e - n_chunks); + if (il < first || il >= il_e) continue; + const int64_t r0 = (int64_t) (il - (il_e - n_chunks)) * engram_chunk_rows; + if (r0 < 0 || r0 >= rows || r0 != engram_kv_done[(size_t) e]) continue; + e_out = e; r0_out = r0; + return true; + } + return false; + }; for (int il = 0; il < w.n_layer; ++il) { const DeepSeek4Layer & L = w.layers[(size_t) il]; @@ -703,10 +773,27 @@ static bool ds4_build_fused_verify_graph( // ── Engram: the n-gram memory enters the residual copies ── if (const int e = engram ? engram->layer_index(il) : -1; e >= 0) { - ggml_tensor * keys = ggml_view_2d(ctx, ex.engram_keys, ex.engram_keys->ne[0], q, - ex.engram_keys->nb[1], (size_t) e * ex.engram_keys->nb[2]); + ggml_tensor * keys = ggml_view_2d(ctx, engram_keys_all, engram_keys_all->ne[0], q, + engram_keys_all->nb[1], (size_t) e * engram_keys_all->nb[2]); ggml_tensor * h3 = ggml_reshape_3d(ctx, ggml_cont(ctx, hc_cur), n_embd, n_hc, q); - ggml_tensor * updated = deepseek4_build_engram_apply(ctx, h3, keys, L, n_embd, n_hc, w.rms_eps); + ggml_tensor * kv_pre = nullptr; + if (!engram_kv_parts[(size_t) e].empty()) { + const int64_t rows = L.engram_wkv->ne[1]; + const int64_t done = engram_kv_done[(size_t) e]; + for (ggml_tensor * part : engram_kv_parts[(size_t) e]) { + kv_pre = kv_pre ? ggml_concat(ctx, kv_pre, part, 0) : part; + } + if (done < rows) { + ggml_tensor * rest = ggml_mul_mat(ctx, + ggml_view_2d(ctx, L.engram_wkv, L.engram_wkv->ne[0], rows - done, + L.engram_wkv->nb[1], (size_t) done * L.engram_wkv->nb[1]), + keys); + ggml_mul_mat_set_prec(rest, GGML_PREC_F32); + kv_pre = ggml_concat(ctx, kv_pre, rest, 0); + } + } + ggml_tensor * updated = deepseek4_build_engram_apply(ctx, h3, keys, L, n_embd, n_hc, w.rms_eps, + nullptr, kv_pre); if (!updated) return false; hc_cur = ggml_reshape_2d(ctx, updated, (int64_t) n_embd * n_hc, q); } @@ -984,7 +1071,8 @@ static bool ds4_build_fused_verify_graph( DeepSeek4PreparedProjectedLane batched_lanes{}; ggml_tensor * lane_contexts = nullptr; if (n_lanes > 1) { - batched_lanes = build_mla_qkv_projection(ctx, normed_all, w, L, q); + batched_lanes = build_mla_qkv_projection(ctx, normed_all, w, L, q, 0, + attn_head_begin, attn_head_count); build_mla_qkv_rope(ctx, batched_lanes, w, ds4_rope_params(w, ratio), q, ex.pos_q, /*fuse_q_rope=*/false); } @@ -992,6 +1080,84 @@ static bool ds4_build_fused_verify_graph( const auto it = std::find(ex.window_layers.begin(), ex.window_layers.end(), il); return it == ex.window_layers.end() ? -1 : (int) (it - ex.window_layers.begin()); }(); + // LUCE_DS4_LANE_BATCH: single-token lanes over one padded compressed + // span hand their attention to one op (build_mla_lane_batch_attention). + bool lane_batch = false; + int lane_batch_padded = 0; + if (n_lanes > 1 && attn_lane_q == 1 && ds4_lane_batch_enabled()) { + lane_batch_padded = ds4_fused_lane_padded(w, *comp_lc, il, kv_start); + bool uniform = true; + for (int lane = 1; lane < n_lanes && uniform; ++lane) { + uniform = ds4_fused_lane_padded(w, *comp_lc, il, kv_start + lane) == + lane_batch_padded; + } + lane_batch = uniform && + ds4_lane_batch_layer_ok(w, il, *lc, *comp_lc, lane_batch_padded, fg.identity_rows); + } + // Past top_k compressed rows the lanes attend over V4.1's shared + // top-k: the index source selects for all lanes in one chain and the + // layers after it reuse that selection. + const bool external_selection = lane_batch && + lane_batch_padded > w.n_indexer_top_k; + std::vector deferrals(lane_batch ? (size_t) n_lanes : 0u); + for (DeepSeek4LaneDeferral & deferral : deferrals) { + deferral.external_selection = external_selection; + } + const int64_t lane_batch_mask_off = mask_off; + // Lane batch over a ratio-1 kv source: one width-q compressor pass + // writes every lane's compressed row and index key ahead of the lanes, + // which then only read them (the batched attention masks each lane's + // later rows, as it does for the per-lane writes). + ggml_tensor * batched_comp_source = nullptr; + ggml_tensor * batched_index_source = nullptr; + // Lane batch over a stateful kv source (ratio > 1): the compressor's kv + // and gate projections are per-token products, so they are evaluated + // at width q and each lane takes its column (prepared_kv/score); the + // state writes and pooling still run per lane, in lane order. F16 + // MMVF picks its block size from K alone at these shapes, so every + // column is the product the lane would have computed. + if (lane_batch && ratio > 1 && attn_lane_q == 1 && + L.attn_compressor_kv && L.attn_compressor_gate && + L.attn_compressor_kv->type == GGML_TYPE_F16 && + L.attn_compressor_gate->type == GGML_TYPE_F16 && + deepseek4_is_kv_source(w, il) && !batched_lanes.compressor_kv) { + batched_lanes.compressor_kv = ggml_mul_mat(ctx, L.attn_compressor_kv, normed_all); + batched_lanes.compressor_score = ggml_mul_mat(ctx, L.attn_compressor_gate, normed_all); + } + if (lane_batch && ratio == 1 && ex.pos_rows && attn_lane_q == 1 && + L.attn_compressor_kv && deepseek4_is_kv_source(w, il)) { + DeepSeek4MlaLaneBindings b0 = deepseek4_contiguous_lane_bindings( + w, il, *lc, *comp_lc, kv_start + q - 1); + const DeepSeek4Layer * latent_index = + (L.indexer_k && b0.index_comp_kv) ? &L : nullptr; + std::vector bi32; + std::vector bi64; + ggml_tensor * last_row = ggml_view_2d(ctx, normed_all, n_embd, 1, normed_all->nb[1], + (size_t) (q - 1) * normed_all->nb[1]); + build_compressor_step(ctx, gf, last_row, + L.attn_compressor_ape, L.attn_compressor_kv, + L.attn_compressor_gate, L.attn_compressor_norm, + *b0.attn_compressor, b0.comp_kv, ratio, w.head_dim, + kv_start + q - 1, w.n_rot, w.rms_eps, + w.compress_rope_freq_base, w.rope_scale_factor, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + (int) w.rope_orig_ctx, + nullptr, nullptr, ex.pos_rows, ex.pos_q, + bi64, bi32, &batched_comp_source, nullptr, + normed_all, q, kv_start, false, nullptr, false, + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, + latent_index, b0.index_comp_kv, &batched_index_source); + if (!bi32.empty() || !bi64.empty()) { + std::fprintf(stderr, + "[ds4-fused-verify] layer %d: lane batch compressor took dynamic inputs\n", il); + return false; + } + if (batched_comp_source == b0.comp_kv || + (latent_index && (!batched_index_source || batched_index_source == b0.index_comp_kv))) { + batched_comp_source = nullptr; // not written: the lanes run it + batched_index_source = nullptr; + } + } for (int lane = 0; lane < n_lanes; ++lane) { const int aq = attn_lane_q; const int a_start = lane * aq; @@ -1066,20 +1232,29 @@ static bool ds4_build_fused_verify_graph( a_kv_start, aq, &ain, i32b, i32ab, i64ab, &f32ab, attention_impl, &ex.boundary_checkpoint.layers[(size_t) il], {}, - &lane_selection[(size_t) lane]); + &lane_selection[(size_t) lane], + attn_head_begin, attn_head_count); } else { const DeepSeek4PreparedProjectedLane lane_proj = ds4_slice_projected_lane(ctx, batched_lanes, w, a_start, ain.rope_pos); DeepSeek4MlaLaneBindings bindings = deepseek4_contiguous_lane_bindings( w, il, *lc, *comp_lc, a_token_pos); bindings.index_selection = &lane_selection[(size_t) lane]; + bindings.head_begin = attn_head_begin; + bindings.head_count = attn_head_count; + if (batched_comp_source) { + bindings.write_comp = false; + bindings.comp_kv = batched_comp_source; + if (batched_index_source) bindings.index_comp_kv = batched_index_source; + } if (!build_mla_attention_lane_core( ctx, gf, normed, w, L, bindings, il, a_kv_start, aq, &ain, - i32b, i32ab, i64ab, &f32ab, attention_impl, &lane_proj, &lane_out)) { + i32b, i32ab, i64ab, &f32ab, attention_impl, &lane_proj, &lane_out, + nullptr, {}, lane_batch ? &deferrals[(size_t) lane] : nullptr)) { lane_out = nullptr; } } - if (!lane_out) return false; + if (lane_batch ? !deferrals[(size_t) lane].deferred : !lane_out) return false; if (!i32b.empty() || !i32ab.empty() || !i64ab.empty() || !f32ab.empty()) { std::fprintf(stderr, @@ -1096,29 +1271,91 @@ static bool ds4_build_fused_verify_graph( ctx, st.state_kv, st.state_kv->ne[0], 1, st.state_kv->nb[1], row_off))); window_score_rows.push_back(ggml_cont(ctx, ggml_view_2d( ctx, st.state_score, st.state_score->ne[0], 1, st.state_score->nb[1], row_off))); + // Outputs: the rollback copies them back on the device. + ggml_set_output(window_kv_rows.back()); + ggml_set_output(window_score_rows.back()); ggml_build_forward_expand(gf, window_kv_rows.back()); ggml_build_forward_expand(gf, window_score_rows.back()); } if (n_lanes == 1) { attn_out = lane_out; - } else { + } else if (!lane_batch) { ggml_build_forward_expand(gf, lane_out); lane_contexts = lane_contexts ? ggml_concat(ctx, lane_contexts, lane_out, 1) : lane_out; } } + if (lane_batch) { + const DeepSeek4LaneDeferral & last = deferrals.back(); + const int64_t n_attn = (int64_t) w.n_swa + last.n_comp_attn; + GGML_ASSERT(lane_batch_mask_off + n_attn * n_lanes == mask_off); + ggml_tensor * selection = nullptr; + if (external_selection) { + ggml_tensor * visibility = deepseek4_is_index_source(w, il) + ? ggml_cont(ctx, ggml_view_2d( + ctx, fg.mask_bundle, lane_batch_padded, n_lanes, + n_attn * sizeof(float), + (size_t) (lane_batch_mask_off + w.n_swa) * sizeof(float))) + : nullptr; + selection = ds4_lane_batch_shared_selection( + ctx, w, L, il, batched_lanes.normalized_q_lora, normed_all, + last.index_comp_source, lane_batch_padded, kv_start, ex.pos_q, + visibility, &lane_batch_carrier); + if (!selection) { + std::fprintf(stderr, + "[ds4-fused-verify] layer %d: lane batch selection failed\n", il); + return false; + } + // The per-lane carriers predate this selection. + if (deepseek4_is_index_source(w, il)) { + std::fill(lane_selection.begin(), lane_selection.end(), nullptr); + } + } else { + selection = ds4_lane_batch_selection( + ctx, deferrals, fg.identity_rows, lane_batch_selections); + } + ggml_tensor * mask_f16 = fg.mask_bundle_f16 + ? ggml_view_2d(ctx, fg.mask_bundle_f16, n_attn, n_lanes, + n_attn * sizeof(ggml_fp16_t), + (size_t) lane_batch_mask_off * sizeof(ggml_fp16_t)) + : ggml_cast(ctx, ggml_view_2d(ctx, fg.mask_bundle, n_attn, n_lanes, + n_attn * sizeof(float), + (size_t) lane_batch_mask_off * sizeof(float)), + GGML_TYPE_F16); + lane_contexts = build_mla_lane_batch_attention( + ctx, gf, w, L, lc->raw_kv, batched_lanes.q, batched_lanes.kv, + ex.saved_rawrows, ex.rawrows, mask_f16, selection, + last.comp_source, last.n_comp_attn, attn_head_begin, attn_head_count); + if (!lane_contexts) return false; + } + if (!external_selection && w.shared_index_topk && ratio > 0 && + deepseek4_is_index_source(w, il)) { + lane_batch_carrier = nullptr; // the lanes selected one by one + } if (n_lanes > 1) { const Ds4RopeParams rope = ds4_rope_params(w, ratio); attn_out = build_tail_rope_3d( - ctx, ggml_reshape_3d(ctx, lane_contexts, w.head_dim, w.n_head, q), ex.neg_q, - w.n_rot, w.head_dim, w.n_head, q, rope.freq, rope.scale, rope.ext, rope.attn, + ctx, ggml_reshape_3d(ctx, lane_contexts, w.head_dim, attn_n_head, q), ex.neg_q, + w.n_rot, w.head_dim, attn_n_head, q, rope.freq, rope.scale, rope.ext, rope.attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope.n_ctx_orig); // The single-token decode graph never takes the grouped output kernel. attn_out = build_mla_output_projection( - ctx, ggml_reshape_2d(ctx, attn_out, (int64_t) w.head_dim * w.n_head, q), - w, L, q, /*allow_grouped=*/false); + ctx, ggml_reshape_2d(ctx, attn_out, (int64_t) w.head_dim * attn_n_head, q), + w, L, q, /*allow_grouped=*/false, 0, attn_head_begin, attn_head_count); } } + // Head split: the attention output is a partial sum over this rank's + // output groups; one in-graph all-reduce makes it whole. + if (cluster_hybrid && attn_head_count > 0) { + // The partial is a fresh product only this node reads: sum it in + // place, without a device copy (a contiguous one skips the cont). + attn_out = ds4_cluster_allreduce_node( + ctx, ggml_is_contiguous(attn_out) ? attn_out : ggml_cont(ctx, attn_out), + *cluster_rt, /*inplace=*/true); + if (!attn_out) return false; + cluster_nodes.push_back(attn_out); + } + // ── Batched HC post (attention) + HC pre (FFN) ── ggml_tensor * attn_batch = ggml_is_contiguous(attn_out) ? attn_out : ggml_cont(ctx, attn_out); @@ -1138,6 +1375,7 @@ static bool ds4_build_fused_verify_graph( // ── Batched FFN ── ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); ggml_tensor * ffn_out = nullptr; + MoeHybridGraphInputs * hyb_ex_inputs = nullptr; const MoeHybridGraphInputs * fused_hc_join_inputs = nullptr; const bool hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && have_token_ids && hash_tables[(size_t) il].loaded; @@ -1293,7 +1531,7 @@ static bool ds4_build_fused_verify_graph( // with whole routes the q=3 verify and q=1 decode are bit-identical). // LUCE_DS4_TP_HEAD4_TAIL2=1 restores the decomposition for bisection. const bool native_route_width = - !same_gpu_runtime || + cluster_hybrid || !same_gpu_runtime || mixed_policy.native_route_width || streamed || whole_routes_for_lanes; if (lane_q > 1 && hybrid_cfg.n_expert_used > 4 && has_hot_routed && @@ -1334,24 +1572,34 @@ static bool ds4_build_fused_verify_graph( } ffn_out = ggml_add(ctx, ffn_out, inputs.output); } else { + // Cluster hybrid exchange: keep the owners' partials apart; + // the reduction below adds them on the secondary owner. + inputs.external_join = cluster_hybrid && cluster_rt && + ds4_cluster_hybrid_exchange_enabled(*cluster_rt) && + L.ffn_gate_shexp && L.ffn_up_shexp && L.ffn_down_shexp && + cluster_rt->shexp_ff_count > 0 && !streamed; if (!build_moe_hybrid_ffn_graph( ctx, hybrid_schedule_graph, hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, selected, router_weights, - lane_q, inputs, true, allow_fused_combine, + lane_q, inputs, /*include_shared=*/!cluster_hybrid, + allow_fused_combine, join_mode, - streamed ? MoeHybridRouteBalance::Disabled - : MoeHybridRouteBalance::Allowed, + (streamed || cluster_hybrid) ? MoeHybridRouteBalance::Disabled + : MoeHybridRouteBalance::Allowed, streamed)) { return false; } ffn_out = inputs.output; + if (inputs.external_join && inputs.peer_raw && inputs.main_output) { + hyb_ex_inputs = &inputs; + } // The staggered pre-mix model joins like single-token decode // (verify lanes and gathered sequences match it exactly). if (lane_q > 1 && !w.hc_staggered_pre && native_route_width && - mixed_policy.fused_hc_join && + !cluster_hybrid && mixed_policy.fused_hc_join && inputs.main_output && inputs.peer_output) { fused_hc_join_inputs = &inputs; } @@ -1363,6 +1611,83 @@ static bool ds4_build_fused_verify_graph( return false; } + // Cluster: sum the ranks' routed partials as a graph node, + // then add this rank's copy of the replicated shared expert. Doing it + // in this order is what keeps the result identical to a single node: + // the shared expert is added exactly once, and the reduction sees + // only routed contributions. + if (cluster_hybrid) { + const bool have_shexp = + L.ffn_gate_shexp && L.ffn_up_shexp && L.ffn_down_shexp; + const int shexp_ff_count = cluster_rt->shexp_ff_count; + if (have_shexp && shexp_ff_count > 0) { + // Sharded: this rank's slice is a partial sum over the down + // projection's contraction axis, so it is added to the routed + // partial and the two ride in ONE reduction. + ggml_tensor * shared = build_shared_ffn_slice( + ctx, ffn_normed, w, L, + cluster_rt->shexp_ff_begin, shexp_ff_count); + if (!shared) return false; + // Queue the shared slice ahead of the owner join: it needs only + // the FFN input, so its kernels run while the secondary owner + // is still computing instead of after the join's wait. + ggml_build_forward_expand(gf, shared); + if (hyb_ex_inputs) { + // dGPU exports hot + shared, the iGPU adds its cold + // partial in the same order and exchanges, the dGPU + // imports the sum (no device join, one link crossing + // each way on the dGPU). + ggml_tensor * combine = nullptr; + ggml_tensor * exported = nullptr; + ggml_tensor * sum = ds4_cluster_hybrid_exchange_nodes( + ctx, hyb_ex_inputs->main_output, shared, hyb_ex_inputs->peer_raw, + *cluster_rt, &combine, &exported); + if (!sum) return false; + ggml_build_forward_expand(gf, exported); + cluster_nodes.push_back(exported); + ggml_set_output(combine); + ggml_build_forward_expand(gf, combine); + hyb_ex_inputs->cold_nodes.push_back(combine); + // An upcoming Engram layer's projection chunk, in the wait + // for the iGPU's sum (ahead of the import in the stream). + int engram_e = -1; + int64_t engram_r0 = 0; + if (engram_chunk_for(il, engram_e, engram_r0)) { + const DeepSeek4Layer & Le = w.layers[(size_t) engram->layer_id(engram_e)]; + const int64_t rows = std::min(engram_chunk_rows, Le.engram_wkv->ne[1] - engram_r0); + ggml_tensor * keys_e = ggml_view_2d(ctx, engram_keys_all, engram_keys_all->ne[0], q, + engram_keys_all->nb[1], + (size_t) engram_e * engram_keys_all->nb[2]); + ggml_tensor * part = ggml_mul_mat(ctx, + ggml_view_2d(ctx, Le.engram_wkv, Le.engram_wkv->ne[0], rows, + Le.engram_wkv->nb[1], (size_t) engram_r0 * Le.engram_wkv->nb[1]), + keys_e); + ggml_mul_mat_set_prec(part, GGML_PREC_F32); + ggml_build_forward_expand(gf, part); + cluster_nodes.push_back(part); + engram_kv_parts[(size_t) engram_e].push_back(part); + engram_kv_done[(size_t) engram_e] = engram_r0 + rows; + } + cluster_nodes.push_back(sum); + ffn_out = sum; + } else { + ffn_out = ggml_add(ctx, ffn_out, shared); + ffn_out = ds4_cluster_allreduce_node(ctx, ffn_out, *cluster_rt, /*inplace=*/true); + if (!ffn_out) return false; + cluster_nodes.push_back(ffn_out); + } + } else { + ffn_out = ds4_cluster_allreduce_node(ctx, ggml_cont(ctx, ffn_out), *cluster_rt, /*inplace=*/true); + if (!ffn_out) return false; + cluster_nodes.push_back(ffn_out); + if (have_shexp) { + ggml_tensor * shared = build_shared_ffn(ctx, ffn_normed, w, L); + if (!shared) return false; + ffn_out = ggml_add(ctx, ffn_out, shared); + } + } + } + // ── Batched HC post (FFN); capture at drafter layers ── if (fused_hc_join_inputs) { hc_cur = ggml_ds4_hc_post_split( @@ -1403,10 +1728,37 @@ static bool ds4_build_fused_verify_graph( ctx, omix, obase, hc_cur, n_hc, hc_out_weights.scale_data[0]); } ggml_tensor * out_normed = build_rms_norm(ctx, final_all, w.out_norm, w.rms_eps); - fg.logits = ds4_mul_mat_columns(ctx, w.output, out_normed, projection_columns); // [n_vocab, q] - ggml_set_output(fg.logits); - ggml_build_forward_expand(gf, fg.logits); - if (paged_mode || ds4_env_flag("LUCE_DS4_GPU_ARGMAX_VERIFY")) { + if (ex.vocab_split && cluster_rt && cluster_rt->size() == 2) { + // LUCE_CLUSTER_VOCAB_SPLIT: each rank projects its half of the + // vocabulary (the Q8_0 rows are the same products as in the full + // head) and the ranks exchange their column maxima as one slot each. + const int rank = cluster_rt->rank(); + const int64_t V = w.output->ne[1]; + const int64_t half = ((V / 2) + 255) / 256 * 256; + const int64_t r0 = rank == 0 ? 0 : half; + const int64_t nr = rank == 0 ? half : V - half; + static bool logged = false; + if (!logged) { + logged = true; + std::fprintf(stderr, "[ds4-fused-verify] vocab split: rank %d takes the argmax over rows " + "%lld..%lld of %lld\n", rank, (long long) r0, (long long) (r0 + nr), (long long) V); + } + ggml_tensor * out_rows = ggml_view_2d(ctx, w.output, w.output->ne[0], nr, w.output->nb[1], + (size_t) r0 * w.output->nb[1]); + fg.logits = ds4_mul_mat_columns(ctx, out_rows, out_normed, projection_columns); // [nr, q] + ggml_tensor * pair = ggml_ds4_argmax_pair(ctx, fg.logits, (int) r0, rank, 2); + pair = ds4_cluster_allreduce_node(ctx, pair, *cluster_rt, /*inplace=*/true); + if (!pair) return false; + cluster_nodes.push_back(pair); + ex.argmax = ggml_ds4_argmax_pick(ctx, pair, 2); + ggml_set_output(ex.argmax); + ggml_build_forward_expand(gf, ex.argmax); + } else { + fg.logits = ds4_mul_mat_columns(ctx, w.output, out_normed, projection_columns); // [n_vocab, q] + ggml_set_output(fg.logits); + ggml_build_forward_expand(gf, fg.logits); + } + if (!ex.vocab_split && (paged_mode || ds4_env_flag("LUCE_DS4_GPU_ARGMAX_VERIFY"))) { ex.argmax = ggml_argmax(ctx, fg.logits); ggml_set_output(ex.argmax); ggml_build_forward_expand(gf, ex.argmax); @@ -1429,17 +1781,10 @@ static bool ds4_build_fused_verify_graph( ggml_set_output(ex.capture); ggml_build_forward_expand(gf, ex.capture); } - const auto concat_rows = [&](const std::vector & rows) { - ggml_tensor * all = nullptr; - for (ggml_tensor * row : rows) all = all ? ggml_concat(ctx, all, row, 1) : row; - if (all) { - ggml_set_output(all); - ggml_build_forward_expand(gf, all); - } - return all; - }; - ex.window_kv = concat_rows(window_kv_rows); - ex.window_score = concat_rows(window_score_rows); + // The window rows stay where the lanes wrote them (no concat, no + // readback): a rejected suffix copies the few it needs back on the device. + ex.window_kv_rows = window_kv_rows; + ex.window_score_rows = window_score_rows; ex.boundary_checkpoint.available = std::any_of( ex.boundary_checkpoint.layers.begin(), ex.boundary_checkpoint.layers.end(), @@ -1533,6 +1878,14 @@ static bool ds4_build_fused_verify_graph( fg.sched, device_join_split); ggml_backend_sched_set_batch_split_copies( fg.sched, batch_split_copies); + // LUCE_DS4_SCHED_DEDICATED_COPIES=1: the verifier's split inputs own + // their storage (a few MiB), so the host never waits for a device + // to drain before queueing the next split. + static const bool dedicated_copies = [] { + const char * v = std::getenv("LUCE_DS4_SCHED_DEDICATED_COPIES"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + ggml_backend_sched_set_dedicated_copies(fg.sched, dedicated_copies); if (mixed_policy.scheduler_trace) { ggml_backend_sched_set_eval_callback( fg.sched, ds4_fused_verify_trace_node, nullptr); @@ -1548,12 +1901,16 @@ static bool ds4_build_fused_verify_graph( } }; for (ggml_tensor * context : attention_contexts) pin_main(context); + // The communicator lives on the primary GPU (ROCm 7.2.4's RCCL has + // no gfx1151 kernels): every collective node runs on its stream. + for (ggml_tensor * node : cluster_nodes) pin_main(node); pin_main(fg.inp_embed); pin_main(ex.engram_keys); pin_main(ex.pos_q); pin_main(ex.neg_q); pin_main(ex.rawrows); pin_main(ex.saved_rawrows); + pin_main(ex.pos_rows); for (const auto & rows : ex.ratio_rows) { pin_main(rows.ape); pin_main(rows.state); @@ -1666,7 +2023,8 @@ static bool ds4_build_fused_verify_graph( ggml_backend_sched_get_n_splits(fg.sched), q); } - ds4_fused_verify_refresh_hybrid_luts(w, hybrid, fg); + ds4_fused_verify_refresh_hybrid_luts(w, hybrid, fg, + /*force_masked_routes=*/cluster_hybrid); } else { if (!fg.sg.alloc) { fg.sg.alloc = ggml_gallocr_new( @@ -1741,6 +2099,27 @@ static void ds4_fused_cache_make_room(Ds4FusedVerifyCache & vc, size_t slot_limi } } +// Upper bound of a step's pinned traffic: every per-step input of the slot, +// and the readbacks. +static size_t ds4_fused_stage_bytes(const DeepSeek4Weights & w, + const DeepSeek4FusedDecodeGraph & fg, + const Ds4FusedVerifyCache::Extra & ex, + int q, const Ds4VerifyHooks * hooks) { + size_t up = 0; + const auto add = [&](const ggml_tensor * t) { if (t) up += ggml_nbytes(t) + 64; }; + add(fg.inp_embed); add(ex.engram_keys); add(ex.pos_q); add(ex.neg_q); + add(ex.rawrows); add(ex.saved_rawrows); add(ex.pos_rows); + for (const auto & rows : ex.ratio_rows) { add(rows.ape); add(rows.state); } + add(fg.i32_bundle); add(fg.i64_bundle); add(fg.identity_rows); add(fg.mask_bundle); + for (const ggml_tensor * hids : fg.hash_ids) add(hids); + size_t down = sizeof(int32_t) * (size_t) q + 256; + if (hooks && hooks->capture_layer_ids) { + down += sizeof(float) * hooks->capture_layer_ids->size() * (size_t) w.n_embd * (size_t) q + 64; + } + // The sparse mask path uploads ranges; give them room beside the rest. + return std::max(up, down) + up / 4; +} + static int ds4_try_fused_verify_step( Ds4FusedVerifyCache & vc, bool capture_cache, @@ -1805,6 +2184,11 @@ static int ds4_try_fused_verify_step( key.push_back(hybrid ? 1 : 0); key.push_back((int64_t) hooks->capture_layer_ids->size()); for (int layer : *hooks->capture_layer_ids) key.push_back(layer); + static const bool vocab_split_env = ds4_env_flag("LUCE_CLUSTER_VOCAB_SPLIT") && + ds4_env_flag("LUCE_DS4_GPU_ARGMAX_VERIFY"); + const bool vocab_split = vocab_split_env && hybrid && hooks->prefer_argmax_only && hooks->argmax_out && + cache.cluster_rt && cache.cluster_rt->size() == 2; + key.push_back(vocab_split ? 1 : 0); const int n_lanes = ds4_fused_attention_lanes(w, q); const int attn_lane_q = q / n_lanes; for (int il = 0; il < w.n_layer; ++il) { @@ -1868,7 +2252,7 @@ static int ds4_try_fused_verify_step( if (!ds4_build_fused_verify_graph(mc, *fg, *ex, backend, w, &cache, hc_weights, hc_out_weights, hash_tables, kv_start, q, token_ids != nullptr, - *hooks->capture_layer_ids, hybrid, + *hooks->capture_layer_ids, hybrid, vocab_split, std::move(key))) { fg->destroy(vc.backend, vc.peer_backend); ex->reset(); @@ -1880,13 +2264,50 @@ static int ds4_try_fused_verify_step( fg->last_use = vc.counter; // ── Fill inputs ── + // With pinned step I/O the mask bundle and the readbacks go through pinned + // slices (one wait for all of them, below) and the Engram rows are read on + // a worker while the graph already runs. + const bool pinned_io = + vc.stage.reserve(backend, ds4_fused_stage_bytes(w, *fg, *ex, q, hooks)); + vc.stage.reset(); + // The async read needs a place in the stream right before its consumer: + // the cluster's Engram broadcast callback (after layer 0's attention), + // which the graph has only when it was built for the hybrid cluster path. + Ds4ClusterRuntime * engram_rt = hybrid && cache.cluster_rt && cache.cluster_rt->size() > 1 + ? cache.cluster_rt : nullptr; + bool engram_async = false; if (ex->engram_keys) { - std::vector keys; - if (!ds4_engram_read_keys(w, cache.engram_tokens, token_ids, kv_start, q, keys, telemetry)) { - return -1; + const DeepSeek4EngramRuntime * engram = w.engram_runtime.get(); + if (pinned_io && engram_rt && engram && token_ids && ex->engram_keys->buffer) { + if (!vc.engram_async) vc.engram_async = std::make_unique(); + engram_async = vc.engram_async->start( + backend, w, cache.engram_tokens, token_ids, kv_start, q, + (size_t) engram->n_layers() * (size_t) q * engram->key_floats()); + } + if (!engram_async) { + std::vector keys; + if (!ds4_engram_read_keys(w, cache.engram_tokens, token_ids, kv_start, q, keys, telemetry)) { + return -1; + } + ds4_fv_set(ex->engram_keys, keys.data(), sizeof(float) * keys.size()); } - ds4_fv_set(ex->engram_keys, keys.data(), sizeof(float) * keys.size()); } + // A started read must be waited for on every exit, or the worker could + // still write the context the caller rolls back. + struct EngramJoin { + Ds4EngramAsync * reader; + DeepSeek4StepTelemetry * telemetry; + bool ok = true; + bool finish() { + if (!reader) return ok; + uint64_t us = 0; + ok = reader->finish(&us); + if (telemetry) telemetry->engram_read_us += us; + reader = nullptr; + return ok; + } + ~EngramJoin() { finish(); } + } engram_join{engram_async ? vc.engram_async.get() : nullptr, telemetry}; const auto set_t0 = Ds4TimingClock::now(); ds4_fv_set(fg->inp_embed, embed, sizeof(float) * (size_t) w.n_embd * q); std::vector iv(q); @@ -1899,6 +2320,10 @@ static int ds4_try_fused_verify_step( ds4_fv_set(ex->rawrows, lv.data(), sizeof(int64_t) * q); for (int i = 0; i < q; ++i) iv[(size_t) i] = (int32_t) lv[(size_t) i]; ds4_fv_set(ex->saved_rawrows, iv.data(), sizeof(int32_t) * q); + if (ex->pos_rows && ex->pos_rows->buffer) { + for (int i = 0; i < q; ++i) lv[(size_t) i] = kv_start + i; + ds4_fv_set(ex->pos_rows, lv.data(), sizeof(int64_t) * q); + } for (const auto & rows : ex->ratio_rows) { for (int i = 0; i < q; ++i) iv[(size_t) i] = (kv_start + i) % rows.ratio; ds4_fv_set(rows.ape, iv.data(), sizeof(int32_t) * q); @@ -2048,10 +2473,10 @@ static int ds4_try_fused_verify_step( for (const MaskRange & range : next_negative) { std::fill(maskv.begin() + range.first, maskv.begin() + range.second, -1.0e30f); - ggml_backend_tensor_set( - fg->mask_bundle, maskv.data() + range.first, - range.first * sizeof(float), - (range.second - range.first) * sizeof(float)); + // Without the pinned block the stage falls back to blocking copies. + vc.stage.set(backend, fg->mask_bundle, maskv.data() + range.first, + range.first * sizeof(float), + (range.second - range.first) * sizeof(float)); } } } @@ -2076,10 +2501,32 @@ static int ds4_try_fused_verify_step( row, (size_t) n_used * sizeof(int32_t)); } - ggml_backend_tensor_set(hids, hash_scratch.data(), 0, - sizeof(int32_t) * (size_t) n_used * q); + ds4_fv_set(hids, hash_scratch.data(), sizeof(int32_t) * (size_t) n_used * q); } } + // One wait for every pinned upload: the scheduler copies some inputs to + // the peer device with a blocking copy that is not ordered behind this + // stream. The Engram keys are queued after it: their upload waits on the + // GPU for the worker, not on the host. + if (pinned_io) ggml_backend_synchronize(backend); + // The root's broadcast sends the keys: its callback queues the upload + // behind the worker's word, so layer 0 runs while the rows are read. + // The other ranks' keys are overwritten by the broadcast, their worker + // only records the tokens in the n-gram context. + struct BroadcastHook { + Ds4ClusterRuntime * rt; + ~BroadcastHook() { if (rt) rt->before_broadcast = nullptr; } + } broadcast_hook{nullptr}; + if (engram_async) { + Ds4EngramAsync * reader = vc.engram_async.get(); + const int rank = engram_rt->rank(); + const size_t key_bytes = ggml_nbytes(ex->engram_keys); + engram_rt->before_broadcast = [reader, rank, key_bytes](void * data, size_t n, void * stream) { + if (rank != 0) return true; + return n * sizeof(float) == key_bytes && reader->enqueue_raw(stream, data, key_bytes); + }; + broadcast_hook.rt = engram_rt; + } if (telemetry) { telemetry->full_graph_set_us += ds4_elapsed_us(set_t0, Ds4TimingClock::now()); } @@ -2116,6 +2563,11 @@ static int ds4_try_fused_verify_step( /*mmvq_batch_invariant=*/n_lanes > 1); const enum ggml_status compute_status = ds4_fused_graph_compute(w, *fg, backend, hybrid, "ds4-fused-verify"); + if (broadcast_hook.rt) { + broadcast_hook.rt->before_broadcast = nullptr; + broadcast_hook.rt = nullptr; + } + if (!engram_join.finish()) return -1; if (compute_status != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[ds4-fused-verify] compute failed\n"); return -1; @@ -2129,10 +2581,34 @@ static int ds4_try_fused_verify_step( const int ncap = (int) hooks->capture_layer_ids->size(); const bool argmax_only = hooks->prefer_argmax_only && hooks->argmax_out && ex->argmax; + // The small readbacks go out together: queue every copy into pinned + // slices, then one wait. + const void * argmax_src = nullptr; + const void * capture_src = nullptr; + const bool want_capture = hooks->capture_out && ex->capture && ncap > 0; + const bool want_window = hooks->window_rows && !ex->window_layers.empty() && + !ex->window_kv_rows.empty() && + ex->window_kv_rows.size() == ex->window_score_rows.size() && + ex->window_kv_rows.size() % ex->window_layers.size() == 0; + if (pinned_io) { + vc.stage.reset(); + if (argmax_only) { + argmax_src = vc.stage.get(backend, ex->argmax, 0, sizeof(int32_t) * (size_t) q); + } + if (want_capture) { + capture_src = vc.stage.get(backend, ex->capture, 0, + sizeof(float) * (size_t) ncap * w.n_embd * q); + } + ggml_backend_synchronize(backend); + } if (argmax_only) { hooks->argmax_out->resize((size_t) q); - ggml_backend_tensor_get(ex->argmax, hooks->argmax_out->data(), 0, - sizeof(int32_t) * (size_t) q); + if (argmax_src) { + std::memcpy(hooks->argmax_out->data(), argmax_src, sizeof(int32_t) * (size_t) q); + } else { + ggml_backend_tensor_get(ex->argmax, hooks->argmax_out->data(), 0, + sizeof(int32_t) * (size_t) q); + } } else if (hooks->all_logits_out) { hooks->all_logits_out->resize((size_t) w.n_vocab * q); ggml_backend_tensor_get(fg->logits, hooks->all_logits_out->data(), 0, @@ -2146,30 +2622,34 @@ static int ds4_try_fused_verify_step( } else { out_logits.clear(); } - if (hooks->capture_out && ex->capture && ncap > 0) { + if (want_capture) { hooks->capture_out->resize((size_t) ncap * w.n_embd * q); - ggml_backend_tensor_get( - ex->capture, hooks->capture_out->data(), 0, - sizeof(float) * hooks->capture_out->size()); + if (capture_src) { + std::memcpy(hooks->capture_out->data(), capture_src, + sizeof(float) * hooks->capture_out->size()); + } else { + ggml_backend_tensor_get( + ex->capture, hooks->capture_out->data(), 0, + sizeof(float) * hooks->capture_out->size()); + } } if (hooks->boundary_checkpoint_out) { *hooks->boundary_checkpoint_out = ex->boundary_checkpoint; } - if (hooks->window_rows && ex->window_kv && ex->window_score) { - // Layer-major [row, token] -> per layer [token][row bytes]. + if (want_window) { + // Layer-major rows -> per layer, one row tensor per batch token. Ds4VerifyWindowRows & rows = *hooks->window_rows; rows.kv.assign((size_t) w.n_layer, {}); rows.score.assign((size_t) w.n_layer, {}); - std::vector kv_all(ggml_nbytes(ex->window_kv)); - std::vector score_all(ggml_nbytes(ex->window_score)); - ggml_backend_tensor_get(ex->window_kv, kv_all.data(), 0, kv_all.size()); - ggml_backend_tensor_get(ex->window_score, score_all.data(), 0, score_all.size()); - const size_t per_layer = kv_all.size() / ex->window_layers.size(); + rows.kv_dev.assign((size_t) w.n_layer, {}); + rows.score_dev.assign((size_t) w.n_layer, {}); + const size_t per_layer = ex->window_kv_rows.size() / ex->window_layers.size(); for (size_t i = 0; i < ex->window_layers.size(); ++i) { const size_t il = (size_t) ex->window_layers[i]; - rows.kv[il].assign(kv_all.begin() + i * per_layer, kv_all.begin() + (i + 1) * per_layer); - rows.score[il].assign(score_all.begin() + i * per_layer, - score_all.begin() + (i + 1) * per_layer); + rows.kv_dev[il].assign(ex->window_kv_rows.begin() + i * per_layer, + ex->window_kv_rows.begin() + (i + 1) * per_layer); + rows.score_dev[il].assign(ex->window_score_rows.begin() + i * per_layer, + ex->window_score_rows.begin() + (i + 1) * per_layer); } } if (telemetry) { diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 84674dd35..159910345 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -9,6 +9,8 @@ // 6. MoE FFN (hash routing + top-k + shared expert + clamped SwiGLU) #include "deepseek4_internal.h" +#include "deepseek4_cluster.h" +#include "cluster/cluster_config.h" #include "deepseek4_norm.h" #include "deepseek4_image_policy.h" #include "deepseek4_vision.h" @@ -27,6 +29,8 @@ #include "../common/moe_hybrid_routing_stats.h" #include "../common/moe_hybrid_stream.h" #include "../common/moe_hybrid_types.h" +#include "../common/pinned_stage.h" +#include "../common/host_stream_flag.h" #include "ggml.h" #include "ggml-alloc.h" @@ -46,6 +50,8 @@ #include #include #include +#include +#include #include #include @@ -2099,6 +2105,12 @@ struct DeepSeek4MlaLaneBindings { }; HistoryMode history_mode = HistoryMode::ContiguousRing; + // Expert-parallel cluster head split: + // build query heads [head_begin, head_begin + head_count) only; 0 = all. + // The output projection then yields a partial sum over this rank's + // output groups, which the caller all-reduces. + int head_begin = 0; + int head_count = 0; // Raw rows before this position are outside the current window // (DeepSeek4LayerCache::swa_floor). int swa_floor = 0; @@ -2156,6 +2168,22 @@ struct DeepSeek4PreparedProjectedLane { ggml_tensor * compressor_score = nullptr; }; +// Lane batch (LUCE_DS4_LANE_BATCH): a tokenwise verify lane runs its cache +// mutations and its row selection, then hands the attention to the caller, +// which attends for every lane in one op (build_mla_lane_batch_attention) +// and writes the raw ring rows after it. +struct DeepSeek4LaneDeferral { + // In: the caller selects the compressed rows of every lane at once (the + // shared top-k of a V4.1 index source and the layers after it). + bool external_selection = false; + // Out. + bool deferred = false; + ggml_tensor * selection = nullptr; // I32 [rows, 1]: identity or shared top-k + ggml_tensor * comp_source = nullptr; // compressed rows after this lane's update + ggml_tensor * index_comp_source = nullptr; // index keys after this lane's update + int n_comp_attn = 0; +}; + // Per-layer RoPE parameters. Compressed layers use YaRN scaling, and // attn_factor cancels the magnitude scaling rope_yarn applies. struct Ds4RopeParams { @@ -2220,12 +2248,24 @@ static DeepSeek4PreparedProjectedLane build_mla_qkv_projection( ggml_tensor * cur, const DeepSeek4Weights & w, const DeepSeek4Layer & L, - int n_tokens, int projection_columns = 0) { + int n_tokens, int projection_columns = 0, + int head_begin = 0, int head_count = 0) { DeepSeek4PreparedProjectedLane out; ggml_tensor * qr = ds4_mul_mat_columns(ctx, L.attn_q_a, cur, projection_columns); qr = build_rms_norm(ctx, qr, L.attn_q_a_norm, w.rms_eps); - ggml_tensor * q = ds4_mul_mat_columns(ctx, L.attn_q_b, qr, projection_columns); - q = ggml_reshape_3d(ctx, q, w.head_dim, w.n_head, n_tokens); + // Head split: a contiguous range of q_b output rows, so the view stays + // contiguous and a codebook registry keyed by pointer still resolves it. + const int n_head = head_count > 0 ? head_count : w.n_head; + if (head_count == 0 && w.cluster_heads_local > 0) { + GGML_ABORT("deepseek4: this cluster rank holds only its attention heads " + "but a graph attends with all of them"); + } + ggml_tensor * q_b_w = head_count > 0 + ? ggml_view_2d(ctx, L.attn_q_b, L.attn_q_b->ne[0], (int64_t) n_head * w.head_dim, + L.attn_q_b->nb[1], (size_t) head_begin * w.head_dim * L.attn_q_b->nb[1]) + : L.attn_q_b; + ggml_tensor * q = ds4_mul_mat_columns(ctx, q_b_w, qr, projection_columns); + q = ggml_reshape_3d(ctx, q, w.head_dim, n_head, n_tokens); // V4 normalizes every query head to unit RMS after wq_b (model.py: // q *= rsqrt(mean(q^2) + eps)); V4.1 dropped that step. if (w.attn_q_head_norm) { @@ -2251,7 +2291,7 @@ static void build_mla_qkv_rope( bool fuse_q_rope) { if (!fuse_q_rope) { p.q = build_tail_rope_3d(ctx, p.q, rope_pos, w.n_rot, w.head_dim, - w.n_head, n_tokens, rope.freq, rope.scale, + (int) p.q->ne[1], n_tokens, rope.freq, rope.scale, rope.ext, rope.attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope.n_ctx_orig); } @@ -2270,24 +2310,48 @@ static ggml_tensor * build_mla_output_projection( const DeepSeek4Weights & w, const DeepSeek4Layer & L, int n_tokens, - bool allow_grouped, int projection_columns = 0) { - const int group_dim = w.head_dim * (w.n_head / w.n_out_group); + bool allow_grouped, int projection_columns = 0, + int head_begin = 0, int head_count = 0) { + const int heads_per_group = w.n_head / w.n_out_group; + const int group_dim = w.head_dim * heads_per_group; + // Head split: this rank's heads are whole output groups, so the first + // stage takes its groups and the second its columns; the product is a + // partial sum over the groups, summed across ranks by the caller. + const bool head_split = head_count > 0 && head_count != w.n_head; + const int n_groups = head_split ? head_count / heads_per_group : w.n_out_group; + const int group_first = head_split ? head_begin / heads_per_group : 0; attn_out = ggml_reshape_3d( - ctx, attn_out, group_dim, w.n_out_group, n_tokens); + ctx, attn_out, group_dim, n_groups, n_tokens); attn_out = ggml_permute(ctx, attn_out, 0, 2, 1, 3); if (n_tokens == 1) { attn_out = ggml_cont(ctx, attn_out); } + // A cluster rank may hold only its own groups (TargetLoadPlan slices). + const int stored_groups = (int) (L.attn_output_a->ne[1] / w.n_lora_o); + GGML_ASSERT(stored_groups == w.n_out_group || (head_split && stored_groups == n_groups)); ggml_tensor * out_a_3d = ggml_reshape_3d( - ctx, L.attn_output_a, group_dim, w.n_lora_o, w.n_out_group); + ctx, L.attn_output_a, group_dim, w.n_lora_o, stored_groups); + if (head_split && stored_groups != n_groups) { + out_a_3d = ggml_view_3d(ctx, out_a_3d, group_dim, w.n_lora_o, n_groups, + out_a_3d->nb[1], out_a_3d->nb[2], + (size_t) group_first * out_a_3d->nb[2]); + } ggml_tensor * attn_low = ds4_mul_mat_columns(ctx, out_a_3d, attn_out, projection_columns); + // attn_low is group-major over ne[0] (n_lora_o is the fast axis), so a + // contiguous group range is a contiguous column range of wo_b; a sliced + // load already holds just these columns. + ggml_tensor * out_b_w = head_split && stored_groups != n_groups + ? ggml_view_2d(ctx, L.attn_output_b, (int64_t) w.n_lora_o * n_groups, + L.attn_output_b->ne[1], L.attn_output_b->nb[1], + ggml_row_size(L.attn_output_b->type, (int64_t) group_first * w.n_lora_o)) + : L.attn_output_b; // The grouped source layout is read by MMQ's activation quantizer and by // nothing else, so a projection stored unquantized (BF16 attention from a // converter that leaves dense tensors alone) takes the plain path. // MXFP8 (native FP8 dense) has no MMQ, so it also takes the plain path. const bool grouped_output_projection = - allow_grouped && n_tokens > 1 && ggml_is_quantized(L.attn_output_b->type) && + allow_grouped && !head_split && n_tokens > 1 && ggml_is_quantized(L.attn_output_b->type) && L.attn_output_b->type != GGML_TYPE_MXFP8 && !ds4_env_flag("LUCE_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION"); if (grouped_output_projection) { @@ -2295,8 +2359,8 @@ static ggml_tensor * build_mla_output_projection( } attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); attn_low = ggml_reshape_2d( - ctx, attn_low, (int64_t) w.n_lora_o * w.n_out_group, n_tokens); - return ds4_mul_mat_columns(ctx, L.attn_output_b, attn_low, projection_columns); + ctx, attn_low, (int64_t) w.n_lora_o * n_groups, n_tokens); + return ds4_mul_mat_columns(ctx, out_b_w, attn_low, projection_columns); } // One lane's column of a batched prologue. These are views only. @@ -2312,7 +2376,7 @@ static DeepSeek4PreparedProjectedLane ds4_slice_projected_lane( batched.normalized_q_lora->nb[1], (size_t) lane * batched.normalized_q_lora->nb[1]); out.q = ggml_view_3d( - ctx, batched.q, w.head_dim, w.n_head, 1, + ctx, batched.q, w.head_dim, batched.q->ne[1], 1, batched.q->nb[1], batched.q->nb[2], (size_t) lane * batched.q->nb[2]); out.kv = ggml_view_2d( @@ -2330,6 +2394,62 @@ static DeepSeek4PreparedProjectedLane ds4_slice_projected_lane( return out; } +// Segmented K/V is a device-class default (HIP backends), with +// GGML_CUDA_MLA_SEGMENTED_KV=0 as the graph-side kill switch. +static bool ds4_segmented_kv_enabled(const DeepSeek4Weights & w) { + const char * segmented_kv_env = + std::getenv("GGML_CUDA_MLA_SEGMENTED_KV"); + const char * backend_name = w.backend ? ggml_backend_name(w.backend) : nullptr; + const bool segmented_kv_default = backend_name && + (std::strstr(backend_name, "HIP") != nullptr || + std::strstr(backend_name, "ROCm") != nullptr); + return segmented_kv_env + ? segmented_kv_env[0] != '\0' && + std::strcmp(segmented_kv_env, "0") != 0 + : segmented_kv_default; +} + +// LUCE_DS4_LANE_BATCH=1: the single-token lanes of a tokenwise verify layer +// attend in one split-KV flash op (build_mla_lane_batch_attention) instead of +// one op chain per lane plus the concatenation of their contexts. +static bool ds4_lane_batch_enabled() { + static const bool enabled = ds4_env_flag("LUCE_DS4_LANE_BATCH"); + return enabled; +} + +// Whether every single-token lane of a tokenwise verify layer takes the +// split-KV selection flash on the F16 segmented cache (the path +// build_mla_attention_lane_core hands to a DeepSeek4LaneDeferral). The lanes +// must share one padded compressed span so their mask rows line up. +static bool ds4_lane_batch_layer_ok(const DeepSeek4Weights & w, int il, + const DeepSeek4LayerCache & lc, + const DeepSeek4LayerCache & comp_lc, + int padded, const ggml_tensor * identity_rows) { + static const bool f16_kv_lanes = ds4_env_flag("LUCE_DS4_F16_KV_LANES"); + static const bool explicit_split = ds4_env_flag("LUCE_DS4_EXPLICIT_SPLIT"); + static const bool i32_repeat = ds4_env_flag("LUCE_CUDA_I32_REPEAT"); + static const bool no_split_kv = ds4_env_flag("GGML_CUDA_MLA_NO_SPLIT_KV") || + ds4_env_flag("GGML_DS4_FA_NO_SPLIT_KV"); + // Lanes past the sparse threshold leave the Explicit implementation. + static const bool sparse_decode = ds4_env_flag("LUCE_DS4_SPARSE_DECODE_FLASH"); + const int ratio = (int) w.compress_ratios[(size_t) il]; + if (!ds4_lane_batch_enabled() || ratio <= 0 || padded <= 1 || sparse_decode || + !w.shared_index_topk || w.head_dim != 512 || w.n_rot != 64 || + !w.fused_verify_f16_kv || !f16_kv_lanes || !explicit_split || + no_split_kv || !ds4_segmented_kv_enabled(w) || + !lc.raw_kv || lc.raw_kv->type != GGML_TYPE_F16 || + !comp_lc.comp_kv || comp_lc.comp_kv->type != GGML_TYPE_F16) { + return false; + } + // Lanes without a shared selection attend through the identity row list; + // batching it repeats that list on the device (an I32 REPEAT). + if (padded <= w.n_indexer_top_k && + (!identity_rows || padded > identity_rows->ne[0] || !i32_repeat)) { + return false; + } + return true; +} + // `lc` owns the layer's raw ring; `comp_lc` holds the compressed rows it // attends over (ds4_comp_cache: itself on V4 and at every kv source). static DeepSeek4MlaLaneBindings deepseek4_contiguous_lane_bindings( @@ -2377,11 +2497,16 @@ static ggml_tensor * build_mla_attention_lane_core( const DeepSeek4PreparedProjectedLane * prepared = nullptr, ggml_tensor ** out_attn_context = nullptr, DeepSeek4SpecBoundaryCheckpointLayer * boundary_checkpoint = nullptr, - vision::ImageSpanView image_spans = {}) { + vision::ImageSpanView image_spans = {}, + DeepSeek4LaneDeferral * deferral = nullptr) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; - const int n_head = w.n_head; + const int n_head = lane.head_count > 0 ? lane.head_count : w.n_head; + ggml_tensor * attn_sinks_w = (L.attn_sinks && lane.head_count > 0) + ? ggml_view_1d(ctx, L.attn_sinks, n_head, + (size_t) lane.head_begin * ggml_type_size(L.attn_sinks->type)) + : L.attn_sinks; const int n_rot = w.n_rot; const int ratio = w.compress_ratios[layer_idx]; const bool gathered_history = lane.history_mode == @@ -2391,7 +2516,8 @@ static ggml_tensor * build_mla_attention_lane_core( // place. Only gathered paged concurrency supplies a q-wide projection. DeepSeek4PreparedProjectedLane projected; if (!prepared) { - projected = build_mla_qkv_projection(ctx, cur, w, L, n_tokens); + projected = build_mla_qkv_projection(ctx, cur, w, L, n_tokens, 0, + lane.head_begin, lane.head_count); } // ── RoPE on Q and KV (tail rotation on last n_rot dims) ──────── @@ -2532,6 +2658,9 @@ static ggml_tensor * build_mla_attention_lane_core( } if (!lane.write_raw) { // Inactive/padding lanes intentionally have no cache mutation. + } else if (deferral) { + // Lane batch: the caller writes every lane's row after the one + // attention op over all lanes has read the ring. } else if (raw_kv_rows) { ggml_tensor * kv_f32 = ggml_is_contiguous(kv) ? kv : ggml_cont(ctx, kv); raw_kv_source = ggml_set_rows(ctx, lane.raw_kv, kv_f32, raw_kv_rows); @@ -2685,7 +2814,10 @@ static ggml_tensor * build_mla_attention_lane_core( // graph (fused, paged) or in Ds4IndexSelectionStore (host paths). const int n_candidate_rows = w.candidate_source_layer >= 0 ? w.candidate_topk_blocks : 0; ggml_tensor * const carrier_in = lane.index_selection ? *lane.index_selection : nullptr; - if (shared_selection && !selects_topk && carrier_in) { + // A lane-batch caller that selects for all lanes stands in for this + // lane's selection (and for the index source's scoring). + const bool external_selection = deferral && deferral->external_selection; + if (shared_selection && !selects_topk && carrier_in && !external_selection) { // A selection exists only while a query sees more than top_k rows, // the same condition under which the index source built it. if (n_comp_live > w.n_indexer_top_k && carrier_in->ne[1] == n_tokens) { @@ -2696,7 +2828,7 @@ static ggml_tensor * build_mla_attention_lane_core( carrier_in->nb[1], 0)); } } - if (selects_topk) { + if (selects_topk && !external_selection) { int n_index_comp = 0; ggml_tensor * index_visibility_mask = nullptr; if (gathered_history) { @@ -2759,8 +2891,8 @@ static ggml_tensor * build_mla_attention_lane_core( // A short-context verify lane without a selection attends to every // compressed row: give it the identity row list so it runs the split-KV // flash schedule on the F16 cache (LUCE_DS4_EXPLICIT_SPLIT). - if (!indexer_topk && cached_inputs && cached_inputs->identity_rows && masked_kv && - n_tokens == 1 && attention_impl == DeepSeek4AttentionImpl::Explicit && + if (!indexer_topk && !external_selection && cached_inputs && cached_inputs->identity_rows && + masked_kv && n_tokens == 1 && attention_impl == DeepSeek4AttentionImpl::Explicit && w.shared_index_topk && head_dim == 512 && n_rot == 64 && !image_spans.size && !gathered_history && cached_inputs->padded_comp > 1 && cached_inputs->padded_comp <= (int) cached_inputs->identity_rows->ne[0]) { @@ -2768,6 +2900,7 @@ static ggml_tensor * build_mla_attention_lane_core( cached_inputs->padded_comp, 1, (size_t) cached_inputs->padded_comp * sizeof(int32_t), 0); } + const bool has_selection = indexer_topk != nullptr || external_selection; // Maskless indexed prefill admission. This repeats the kernel's // ratio4_causal support check in fattn.cu exactly (indexed-row capacity, // chronological prior window, completed compressed-row frontier): the @@ -2903,30 +3036,21 @@ static ggml_tensor * build_mla_attention_lane_core( // K/V in F16 there. static const bool explicit_split = ds4_env_flag("LUCE_DS4_EXPLICIT_SPLIT"); const bool explicit_selection_f16_kv = explicit_split && fused_explicit_f16_kv && - w.shared_index_topk && indexer_topk && head_dim == 512 && + w.shared_index_topk && has_selection && head_dim == 512 && n_rot == 64 && !image_spans.size; // Segmented K/V is a device-class default (HIP backends), with // GGML_CUDA_MLA_SEGMENTED_KV=0 as the graph-side kill switch. The two // NO_SPLIT_KV flags below are the kernel's own split-KV kill switches // (ds4-env.cuh); repeating them here keeps the graph from handing // segments to a kernel that will not consume them. - const char * segmented_kv_env = - std::getenv("GGML_CUDA_MLA_SEGMENTED_KV"); - const char * backend_name = w.backend ? ggml_backend_name(w.backend) : nullptr; - const bool segmented_kv_default = backend_name && - (std::strstr(backend_name, "HIP") != nullptr || - std::strstr(backend_name, "ROCm") != nullptr); - const bool segmented_kv_enabled = segmented_kv_env - ? segmented_kv_env[0] != '\0' && - std::strcmp(segmented_kv_env, "0") != 0 - : segmented_kv_default; + const bool segmented_kv_enabled = ds4_segmented_kv_enabled(w); // Ratio-4 verification already has an exact mask-derived selected-row // list. Keep raw, compressed, and preserved overwritten rows as separate // dependencies and let the native split-KV kernel address them directly. // This removes two O(context) concatenations per indexed layer/step. const bool segmented_sparse_f16_kv = (fused_sparse_f16_kv || explicit_selection_f16_kv) && - segmented_kv_enabled && indexer_topk && n_tokens <= 8 && + segmented_kv_enabled && has_selection && n_tokens <= 8 && n_comp_attn > 0 && comp_history_source && (old_rows_scratch_f16 || n_comp_attn > 1) && !ds4_env_flag("GGML_CUDA_MLA_NO_SPLIT_KV") && @@ -3027,7 +3151,7 @@ static ggml_tensor * build_mla_attention_lane_core( // reads only the selected compressed rows: the lane's mask keeps the // causal visibility and the selection names the rows (the 2K numerical // bands attach each band's slice below). - const bool selection_flash = w.shared_index_topk && indexer_topk && + const bool selection_flash = w.shared_index_topk && has_selection && head_dim == 512 && n_rot == 64 && !image_spans.size; const bool direct_indexer_topk = ratio == 4 ? indexer_topk && !image_spans.size && @@ -3191,6 +3315,36 @@ static ggml_tensor * build_mla_attention_lane_core( ctx, ggml_cont(ctx, score_mask), indexer_topk, n_raw); } } + if (deferral) { + // Lane batch: the caller admits only layers whose lanes take the + // split-KV selection flash below on the F16 segmented cache, so the + // batched op reproduces each lane's op exactly. Anything else is a + // contract violation; fail the build rather than change arithmetic. + const bool selection_ok = external_selection + ? indexer_topk == nullptr + : indexer_topk && indexer_topk->ne[1] == 1; + const bool batchable = n_tokens == 1 && masked_kv && score_mask && + selection_flash && direct_indexer_topk && !exact_numerical_bands && + !maskless_direct && explicit_selection_f16_kv && segmented_sparse_f16_kv && + segmented_kv_comp && segmented_kv_tail && !old_rows_scratch_f16 && + attention_impl == DeepSeek4AttentionImpl::Explicit && n_raw == w.n_swa && + raw_score_capacity == w.n_swa && comp_history_source && n_comp_attn > 1 && + selection_ok; + if (!batchable) { + std::fprintf(stderr, + "[deepseek4] layer %d: lane batch needs the split-KV selection " + "flash path (selection=%d external=%d segmented=%d f16=%d)\n", + layer_idx, indexer_topk != nullptr, external_selection ? 1 : 0, + segmented_sparse_f16_kv ? 1 : 0, explicit_selection_f16_kv ? 1 : 0); + return nullptr; + } + deferral->deferred = true; + deferral->selection = indexer_topk; + deferral->comp_source = comp_history_source; + deferral->index_comp_source = index_comp_history_source; + deferral->n_comp_attn = n_comp_attn; + return comp_history_source; + } ggml_tensor * context = nullptr; bool inverse_rope_fused = false; // Decode normally keeps the cheaper explicit path. Once the trained @@ -3267,7 +3421,7 @@ static ggml_tensor * build_mla_attention_lane_core( ctx, q_band, k_band, k_band, mask_fa, kq_scale, 0.0f, 0.0f); if (L.attn_sinks) { - ggml_flash_attn_ext_add_sinks(result, L.attn_sinks); + ggml_flash_attn_ext_add_sinks(result, attn_sinks_w); } ggml_flash_attn_ext_set_prec(result, GGML_PREC_F32); ggml_flash_attn_ext_set_ds4_sparse( @@ -3390,7 +3544,7 @@ static ggml_tensor * build_mla_attention_lane_core( context = ggml_flash_attn_ext(ctx, q_fa, k_fa, v_fa, mask_fa, kq_scale, 0.0f, 0.0f); if (L.attn_sinks) { - ggml_flash_attn_ext_add_sinks(context, L.attn_sinks); + ggml_flash_attn_ext_add_sinks(context, attn_sinks_w); } ggml_flash_attn_ext_set_prec(context, GGML_PREC_F32); // Always publish the raw/compressed boundary. A zero keep count leaves @@ -3461,13 +3615,21 @@ static ggml_tensor * build_mla_attention_lane_core( ggml_tensor * probs = nullptr; // Keep long histories on the general softmax path; the fused sink // kernel requires shared memory for every attention column. - if (L.attn_sinks && !score_mask && n_tokens == 1 && n_attn <= 2048) { - // Single-token lanes: scale, sink concat, softmax, and the strided - // view (which forced a 2D staging copy before the PV matmul) fold - // into one launch. The sink is a virtual last column, so every - // reduction matches the concat form bit for bit. - ggml_tensor * sinks = ggml_reshape_1d(ctx, L.attn_sinks, n_head); - probs = ggml_soft_max_ext_sink_col(ctx, scores, sinks, kq_scale); + const bool one_mask_row = score_mask && score_mask->ne[0] == n_attn && + ggml_nrows(score_mask) == 1 && score_mask->nb[0] == ggml_type_size(score_mask->type); + if (L.attn_sinks && (!score_mask || one_mask_row) && n_tokens == 1 && n_attn <= 2048) { + // Single-token lanes: scale, the additive mask, sink concat, + // softmax, and the strided view (which forced a 2D staging copy + // before the PV matmul) fold into one launch. The sink is a + // virtual last column, so every reduction matches the concat form + // bit for bit; with a mask the heads ride on dim 2 so its one row + // broadcasts over them as the add did. + ggml_tensor * sinks = ggml_reshape_1d(ctx, attn_sinks_w, n_head); + probs = score_mask + ? ggml_reshape_2d(ctx, ggml_soft_max_ext_sink_col_mask( + ctx, ggml_reshape_3d(ctx, scores, n_attn, 1, n_head), + score_mask, sinks, kq_scale), n_attn, n_head) + : ggml_soft_max_ext_sink_col(ctx, scores, sinks, kq_scale); } else { scores = ggml_scale(ctx, scores, kq_scale); if (score_mask) { @@ -3483,7 +3645,7 @@ static ggml_tensor * build_mla_attention_lane_core( } } if (L.attn_sinks) { - ggml_tensor * sink_scores = ggml_reshape_2d(ctx, L.attn_sinks, + ggml_tensor * sink_scores = ggml_reshape_2d(ctx, attn_sinks_w, 1, n_head); if (n_tokens > 1) { ggml_tensor * sink_shape = ggml_new_tensor_2d( @@ -3537,7 +3699,183 @@ static ggml_tensor * build_mla_attention_lane_core( } return build_mla_output_projection(ctx, attn_out, w, L, n_tokens, - /*allow_grouped=*/true); + /*allow_grouped=*/true, 0, + lane.head_begin, lane.head_count); +} + +// The lane batch's I32 [rows, lanes] selection, cached across the layers of +// one graph: lanes on the identity row list share one device repeat of it per +// span, and lanes on shared top-k selections concatenate them once per index +// source (the layers after it carry the same tensors). +struct Ds4LaneBatchSelections { + std::vector> identity; + std::vector, ggml_tensor *>> shared; +}; + +static ggml_tensor * ds4_lane_batch_selection( + ggml_context * ctx, + const std::vector & lanes, + ggml_tensor * identity_rows, + Ds4LaneBatchSelections & cache) { + GGML_ASSERT(!lanes.empty() && lanes[0].selection); + const int64_t rows = lanes[0].selection->ne[0]; + bool identity = identity_rows != nullptr; + std::vector key; + key.reserve(lanes.size()); + for (const DeepSeek4LaneDeferral & lane : lanes) { + GGML_ASSERT(lane.selection && lane.selection->type == GGML_TYPE_I32 && + lane.selection->ne[0] == rows && lane.selection->ne[1] == 1); + identity = identity && lane.selection->view_src == identity_rows && + lane.selection->view_offs == 0; + key.push_back(lane.selection); + } + if (identity) { + for (const auto & entry : cache.identity) { + if (entry.first == rows) return entry.second; + } + ggml_tensor * list = ggml_view_2d( + ctx, identity_rows, rows, 1, (size_t) rows * sizeof(int32_t), 0); + ggml_tensor * all = ggml_repeat( + ctx, list, ggml_new_tensor_2d(ctx, GGML_TYPE_I32, rows, (int64_t) lanes.size())); + cache.identity.emplace_back(rows, all); + return all; + } + for (const auto & entry : cache.shared) { + if (entry.first == key) return entry.second; + } + ggml_tensor * all = nullptr; + for (ggml_tensor * selection : key) { + all = all ? ggml_concat(ctx, all, selection, 1) : selection; + } + cache.shared.emplace_back(std::move(key), all); + return all; +} + +// Lane batch over V4.1's shared top-k (external_selection lanes): an index +// source scores and selects the compressed rows of every lane in one indexer +// chain, whose per-token rows (projections, RoPE, QAT, scores, top-k) equal +// each lane's own single-token chain; the layers after it attend over that +// selection. `carrier` holds the latest source's [top_k (+ candidate +// blocks), lanes] selection, as lane_selection does per lane. +static ggml_tensor * ds4_lane_batch_shared_selection( + ggml_context * ctx, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int il, + ggml_tensor * qr, // [n_lora_q, lanes] + ggml_tensor * cur, // [n_embd, lanes] + ggml_tensor * index_comp_source, // F16 index keys after the lanes' updates + int n_comp, // the lanes' padded compressed span + int kv_start, + ggml_tensor * rope_pos, // I32 [lanes] + ggml_tensor * visibility_mask, // F32 [n_comp, lanes] + ggml_tensor ** carrier) { + const int top_k = w.n_indexer_top_k; + const int64_t lanes = qr->ne[1]; + const int ratio = (int) w.compress_ratios[(size_t) il]; + const int n_candidate_rows = w.candidate_source_layer >= 0 ? w.candidate_topk_blocks : 0; + if (deepseek4_is_index_source(w, il)) { + if (!index_comp_source || index_comp_source->type != GGML_TYPE_F16) return nullptr; + const bool uses_candidates = n_candidate_rows > 0 && il >= w.candidate_source_layer; + DeepSeek4IndexCandidates candidates; + candidates.source = uses_candidates && il == w.candidate_source_layer; + if (uses_candidates && !candidates.source && *carrier && + (*carrier)->ne[0] == top_k + n_candidate_rows) { + candidates.blocks = ggml_view_2d( + ctx, *carrier, n_candidate_rows, (*carrier)->ne[1], (*carrier)->nb[1], + (size_t) top_k * sizeof(int32_t)); + } + std::vector bindings; + ggml_tensor * selected = deepseek4_build_indexer_topk( + ctx, qr, cur, w, L, index_comp_source, n_comp, kv_start, (int) lanes, + ratio, rope_pos, visibility_mask, bindings, + uses_candidates ? &candidates : nullptr); + if (!selected || !bindings.empty() || selected->ne[0] != top_k || + selected->ne[1] != lanes) { + return nullptr; + } + *carrier = candidates.blocks + ? ggml_concat(ctx, selected, ggml_cont(ctx, candidates.blocks), 0) + : selected; + return selected; + } + if (!*carrier || (*carrier)->ne[1] != lanes) return nullptr; + return (*carrier)->ne[0] == top_k ? *carrier + : ggml_cont(ctx, ggml_view_2d(ctx, *carrier, top_k, lanes, (*carrier)->nb[1], 0)); +} + +// Attention of every single-token lane of a tokenwise verify layer in one +// split-KV flash op, after each lane ran build_mla_attention_lane_core with a +// DeepSeek4LaneDeferral (its cache updates and row selection, no raw write). +// Lane t must see the raw ring as its own decode step does: the rows lanes +// 0..t wrote and the old contents of the rows later lanes overwrite. The op +// reads the old ring and takes lanes 0..t's rows from `kv` in place +// (ggml_flash_attn_ext_set_ds4_lane_rows); the ring is written after it. +// Mask rows, selections and compressed rows stay per lane, so every lane's +// context is bit-identical to its own op, without the per-lane launches and +// the concatenation of the lane contexts. Returns [head_dim * n_head, lanes] +// before the inverse RoPE, like the lanes' out_attn_context. +static ggml_tensor * build_mla_lane_batch_attention( + ggml_context * ctx, + ggml_cgraph * gf, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + ggml_tensor * raw_kv, // F16 [head_dim, n_swa] ring of this layer + ggml_tensor * q, // [head_dim, n_head, lanes] after RoPE + ggml_tensor * kv, // F32 [head_dim, lanes] after RoPE + ggml_tensor * ring_rows_i32, // [lanes] ring row of each lane + ggml_tensor * ring_rows_i64, // [lanes] the same rows for SET_ROWS + ggml_tensor * mask_f16, // F16 [n_swa + compressed, lanes] + ggml_tensor * selection, // I32 [rows, lanes] + ggml_tensor * comp_source, // compressed rows after the lanes' updates + int n_comp_attn, + int head_begin = 0, // cluster head split: this rank's heads + int head_count = 0) { + const int head_dim = w.head_dim; + const int n_head = head_count > 0 ? head_count : w.n_head; + const int64_t lanes = q->ne[2]; + GGML_ASSERT(q->ne[1] == n_head); + GGML_ASSERT(raw_kv->type == GGML_TYPE_F16 && raw_kv->ne[1] == w.n_swa); + GGML_ASSERT(kv->ne[1] == lanes && selection->ne[1] == lanes); + GGML_ASSERT(mask_f16->type == GGML_TYPE_F16 && mask_f16->ne[1] == lanes && + mask_f16->ne[0] == (int64_t) w.n_swa + n_comp_attn); + GGML_ASSERT(n_comp_attn > 1); + + ggml_tensor * q_fa = ggml_permute(ctx, q, 0, 2, 1, 3); + ggml_tensor * ring = ggml_view_2d(ctx, raw_kv, head_dim, w.n_swa, raw_kv->nb[1], 0); + ggml_tensor * k_fa = ggml_reshape_3d(ctx, ring, head_dim, w.n_swa, 1); + // The lanes' segment split: every compressed row but the last, then the + // last one (as each single-token lane presents them). + ggml_tensor * comp_head = ggml_view_2d( + ctx, comp_source, head_dim, n_comp_attn - 1, comp_source->nb[1], 0); + ggml_tensor * comp_tail = ggml_view_2d( + ctx, comp_source, head_dim, 1, comp_source->nb[1], + (size_t) (n_comp_attn - 1) * comp_source->nb[1]); + + ggml_tensor * context = ggml_flash_attn_ext( + ctx, q_fa, k_fa, k_fa, mask_f16, + 1.0f / sqrtf((float) head_dim), 0.0f, 0.0f); + if (L.attn_sinks) { + ggml_flash_attn_ext_add_sinks(context, head_count > 0 + ? ggml_view_1d(ctx, L.attn_sinks, n_head, + (size_t) head_begin * ggml_type_size(L.attn_sinks->type)) + : L.attn_sinks); + } + ggml_flash_attn_ext_set_prec(context, GGML_PREC_F32); + ggml_flash_attn_ext_set_ds4_sparse( + context, w.n_swa, w.n_swa, -(int) selection->ne[0], 32); + ggml_flash_attn_ext_set_ds4_kv_segments(context, comp_head, comp_tail); + ggml_flash_attn_ext_set_ds4_indexer_topk(context, selection); + // The rows the lanes write, in the cache format the ring holds. + ggml_tensor * lane_kv = ggml_cast(ctx, kv, GGML_TYPE_F16); + ggml_flash_attn_ext_set_ds4_lane_rows(context, ring_rows_i32, lane_kv); + ggml_build_forward_expand(gf, context); + + // Write the ring after the op read its old rows (same stream, graph order). + ggml_tensor * kv_rows = ggml_is_contiguous(kv) ? kv : ggml_cont(ctx, kv); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, raw_kv, kv_rows, ring_rows_i64)); + + return ggml_reshape_2d(ctx, context, (int64_t) head_dim * n_head, lanes); } // Legacy contiguous-cache adapter. Both decode and the consecutive q>1 @@ -3562,10 +3900,13 @@ static ggml_tensor * build_mla_attention( DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit, DeepSeek4SpecBoundaryCheckpointLayer * boundary_checkpoint = nullptr, vision::ImageSpanView image_spans = {}, - ggml_tensor ** index_selection = nullptr) { + ggml_tensor ** index_selection = nullptr, + int head_begin = 0, int head_count = 0) { DeepSeek4MlaLaneBindings lane = deepseek4_contiguous_lane_bindings( w, layer_idx, lc, comp_lc, kv_start + n_tokens - 1); lane.index_selection = index_selection; + lane.head_begin = head_begin; + lane.head_count = head_count; return build_mla_attention_lane_core( ctx, gf, cur, w, L, lane, layer_idx, kv_start, n_tokens, cached_inputs, i32_inputs, i32_array_inputs, i64_array_inputs, @@ -3754,6 +4095,9 @@ struct DeepSeek4PrefillHcPostGraph { ggml_tensor * block_out_cold = nullptr; ggml_tensor * split = nullptr; bool owner_join = false; + // Cluster prefill: the joined hot+cold routed partial is all-reduced in + // the graph before HC-post (the device-join path of a cluster rank). + bool cluster_reduce = false; bool valid() const { return owner_ctx && backend && n_tokens > 0 && sg.ctx && sg.gf && @@ -3772,6 +4116,7 @@ struct DeepSeek4PrefillHcPostGraph { block_out_cold = nullptr; split = nullptr; owner_join = false; + cluster_reduce = false; } }; @@ -4351,6 +4696,65 @@ static ggml_tensor * build_shared_ffn( return ggml_mul_mat(ctx, L.ffn_down_shexp, mid_sh); } +// Cluster shared-expert sharding: build only the intermediate range +// [ff_begin, ff_begin + ff_count) of the shared expert. The result is a +// partial sum of the shared expert over that range, because the down +// projection contracts over the intermediate axis, so the caller adds it to +// the routed partial before the all-reduce and the reduction sums both in one +// message. +// +// gate and up are sliced by output row, a contiguous view. When the range +// falls on whole quant blocks of `down`, `down` is taken as a column view (row +// stride unchanged) so only the slice's weights are read. Otherwise (the +// drafter's shared expert) the intermediate is zero-padded back to full width +// and the whole `down` is applied; the zeros contribute nothing. +static ggml_tensor * build_shared_ffn_slice( + ggml_context * ctx, + ggml_tensor * cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int ff_begin, + int ff_count) { + if (ff_count <= 0 || !L.ffn_gate_shexp || !L.ffn_up_shexp || !L.ffn_down_shexp) { + return build_shared_ffn(ctx, cur, w, L); + } + const int64_t n_ff = L.ffn_gate_shexp->ne[1]; + if (ff_begin < 0 || (int64_t) ff_begin + ff_count > n_ff) { + return build_shared_ffn(ctx, cur, w, L); + } + auto row_slice = [&](ggml_tensor * t) { + return ggml_view_2d(ctx, t, t->ne[0], (int64_t) ff_count, t->nb[1], + (size_t) ff_begin * t->nb[1]); + }; + ggml_tensor * gate_sh = ggml_mul_mat(ctx, row_slice(L.ffn_gate_shexp), cur); + ggml_tensor * up_sh = ggml_mul_mat(ctx, row_slice(L.ffn_up_shexp), cur); + ggml_tensor * mid_sh = build_clamped_swiglu(ctx, gate_sh, up_sh, w.swiglu_clamp_exp); + if ((int64_t) ff_count == n_ff) { + // The layer holds only this slice (a sliced load). + return ggml_mul_mat(ctx, L.ffn_down_shexp, mid_sh); + } + const int64_t blck = ggml_blck_size(L.ffn_down_shexp->type); + if (ff_begin % blck == 0 && ff_count % blck == 0) { + // On the RDNA4 Q8_0 matvec the slice's blocks keep the lanes they take + // in the whole rows, so the product equals the whole rows times the + // intermediate zero-padded outside the slice; other kernels sum the + // slice in their own order. + ggml_tensor * down = ggml_view_2d(ctx, L.ffn_down_shexp, ff_count, L.ffn_down_shexp->ne[1], + L.ffn_down_shexp->nb[1], + ggml_row_size(L.ffn_down_shexp->type, ff_begin)); + ggml_tensor * out = ggml_mul_mat(ctx, down, mid_sh); + ggml_backend_cuda_mul_mat_whole_row_lanes(out); + return out; + } + ggml_tensor * zeros = ggml_scale(ctx, mid_sh, 0.0f); + ggml_tensor * padded = nullptr; + for (int64_t off = 0; off < n_ff; off += ff_count) { + ggml_tensor * piece = (off == (int64_t) ff_begin) ? mid_sh : zeros; + padded = padded ? ggml_concat(ctx, padded, piece, 0) : piece; + } + return ggml_mul_mat(ctx, L.ffn_down_shexp, padded); +} + bool init_deepseek4_streamed_expert_cache( const DeepSeek4Weights & w, const MoeHybridStorage & hybrid, @@ -6164,6 +6568,18 @@ bool deepseek4_step( Ds4VerifyHooks * verify_hooks, MoeExpertComputeRuntime * expert_runtime, bool need_logits) { + if (cache.cluster_rt && moe_hybrid != nullptr) { + // Cluster rank: the legacy CPU-HC hybrid step has no all-reduce + // insertion point. Route through the per-layer forward, which reduces + // the routed partial inside eval_ds4_layer_range_hybrid_ffn. + std::vector hc_state; + (void) stream_engine; (void) need_logits; + return deepseek4_step_layer_range( + backend, device, w, cache, hc_state, embed, n_tokens, kv_start, + 0, w.n_layer, &out_logits, token_ids, telemetry, + /*allow_decode_graph_reuse=*/verify_hooks == nullptr, verify_hooks, + moe_hybrid, expert_runtime, routing_stats); + } if (w.moe_hybrid && moe_hybrid != nullptr) { if (!deepseek4_cuda_hc_set_device(device)) { std::fprintf(stderr, @@ -6356,6 +6772,131 @@ struct DeepSeek4FusedDecodeCache { // Fused verification graphs retain allocators, schedulers, peer events, and // model tensor pointers. Keep them under DeepSeek4Cache ownership so park, // reload, and shutdown destroy them before either GPU backend is released. +// Reads a verify step's Engram rows on a worker thread while the GPU already +// runs the graph. The keys reach the device through a pinned buffer whose +// upload waits on the stream for a word the worker raises once the rows are +// decoded, so the host never blocks on the table reads (preads that mostly +// miss the page cache, on the rank that holds the tables, while every other +// rank waits for its broadcast). The first consumer is the +// Engram broadcast after layer 0's attention, which hides the read. +class Ds4EngramAsync { +public: + Ds4EngramAsync() = default; + Ds4EngramAsync(const Ds4EngramAsync &) = delete; + Ds4EngramAsync & operator=(const Ds4EngramAsync &) = delete; + ~Ds4EngramAsync() { stop(); } + + // Starts the read of `count` tokens at `first_pos` into the pinned keys + // (`floats` floats). False: the async path is unavailable, read in line. + bool start(ggml_backend_t backend, const DeepSeek4Weights & w, DeepSeek4EngramTokens & ctx, + const int32_t * tokens, int first_pos, int count, size_t floats) { + if (!ggml_backend_is_cuda(backend) || !ggml_backend_cuda_get_stream(backend)) return false; + if (!flag_.ready() && !flag_.init()) return false; + if (!keys_.reserve(backend, floats * sizeof(float))) return false; + keys_.reset(); + float * dst = static_cast(keys_.take(floats * sizeof(float))); + if (!dst) return false; + if (!thread_.joinable()) thread_ = std::thread([this] { run(); }); + { + std::lock_guard lock(mutex_); + job_w_ = &w; + job_ctx_ = &ctx; + job_tokens_.assign(tokens, tokens + count); + job_first_ = first_pos; + job_count_ = count; + job_dst_ = dst; + job_floats_ = floats; + ++seq_; + pending_ = true; + done_ = false; + } + cv_.notify_all(); + return true; + } + + // Queues the keys upload on `backend`'s stream behind the worker's word. + bool enqueue(ggml_backend_t backend, ggml_tensor * keys) { + if (!keys || !keys->buffer || ggml_nbytes(keys) < job_floats_ * sizeof(float)) return false; + if (!flag_.enqueue_wait(ggml_backend_cuda_get_stream(backend), seq_)) return false; + ggml_backend_tensor_set_async(backend, keys, job_dst_, 0, job_floats_ * sizeof(float)); + return true; + } + + // The same upload into a raw device pointer on `stream` (the Engram + // broadcast's buffer, queued from its graph callback). + bool enqueue_raw(void * stream, void * dst, size_t bytes) { + if (!stream || !dst || bytes > job_floats_ * sizeof(float)) return false; + return flag_.enqueue_wait(stream, seq_) && + HostStreamFlag::copy_to_device(stream, dst, job_dst_, bytes); + } + + // Waits for the worker. False when the table read failed (the word was + // raised anyway so the stream never hangs, but the step is invalid). + bool finish(uint64_t * read_us) { + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { return done_; }); + if (read_us) *read_us += read_us_; + if (!ok_) std::fprintf(stderr, "[ds4-fused-verify] async engram read: %s\n", err_.c_str()); + return ok_; + } + +private: + void run() { + std::unique_lock lock(mutex_); + for (;;) { + cv_.wait(lock, [this] { return pending_ || stopping_; }); + if (stopping_) return; + pending_ = false; + const uint32_t seq = seq_; + lock.unlock(); + const auto t0 = std::chrono::steady_clock::now(); + std::string err; + const DeepSeek4EngramRuntime * engram = job_w_->engram_runtime.get(); + const bool ok = engram && engram->prepare(*job_ctx_, job_tokens_.data(), job_first_, + (size_t) job_count_, job_dst_, &err); + const uint64_t us = (uint64_t) std::chrono::duration_cast( + std::chrono::steady_clock::now() - t0).count(); + flag_.signal(seq); + lock.lock(); + ok_ = ok; + err_ = engram ? err : std::string("no Engram runtime"); + read_us_ = us; + done_ = true; + cv_.notify_all(); + } + } + + void stop() { + if (!thread_.joinable()) return; + { + std::lock_guard lock(mutex_); + stopping_ = true; + } + cv_.notify_all(); + thread_.join(); + } + + HostStreamFlag flag_; + PinnedStage keys_; + std::thread thread_; + std::mutex mutex_; + std::condition_variable cv_; + bool pending_ = false; + bool done_ = true; + bool stopping_ = false; + bool ok_ = true; + std::string err_; + uint64_t read_us_ = 0; + uint32_t seq_ = 0; + const DeepSeek4Weights * job_w_ = nullptr; + DeepSeek4EngramTokens * job_ctx_ = nullptr; + std::vector job_tokens_; + int job_first_ = 0; + int job_count_ = 0; + float * job_dst_ = nullptr; + size_t job_floats_ = 0; +}; + struct Ds4FusedVerifyCache { // Resident verifier shapes for the adaptive-width profile: four widths // (q2..q5) across the four ratio-4 phases plus the q5 two-boundary and @@ -6401,6 +6942,7 @@ struct Ds4FusedVerifyCache { ggml_tensor * engram_keys = nullptr; // f32 [cols*key_len, q, n_engram_layers] ggml_tensor * rawrows = nullptr; // i64 [1,q] ggml_tensor * saved_rawrows = nullptr; // i32 [q], gather before ring writes + ggml_tensor * pos_rows = nullptr; // i64 [q], kv_start + i: a ratio-1 lane batch's rows // Per distinct compress ratio: each token's APE row (i32 [q]) and // compressor state row (i64 [1,q]). struct RatioRows { @@ -6417,13 +6959,14 @@ struct Ds4FusedVerifyCache { } ggml_tensor * capture = nullptr; // f32 [n_embd*ncap,q], token-major ggml_tensor * argmax = nullptr; // i32 [q], optional greedy output + bool vocab_split = false; // built with the cluster vocabulary split head DeepSeek4SpecBoundaryCheckpoint boundary_checkpoint; // Tokenwise verify: the compressor window row each token wrote, per // window-state layer (window_layers), [row, layer-major token] for // kv and score, read back as Ds4VerifyWindowRows. std::vector window_layers; - ggml_tensor * window_kv = nullptr; - ggml_tensor * window_score = nullptr; + std::vector window_kv_rows; // layer-major, one per token + std::vector window_score_rows; // Reused host staging for the context-sized additive attention mask. // Keeping it per slot removes allocation churn in both full and // sparse-range mask update modes. @@ -6443,6 +6986,10 @@ struct Ds4FusedVerifyCache { void reset() { *this = Extra{}; } }; std::array extra; + // Pinned staging for a step's uploads and readbacks, and the Engram + // reader (both shared by the slots: one step runs at a time). + PinnedStage stage; + std::unique_ptr engram_async; void destroy() { for (auto & slot : slots) slot.destroy(backend, peer_backend); @@ -6996,9 +7543,12 @@ static bool build_prefill_hc_post_graph( ggml_backend_t backend, const DeepSeek4Weights & w, int n_tokens, - bool owner_join = false) { + bool owner_join = false, + Ds4ClusterRuntime * cluster_rt = nullptr) { + const bool cluster_reduce = owner_join && cluster_rt && cluster_rt->size() > 1; if (out.valid() && out.owner_ctx == w.ctx && out.backend == backend && - out.n_tokens == n_tokens && out.owner_join == owner_join) { + out.n_tokens == n_tokens && out.owner_join == owner_join && + out.cluster_reduce == cluster_reduce) { return true; } out.free(); @@ -7032,6 +7582,23 @@ static bool build_prefill_hc_post_graph( ggml_tensor * hc_block_out = out.block_out_cold ? ggml_add(out.sg.ctx, out.block_out, out.block_out_cold) : out.block_out; + if (cluster_reduce) { + // The bf16 prefill wire needs a device scratch of n elements. + std::string scratch_err; + if (!cluster_rt->ensure_scratch(backend, (size_t) w.n_embd * (size_t) n_tokens, &scratch_err)) { + std::fprintf(stderr, "[deepseek4-cluster] prefill scratch failed: %s\n", scratch_err.c_str()); + out.free(); + return false; + } + // Every rank holds the routed partial of the experts it owns (rank 0's + // also carries the shared expert): one in-graph all-reduce makes the + // FFN output whole before HC-post, on the backend stream. + hc_block_out = ds4_cluster_allreduce_node(out.sg.ctx, ggml_cont(out.sg.ctx, hc_block_out), *cluster_rt, /*inplace=*/true); + if (!hc_block_out) { + out.free(); + return false; + } + } out.sg.hidden_states = ggml_ds4_hc_post( out.sg.ctx, out.residual_hc, hc_block_out, out.split, w.n_hc); out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 1024, false); @@ -7048,6 +7615,7 @@ static bool build_prefill_hc_post_graph( out.backend = backend; out.n_tokens = n_tokens; out.owner_join = owner_join; + out.cluster_reduce = cluster_reduce; return true; } @@ -7471,7 +8039,11 @@ static bool eval_ds4_layer_range_hybrid_ffn( std::vector & out, DeepSeek4StepTelemetry * telemetry, const MoeHybridDeviceOutputs * device_outputs = nullptr, - int kv_start = 0, vision::ImageSpanView image_spans = {}) { + int kv_start = 0, vision::ImageSpanView image_spans = {}, + // Cluster expert-parallel (per-layer path): mask routes to experts this rank + // does not evaluate, all-reduce the routed partial across ranks, then + // add the locally computed shared expert. nullptr = single box. + Ds4ClusterRuntime * cluster_rt = nullptr) { const bool trace_prefill = ds4_env_flag("LUCE_DS4_PREFILL_TRACE"); if (trace_prefill) { std::fprintf(stderr, @@ -7498,8 +8070,13 @@ static bool eval_ds4_layer_range_hybrid_ffn( const bool device_input_enabled = !device_input_env || !*device_input_env || std::strcmp(device_input_env, "0") != 0; + // A cluster rank on the host path adds the shared expert on the host after + // the reduction, which needs the normalized activation back on the host. + // With the device join the reduction is a graph node in HC-post instead. + const bool cluster_device_join = + cluster_rt && cluster_rt->size() > 1 && device_outputs && device_outputs->valid(); const bool device_ffn_input = - device_input_enabled && layer_storage.n_streamed == 0 && + (!cluster_rt || cluster_device_join) && device_input_enabled && layer_storage.n_streamed == 0 && !expert_compute && moe_expert_major_prefill_enabled(n_tokens) && layer_storage.cold_backend_kind == MoeHybridColdBackend::Gpu && layer_storage.cold_backend && layer_storage.cold_backend != backend && @@ -7801,6 +8378,33 @@ static bool eval_ds4_layer_range_hybrid_ffn( MoeHybridConfig cfg = make_ds4_moe_hybrid_config(w); cfg.n_expert_used = route_width; MoeLayerDesc desc = make_ds4_moe_layer_desc(L); + // Cluster: the shared expert is replicated on every rank and must not be + // part of the reduced partial. Evaluate the routed experts with the shexp + // tensors cleared (per-layer equivalent of include_shared=false) and add + // the full-desc shared term after the all-reduce below. Routes to + // experts another rank evaluates become id -1 / weight 0, which the + // hybrid evaluators skip (MoeHybridLayerStorage::foreign_routes). + // Slot = token index inside this batch (identical on all ranks). + const MoeLayerDesc shared_desc = desc; + if (cluster_rt) { + ds4_cluster_mask_routes(*cluster_rt, layer, selected.data(), weights.data(), + route_width, n_tokens); + // Device join: rank 0's hot owner carries the shared expert, so the + // in-graph reduction counts it exactly once. + // A sliced shared expert is a partial on every rank instead. + if (!(cluster_device_join && + (cluster_rt->rank() == 0 || w.cluster_shexp_ff_local > 0))) { + if (w.cluster_shexp_ff_local > 0) { + std::fprintf(stderr, "[deepseek4-cluster] layer %d: a sliced shared expert " + "needs the device join\n", layer); + return false; + } + desc.ffn_gate_shexp = nullptr; + desc.ffn_up_shexp = nullptr; + desc.ffn_down_shexp = nullptr; + desc.ffn_gate_inp_shexp = nullptr; + } + } ggml_gallocr_t * hot_alloc = persistent_owner_alloc ? &hybrid.prefill_hot_alloc : nullptr; ggml_gallocr_t * cold_alloc = persistent_owner_alloc @@ -7811,7 +8415,22 @@ static bool eval_ds4_layer_range_hybrid_ffn( layer); } const auto owners_t0 = Ds4TimingClock::now(); - const bool ok = eval_ds4_hybrid( + // Cluster: a rank can end up with no active route in a layer (every one + // of the 6 routes owned elsewhere). No evaluator is entered then; the + // partial is zero and only the all-reduce and the shared expert below + // contribute. + int cluster_active_routes = -1; + if (cluster_rt) { + cluster_active_routes = 0; + for (int32_t id : selected) { + if (id >= 0) ++cluster_active_routes; + } + } + const bool skip_local_experts = cluster_rt && !cluster_device_join && cluster_active_routes == 0; + if (skip_local_experts) { + out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); + } + const bool ok = skip_local_experts || eval_ds4_hybrid( backend, hybrid.cpu_backend, cfg, desc, &hybrid, hybrid.layers[(size_t)layer], nullptr, layer, n_embd, route_width, @@ -7829,7 +8448,50 @@ static bool eval_ds4_layer_range_hybrid_ffn( layer, ok ? "ok" : "failed", ds4_elapsed_us(owners_t0, Ds4TimingClock::now()) / 1000.0); } - return ok; + if (!ok || !cluster_rt || cluster_device_join) return ok; + + // ── Cluster (per-layer path): reduce the routed partial, then add the shared expert ── + // `out` is the host copy of this rank's routed partial [n_embd, n_tokens] + // (a cluster rank never takes the device_ffn_input / device-join + // branches). The all-reduce stages it through the runtime's device + // scratch on the backend stream and waits, so the sum is identical on + // every rank before HC-post consumes it. + if (out.size() != (size_t) n_embd * (size_t) n_tokens) { + std::fprintf(stderr, + "[deepseek4-cluster] layer %d: routed partial has %zu floats, expected %zu\n", + layer, out.size(), (size_t) n_embd * (size_t) n_tokens); + return false; + } + std::string cluster_err; + if (!ds4_cluster_allreduce_layer_host( + *cluster_rt, backend, out.data(), + n_embd, n_tokens, layer, telemetry, &cluster_err)) { + std::fprintf(stderr, "[deepseek4-cluster] layer %d all-reduce failed: %s\n", + layer, cluster_err.c_str()); + return false; + } + if (shared_desc.has_shared_expert()) { + if (normed_host.empty()) { + std::fprintf(stderr, + "[deepseek4-cluster] layer %d: shared expert needs the host " + "normalized activation\n", layer); + return false; + } + const auto shared_t0 = Ds4TimingClock::now(); + std::vector shared_out; + if (!eval_moe_shared_expert_batched( + backend, cfg, shared_desc, hybrid.layers[(size_t)layer], + normed_host.data(), n_tokens, shared_out, &cluster_err)) { + std::fprintf(stderr, "[deepseek4-cluster] layer %d shared expert failed: %s\n", + layer, cluster_err.c_str()); + return false; + } + for (size_t i = 0; i < out.size(); ++i) out[i] += shared_out[i]; + if (telemetry) { + telemetry->ffn_eval_us += ds4_elapsed_us(shared_t0, Ds4TimingClock::now()); + } + } + return true; } // Exact-order prefill control: retain the layer-major HC/FFN schedule, but run @@ -8137,6 +8799,10 @@ static bool ds4_run_verify_attention( window_rows->kv.resize((size_t) il + 1); window_rows->score.resize((size_t) il + 1); } + if (window_rows->kv_dev.size() > (size_t) il) { + window_rows->kv_dev[(size_t) il].clear(); + window_rows->score_dev[(size_t) il].clear(); + } const size_t row_bytes = ggml_nbytes(window_kv[0]); window_rows->kv[(size_t) il].resize(row_bytes * (size_t) n_tokens); window_rows->score[(size_t) il].resize(row_bytes * (size_t) n_tokens); @@ -9531,7 +10197,7 @@ bool deepseek4_paged_gathered_step( mc, *fg, *ex, backend, w, nullptr, rt->model.hc_layer_weights, rt->model.hc_output_weights, rt->model.hash_routing_tables, 0, (int) lanes, - token_ids != nullptr, {}, hybrid, std::move(key), + token_ids != nullptr, {}, hybrid, /*vocab_split=*/false, std::move(key), &cache, &prepared, bucket_history)) { std::fprintf(stderr, "[deepseek4-paged] failed to build gathered graph " @@ -9996,10 +10662,13 @@ bool deepseek4_step_layer_range( const auto step_t0 = Ds4TimingClock::now(); // This band's HC graph slot may still hold an earlier band's pending // join: complete it before the slot's graphs are reused. - if (cache.layer_major_band && cache.layer_major_band->pipeline && - !ds4_pipeline_finish(cache.layer_major_band->pipeline->slots[cache.layer_major_band->graph_slot], - backend)) { - return false; + if (cache.layer_major_band && cache.layer_major_band->pipeline) { + const auto join_t0 = Ds4TimingClock::now(); + if (!ds4_pipeline_finish(cache.layer_major_band->pipeline->slots[cache.layer_major_band->graph_slot], + backend)) { + return false; + } + if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(join_t0, Ds4TimingClock::now()); } bool image_batch = false; @@ -10027,11 +10696,18 @@ bool deepseek4_step_layer_range( const int n_hc = w.n_hc; const int hc_dim = n_hc * n_embd; const bool is_last_shard = (layer_end >= w.n_layer); + // Expert-parallel cluster: the whole-model fused verify/decode graph + // carries one in-graph all-reduce per MoE layer once the communicator is + // attached; otherwise the per-layer loop below reduces host-side. + const bool cluster_mode = cache.cluster_rt != nullptr; + const bool cluster_fused_ok = + !cluster_mode || ds4_cluster_fused_graph_available(cache.cluster_rt); const bool fused_hybrid_ready = moe_hybrid && !expert_runtime && moe_hybrid->materialized_cold_experts && moe_hybrid->cold_backend_kind == MoeHybridColdBackend::Gpu && - moe_hybrid->cold_backend && moe_hybrid->cold_backend != backend; + moe_hybrid->cold_backend && moe_hybrid->cold_backend != backend && + cluster_fused_ok; const bool wide_verify_candidate = n_tokens == DS4_Q5_VERIFY_TOKENS && ds4_env_flag("LUCE_DS4_Q5_VERIFY"); @@ -10450,6 +11126,12 @@ bool deepseek4_step_layer_range( scratch.hash_expert_ids, embed, n_tokens, kv_start, *out_logits, token_ids, fused_graph_hooks, telemetry, fused_hybrid_ready ? moe_hybrid : nullptr, routing_stats); + if (cluster_mode && !cache.cluster_rt->node_error.empty()) { + std::fprintf(stderr, "[deepseek4-cluster] in-graph collective failed: %s\n", + cache.cluster_rt->node_error.c_str()); + cache.cluster_rt->node_error.clear(); + return false; + } if (vrc < 0) return false; if (vrc > 0) { const int np = kv_start + n_tokens; @@ -10532,8 +11214,25 @@ bool deepseek4_step_layer_range( } if (band_engram_keys) *band_engram_keys = std::move(engram_keys_local); } - const std::vector & engram_keys = + std::vector & engram_keys_mut = band_engram_keys && range_has_engram ? *band_engram_keys : engram_keys_local; + // The band's keys hold every Engram layer's rows: one broadcast per band. + uint8_t * keys_shared = band_engram_keys && range_has_engram && layer_major_band + ? layer_major_band->engram_keys_shared : nullptr; + if (cache.cluster_rt && range_has_engram && !(keys_shared && *keys_shared)) { + // Engram tables live on rank 0: its rows replace every other rank's. + const auto bcast_t0 = Ds4TimingClock::now(); + std::string bcast_err; + if (!ds4_cluster_broadcast_host(*cache.cluster_rt, backend, engram_keys_mut.data(), + engram_keys_mut.size(), &bcast_err)) { + std::fprintf(stderr, "[deepseek4-cluster] Engram row broadcast failed: %s\n", + bcast_err.c_str()); + return false; + } + if (keys_shared) *keys_shared = 1; + if (telemetry) telemetry->engram_read_us += ds4_elapsed_us(bcast_t0, Ds4TimingClock::now()); + } + const std::vector & engram_keys = engram_keys_mut; if (!layer_range_cache.index_selection.ensure( backend, w, layer_major_band ? layer_major_band->selection_columns : n_tokens)) { std::fprintf(stderr, "[deepseek4] index selection store allocation failed\n"); @@ -10617,7 +11316,7 @@ bool deepseek4_step_layer_range( prefill_hc_post_graph, backend, w, n_tokens) || (moe_hybrid && !build_prefill_hc_post_graph( prefill_moe_hc_post_graph, backend, w, n_tokens, - /*owner_join=*/true))) { + /*owner_join=*/true, cache.cluster_rt))) { std::fprintf(stderr, "[deepseek4-prefill] batched GPU HC initialization failed\n"); return false; @@ -11063,6 +11762,13 @@ bool deepseek4_step_layer_range( ggml_tensor * store_view = layer_range_cache.index_selection.view( ctx, layer_major_band ? layer_major_band->selection_first : 0, n_tokens); ggml_tensor * index_selection = store_view; + // Cluster head split in the uncached graphs + // (prefill bands, short first decodes): each rank holds and + // attends with its own output groups, and one all-reduce sums + // the partials. + Ds4ClusterRuntime * split_rt = cache.cluster_rt; + const bool split_heads = split_rt && split_rt->size() > 1 && + split_rt->attn_head_count > 0 && w.cluster_heads_local > 0; attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, comp_lc, il, kv_start, n_tokens, nullptr, i32_inputs, i32_array_inputs, @@ -11071,8 +11777,22 @@ bool deepseek4_step_layer_range( attention_impl, /*boundary_checkpoint=*/nullptr, image_batch ? image_spans : vision::ImageSpanView{}, - &index_selection); + &index_selection, + split_heads ? split_rt->attn_head_begin : 0, + split_heads ? split_rt->attn_head_count : 0); if (!attn_out) { ggml_free(ctx); return false; } + if (split_heads) { + std::string scratch_err; + if (!split_rt->ensure_scratch(backend, (size_t) n_embd * (size_t) n_tokens, + &scratch_err)) { + std::fprintf(stderr, "[deepseek4-cluster] prefill attention scratch: %s\n", + scratch_err.c_str()); + ggml_free(ctx); + return false; + } + attn_out = ds4_cluster_allreduce_node(ctx, ggml_cont(ctx, attn_out), *split_rt, /*inplace=*/true); + if (!attn_out) { ggml_free(ctx); return false; } + } ds4_publish_index_selection(ctx, gf, w, il, store_view, index_selection); ggml_set_output(attn_out); ggml_build_forward_expand(gf, attn_out); @@ -11596,7 +12316,8 @@ bool deepseek4_step_layer_range( *moe_hybrid, expert_runtime, routing_stats, ffn_out_host, telemetry, ffn_device_join ? &owner_outputs : nullptr, - kv_start, image_batch ? image_spans : vision::ImageSpanView{}); + kv_start, image_batch ? image_spans : vision::ImageSpanView{}, + cache.cluster_rt); if (!ffn_eval_ok) { if (pipe) ds4_pipeline_finish_all(*pipe, backend, false); std::fprintf(stderr, @@ -11997,17 +12718,34 @@ bool deepseek4_prefill_layer_major( // turn. const int pipeline_bands = deepseek4_prefill_pipeline_bands(); const bool pipeline_requested = pipeline_bands > 0; - const std::vector run_bands = pipeline_requested + // A cluster rank pipelines only when every rank does: every band is a + // collective, so the ranks agree on pipelining (the parts depend on the + // agreed chunks alone, so they match). Without cluster pipelining a rank + // keeps the caller's bands whole. + const bool cluster_multi = cache.cluster_rt && cache.cluster_rt->size() > 1; + const bool cluster_pipeline = !cluster_multi || ds4_cluster_prefill_pipeline_enabled(bands); + const std::vector run_bands = pipeline_requested && cluster_pipeline ? deepseek4_pipeline_parts(bands, pipeline_bands) : bands; Ds4PrefillPipeline pipeline_state; - const bool pipelined = pipeline_requested && device_residual && !cache.pipeline_off; + bool pipelined = pipeline_requested && cluster_pipeline && device_residual && !cache.pipeline_off; + if (cluster_multi && pipeline_requested && cluster_pipeline) { + std::vector flags; + std::string agree_err; + if (!ds4_cluster_allgather_i32(*cache.cluster_rt, backend, pipelined ? 1 : 0, flags, &agree_err)) { + std::fprintf(stderr, "[deepseek4-cluster] pipeline agreement failed: %s\n", agree_err.c_str()); + return false; + } + for (int32_t v : flags) if (v <= 0) pipelined = false; + if (flags.empty()) pipelined = false; + } { static bool logged = false; if (pipeline_requested && !logged) { logged = true; std::fprintf(stderr, "[deepseek4] layer-major prefill pipeline %s (%zu bands)\n", pipelined ? "active" + : !cluster_pipeline ? "inactive (cluster bands too short)" : !device_residual ? "inactive (needs the device residual)" : "inactive (off after a failed prefill)", run_bands.size()); @@ -12031,6 +12769,7 @@ bool deepseek4_prefill_layer_major( // pass, and reused by the band's second Engram layer. Ds4EngramReadAhead engram_ahead; std::vector> & band_keys = engram_ahead.keys; + std::vector band_keys_shared(run_bands.size(), 0); if (ds4_engram_on_device() && w.engram_runtime && token_ids) { engram_ahead.start(w, context_before, token_ids, kv_start, run_bands); } @@ -12056,6 +12795,7 @@ bool deepseek4_prefill_layer_major( band.staggered_pre = &band_pre; band.selection_first = first; band.engram_keys = band_keys.empty() ? nullptr : &band_keys[bi]; + band.engram_keys_shared = band_keys.empty() ? nullptr : &band_keys_shared[bi]; cache.layer_major_band = &band; const float * input = il == 0 ? embed + (size_t) first * n_embd : device_residual ? nullptr @@ -12388,9 +13128,14 @@ static ggml_tensor * build_dspark_attention( ggml_tensor * neg_block, // I32[block] -(block positions) ggml_tensor * pos_ctx, // I32[ctx_len] absolute positions committed-ctx_len..committed-1 ggml_tensor * attn_mask, // F32[ctx_len+block], 0 or -inf - bool flash) { // D=512 flash kernel instead of BLAS products + bool flash, // D=512 flash kernel instead of BLAS products + ggml_tensor ** flash_mask_cache = nullptr, // the F16 flash mask, shared by the layers + int head_begin = 0, // cluster split: this rank's heads; the result is + int head_count = 0) { // then a partial sum over its output groups const int head_dim = w.head_dim; - const int n_head = w.n_head; + const int n_head_all = w.n_head; + const bool head_split = head_count > 0 && head_count != n_head_all; + const int n_head = head_split ? head_count : n_head_all; const int n_rot = w.n_rot; const int n_lora_o = w.n_lora_o; const int n_out_group = w.n_out_group; @@ -12403,7 +13148,14 @@ static ggml_tensor * build_dspark_attention( // ── Q path (block queries) ────────────────────────────────────── ggml_tensor * qr = build_rms_norm(ctx, ggml_mul_mat(ctx, L.attn_q_a, cur), L.attn_q_a_norm, eps); - ggml_tensor * q = ggml_mul_mat(ctx, L.attn_q_b, qr); // [n_head*head_dim, block] + ggml_tensor * q_b_w = head_split + ? ggml_view_2d(ctx, L.attn_q_b, L.attn_q_b->ne[0], (int64_t) n_head * head_dim, + L.attn_q_b->nb[1], (size_t) head_begin * head_dim * L.attn_q_b->nb[1]) + : L.attn_q_b; + ggml_tensor * sinks = L.attn_sinks && head_split + ? ggml_view_1d(ctx, L.attn_sinks, n_head, (size_t) head_begin * ggml_element_size(L.attn_sinks)) + : L.attn_sinks; + ggml_tensor * q = ggml_mul_mat(ctx, q_b_w, qr); // [n_head*head_dim, block] q = ggml_reshape_3d(ctx, q, head_dim, n_head, block); if (w.attn_q_head_norm) { q = ggml_rms_norm(ctx, q, eps); // per-head unweighted (V4) @@ -12441,14 +13193,21 @@ static ggml_tensor * build_dspark_attention( ggml_tensor * kv_fa = ggml_reshape_3d(ctx, kv_attn, head_dim, n_attn, 1); ggml_tensor * mask_fa = nullptr; if (attn_mask) { - mask_fa = ggml_cast(ctx, ggml_repeat(ctx, ggml_reshape_2d(ctx, attn_mask, n_attn, 1), - ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_attn, block)), - GGML_TYPE_F16); + // Every layer attends over the same rows: build the F16 mask once. + if (flash_mask_cache && *flash_mask_cache && + (*flash_mask_cache)->ne[0] == n_attn && (*flash_mask_cache)->ne[1] == block) { + mask_fa = *flash_mask_cache; + } else { + mask_fa = ggml_cast(ctx, ggml_repeat(ctx, ggml_reshape_2d(ctx, attn_mask, n_attn, 1), + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_attn, block)), + GGML_TYPE_F16); + if (flash_mask_cache) *flash_mask_cache = mask_fa; + } } context = ggml_flash_attn_ext(ctx, q_fa, kv_fa, kv_fa, mask_fa, 1.0f / sqrtf((float) head_dim), 0.0f, 0.0f); - if (L.attn_sinks) { - ggml_flash_attn_ext_add_sinks(context, L.attn_sinks); + if (sinks) { + ggml_flash_attn_ext_add_sinks(context, sinks); } ggml_flash_attn_ext_set_prec(context, GGML_PREC_F32); ggml_flash_attn_ext_set_ds4_sparse(context, n_attn, n_attn, 0, 32); @@ -12466,8 +13225,8 @@ static ggml_tensor * build_dspark_attention( scores = ggml_add(ctx, scores, ggml_repeat(ctx, attn_mask, scores)); } ggml_tensor * probs = nullptr; - if (L.attn_sinks) { - ggml_tensor * sink = ggml_reshape_2d(ctx, L.attn_sinks, 1, n_head); + if (sinks) { + ggml_tensor * sink = ggml_reshape_2d(ctx, sinks, 1, n_head); ggml_tensor * sink_shape = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n_head * block); sink = ggml_repeat(ctx, sink, sink_shape); ggml_tensor * sws = ggml_concat(ctx, scores, sink, 0); // [n_attn+1, n_head*block] @@ -12486,14 +13245,62 @@ static ggml_tensor * build_dspark_attention( rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); ggml_tensor * attn_out = ggml_reshape_2d(ctx, context, head_dim * n_head, block); - const int group_dim = head_dim * (n_head / n_out_group); - attn_out = ggml_reshape_3d(ctx, attn_out, group_dim, n_out_group, block); - attn_out = ggml_cont(ctx, ggml_permute(ctx, attn_out, 0, 2, 1, 3)); // [group_dim, block, n_out_group] + const int heads_per_group = n_head_all / n_out_group; + const int group_dim = head_dim * heads_per_group; + const int n_groups = head_split ? n_head / heads_per_group : n_out_group; + const int group_first = head_split ? head_begin / heads_per_group : 0; + attn_out = ggml_reshape_3d(ctx, attn_out, group_dim, n_groups, block); + attn_out = ggml_cont(ctx, ggml_permute(ctx, attn_out, 0, 2, 1, 3)); // [group_dim, block, n_groups] ggml_tensor * out_a_3d = ggml_reshape_3d(ctx, L.attn_output_a, group_dim, n_lora_o, n_out_group); - ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); // [n_lora_o, block, n_out_group] - attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); // [n_lora_o, n_out_group, block] - attn_low = ggml_reshape_2d(ctx, attn_low, n_lora_o * n_out_group, block); - return ggml_mul_mat(ctx, L.attn_output_b, attn_low); // [n_embd, block] + if (head_split) { + out_a_3d = ggml_view_3d(ctx, out_a_3d, group_dim, n_lora_o, n_groups, + out_a_3d->nb[1], out_a_3d->nb[2], + (size_t) group_first * out_a_3d->nb[2]); + } + ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); // [n_lora_o, block, n_groups] + attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); // [n_lora_o, n_groups, block] + attn_low = ggml_reshape_2d(ctx, attn_low, n_lora_o * n_groups, block); + ggml_tensor * out_b_w = head_split + ? ggml_view_2d(ctx, L.attn_output_b, (int64_t) n_lora_o * n_groups, + L.attn_output_b->ne[1], L.attn_output_b->nb[1], + ggml_row_size(L.attn_output_b->type, (int64_t) group_first * n_lora_o)) + : L.attn_output_b; + return ggml_mul_mat(ctx, out_b_w, attn_low); // [n_embd, block] +} + +// Cluster split MoE: the router and its top-k run on every rank; each rank +// evaluates the routes that land on its experts (the others map to -1, which +// the MoE MMVQ skips and the masked combine drops) plus its slice of the +// shared expert. The result is this rank's partial of the FFN output. +static ggml_tensor * build_dspark_moe_split(ggml_context * ctx, + ggml_tensor * cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int n_tokens, + ggml_tensor * expert_lut, // I32 [1, n_expert, T] + ggml_tensor * expert_valid, // F32 [1, n_expert, T] + int shexp_begin, int shexp_count) { + const int n_embd = w.n_embd; + const int n_ff_exp = w.n_ff_exp; + Ds4MoeRouting routing = build_moe_routing(ctx, cur, w, L, n_tokens, nullptr); + const int n_used = (int) routing.selected->ne[0]; + ggml_tensor * local_ids = ggml_reshape_2d( + ctx, ggml_get_rows(ctx, expert_lut, routing.selected), n_used, n_tokens); + ggml_tensor * cur_3d = ggml_reshape_3d(ctx, cur, n_embd, 1, n_tokens); + ggml_tensor * gate_e = ggml_mul_mat_id(ctx, L.ffn_gate_exps, cur_3d, local_ids); + ggml_tensor * up_e = ggml_mul_mat_id(ctx, L.ffn_up_exps, cur_3d, local_ids); + ggml_mul_mat_set_mixed_mmq(gate_e, w.mixed_mmq_policy); + ggml_mul_mat_set_mixed_mmq(up_e, w.mixed_mmq_policy); + gate_e = ggml_reshape_3d(ctx, gate_e, n_ff_exp, n_used, n_tokens); + up_e = ggml_reshape_3d(ctx, up_e, n_ff_exp, n_used, n_tokens); + ggml_tensor * mid_e = build_clamped_swiglu(ctx, gate_e, up_e, w.swiglu_clamp_exp); + ggml_tensor * down_e = ggml_mul_mat_id(ctx, L.ffn_down_exps, mid_e, local_ids); + ggml_mul_mat_set_mixed_mmq(down_e, w.mixed_mmq_policy); + down_e = ggml_reshape_3d(ctx, down_e, n_embd, n_used, n_tokens); + ggml_tensor * routed = ggml_moe_combine_masked(ctx, down_e, routing.weights, expert_valid, + routing.selected); + ggml_tensor * shared = build_shared_ffn_slice(ctx, cur, w, L, shexp_begin, shexp_count); + return shared ? ggml_add(ctx, routed, shared) : nullptr; } // Read a small F32 GPU tensor (HC scale, [k]) into host floats. @@ -12514,122 +13321,140 @@ static void ds4_read_f32(ggml_tensor * t, float * dst, int k) { namespace { // Cache DSpark context KV projections across speculative steps. Only newly -// committed target-feature columns are projected on the draft GPU; the final -// normed+RoPE KV window is small enough (<1 MiB at n_swa=128) to retain in a -// host-side ring and upload as one compact draft-graph input. -struct DsparkContextKvProjector { +// committed target-feature columns are projected on the draft GPU, into a +// normed+RoPE KV window kept on that device. +// One projection graph for a given count of new columns. +struct DsparkContextKvGraph { int n_cols = -1; - const void * drafter = nullptr; - ggml_backend_t backend = nullptr; - std::vector arena; ggml_context * ctx = nullptr; ggml_gallocr_t alloc = nullptr; ggml_cgraph * gf = nullptr; ggml_tensor * inp_features = nullptr; ggml_tensor * positions = nullptr; - ggml_tensor * out = nullptr; + ggml_tensor * out = nullptr; // [head_dim, n_cols, n_layer] + void release() { + if (alloc) ggml_gallocr_free(alloc); + if (ctx) ggml_free(ctx); + *this = DsparkContextKvGraph{}; + } +}; + +struct DsparkContextKvProjector { + const void * drafter = nullptr; + ggml_backend_t backend = nullptr; + // A step projects accepted + 1 new columns, so the count changes from step + // to step: one graph per count (slot = n_cols for 1..7, slot 0 for larger + // counts such as the first full-window projection) instead of rebuilding + // and re-planning one graph whenever the count differs from the last. + static constexpr int kGraphSlots = 8; + DsparkContextKvGraph graphs[kGraphSlots]; int end_pos = -1; int valid = 0; - std::vector host_kv; - std::vector projected; + std::vector host_kv; // the all-zero window before any context + PinnedStage stage; // inputs of one projection + + // The window stays on the draft device, + // [head_dim, n_swa, n_layer] F32 with the valid rows end-aligned (the + // layout the draft graph reads), double-buffered so each step's shift and + // append is one batched copy from one buffer into the other. + ggml_context * win_ctx = nullptr; + ggml_backend_buffer_t win_buf = nullptr; + ggml_tensor * win[2] = {nullptr, nullptr}; + int win_cur = 0; + void release() { + for (DsparkContextKvGraph & g : graphs) g.release(); + if (win_buf) ggml_backend_buffer_free(win_buf); + if (win_ctx) ggml_free(win_ctx); + *this = DsparkContextKvProjector{}; + } }; thread_local DsparkContextKvProjector g_dspark_ctx_kv; +// Projects n_cols new feature columns to the drafter's per-layer context K/V, +// [head_dim, n_cols, n_layer] in *dev_out, queued on backend's stream. static bool dspark_project_context_columns( ggml_backend_t backend, const DSparkDrafter & d, const float * features, int n_cols, int first_pos, - std::vector & out) { + ggml_tensor ** dev_out) { if (!backend || !features || n_cols <= 0) return false; const DeepSeek4Weights & w = d.core; const int n_embd = w.n_embd; const int fc_in = d.n_target_layers * n_embd; const int head_dim = w.head_dim; DsparkContextKvProjector & P = g_dspark_ctx_kv; + if (P.drafter != (const void *) &d || P.backend != backend) { + for (DsparkContextKvGraph & g : P.graphs) g.release(); + P.drafter = (const void *) &d; + P.backend = backend; + } + DsparkContextKvGraph & G = + P.graphs[n_cols < DsparkContextKvProjector::kGraphSlots ? n_cols : 0]; - if (!P.ctx || P.n_cols != n_cols || P.drafter != (const void *) &d || - P.backend != backend) { - if (P.alloc && P.backend != backend) { - ggml_gallocr_free(P.alloc); - P.alloc = nullptr; - } - if (P.ctx) { - ggml_free(P.ctx); - P.ctx = nullptr; - } - P.gf = nullptr; - P.inp_features = nullptr; - P.positions = nullptr; - P.out = nullptr; - if (P.arena.empty()) P.arena.resize(32u * 1024 * 1024); + if (!G.ctx || G.n_cols != n_cols) { + G.release(); ggml_init_params ip{}; - ip.mem_size = P.arena.size(); - ip.mem_buffer = P.arena.data(); + ip.mem_size = ggml_tensor_overhead() * 512 + ggml_graph_overhead_custom(4096, false); + ip.mem_buffer = nullptr; ip.no_alloc = true; - P.ctx = ggml_init(ip); - if (!P.ctx) return false; - P.gf = ggml_new_graph_custom(P.ctx, 4096, false); - P.inp_features = ggml_new_tensor_2d( - P.ctx, GGML_TYPE_F32, fc_in, n_cols); - ggml_set_input(P.inp_features); - P.positions = ggml_new_tensor_1d( - P.ctx, GGML_TYPE_I32, n_cols); - ggml_set_input(P.positions); + G.ctx = ggml_init(ip); + if (!G.ctx) return false; + G.gf = ggml_new_graph_custom(G.ctx, 4096, false); + G.inp_features = ggml_new_tensor_2d( + G.ctx, GGML_TYPE_F32, fc_in, n_cols); + ggml_set_input(G.inp_features); + G.positions = ggml_new_tensor_1d( + G.ctx, GGML_TYPE_I32, n_cols); + ggml_set_input(G.positions); ggml_tensor * feature_norm = - ggml_rms_norm(P.ctx, P.inp_features, w.rms_eps); + ggml_rms_norm(G.ctx, G.inp_features, w.rms_eps); ggml_tensor * main_x = build_rms_norm( - P.ctx, ggml_mul_mat(P.ctx, d.main_proj, feature_norm), + G.ctx, ggml_mul_mat(G.ctx, d.main_proj, feature_norm), d.main_norm, w.rms_eps); ggml_tensor * stacked = nullptr; for (int il = 0; il < w.n_layer; ++il) { const DeepSeek4Layer & L = w.layers[(size_t) il]; ggml_tensor * kv = build_rms_norm( - P.ctx, ggml_mul_mat(P.ctx, L.attn_kv, main_x), + G.ctx, ggml_mul_mat(G.ctx, L.attn_kv, main_x), L.attn_kv_a_norm, w.rms_eps); kv = build_tail_rope_2d( - P.ctx, kv, P.positions, w.n_rot, head_dim, n_cols, + G.ctx, kv, G.positions, w.n_rot, head_dim, n_cols, w.rope_freq_base, 1.0f, 0.0f, 1.0f, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, (int) w.rope_orig_ctx); - kv = ggml_reshape_3d(P.ctx, kv, head_dim, n_cols, 1); - stacked = stacked ? ggml_concat(P.ctx, stacked, kv, 2) : kv; - } - P.out = ggml_cont(P.ctx, stacked); - ggml_set_output(P.out); - ggml_build_forward_expand(P.gf, P.out); - if (!P.alloc) { - P.alloc = ggml_gallocr_new( - ggml_backend_get_default_buffer_type(backend)); - } - if (!P.alloc || !ggml_gallocr_alloc_graph(P.alloc, P.gf)) { - ggml_free(P.ctx); - P.ctx = nullptr; - P.gf = nullptr; + kv = ggml_reshape_3d(G.ctx, kv, head_dim, n_cols, 1); + stacked = stacked ? ggml_concat(G.ctx, stacked, kv, 2) : kv; + } + G.out = ggml_cont(G.ctx, stacked); + ggml_set_output(G.out); + ggml_build_forward_expand(G.gf, G.out); + G.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!G.alloc || !ggml_gallocr_alloc_graph(G.alloc, G.gf)) { + G.release(); return false; } - P.n_cols = n_cols; - P.drafter = (const void *) &d; - P.backend = backend; + G.n_cols = n_cols; } - ggml_backend_tensor_set( - P.inp_features, features, 0, - sizeof(float) * (size_t) fc_in * n_cols); + // The inputs go through pinned slices; without the pinned block the stage + // falls back to blocking copies. + const size_t in_bytes = sizeof(float) * (size_t) fc_in * n_cols; + P.stage.reserve(backend, in_bytes + sizeof(int32_t) * (size_t) n_cols + 4096); + P.stage.reset(); std::vector pos((size_t) n_cols); for (int i = 0; i < n_cols; ++i) pos[(size_t) i] = first_pos + i; - ggml_backend_tensor_set( - P.positions, pos.data(), 0, sizeof(int32_t) * pos.size()); - if (ggml_backend_graph_compute(backend, P.gf) != GGML_STATUS_SUCCESS) { + P.stage.set(backend, G.inp_features, features, 0, in_bytes); + P.stage.set(backend, G.positions, pos.data(), 0, sizeof(int32_t) * pos.size()); + if (ggml_backend_graph_compute_async(backend, G.gf) != GGML_STATUS_SUCCESS) { + ggml_backend_synchronize(backend); return false; } - out.resize((size_t) head_dim * n_cols * w.n_layer); - ggml_backend_tensor_get( - P.out, out.data(), 0, sizeof(float) * out.size()); + *dev_out = G.out; return true; } @@ -12639,64 +13464,86 @@ static bool dspark_update_context_kv_cache( const float * ctx_features, int ctx_len, int committed, - const std::vector ** out_host_kv) { + const std::vector ** out_host_kv, + ggml_tensor ** out_dev_kv) { const DeepSeek4Weights & w = d.core; const int n_swa = w.n_swa; const int head_dim = w.head_dim; const int fc_in = d.n_target_layers * w.n_embd; DsparkContextKvProjector & P = g_dspark_ctx_kv; - if (ctx_len <= 0) { + if (ctx_len > 0) { + if (!ctx_features || ctx_len > n_swa) return false; + // Two device windows, valid rows end-aligned. Zeroed when (re)made; + // the valid span only grows until it fills, so the rows ahead of it + // stay zero in both buffers. + if (!P.win_buf || P.backend != backend || P.drafter != (const void *) &d) { + if (P.win_buf) ggml_backend_buffer_free(P.win_buf); + if (P.win_ctx) ggml_free(P.win_ctx); + P.win_buf = nullptr; P.win_ctx = nullptr; P.win[0] = P.win[1] = nullptr; + ggml_init_params ip{}; + ip.mem_size = ggml_tensor_overhead() * 4; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + P.win_ctx = ggml_init(ip); + if (!P.win_ctx) return false; + for (int b = 0; b < 2; ++b) { + P.win[b] = ggml_new_tensor_3d(P.win_ctx, GGML_TYPE_F32, head_dim, n_swa, w.n_layer); + } + P.win_buf = ggml_backend_alloc_ctx_tensors(P.win_ctx, backend); + if (!P.win_buf) return false; + P.end_pos = -1; // force the full projection below + } + int n_new = committed - P.end_pos; + const bool rebuild = + P.drafter != (const void *) &d || P.backend != backend || + P.end_pos < 0 || n_new <= 0 || n_new > ctx_len || + std::min(n_swa, P.valid + n_new) != ctx_len; + if (rebuild) { + ggml_backend_buffer_clear(P.win_buf, 0); + P.win_cur = 0; + P.valid = 0; + n_new = ctx_len; + } + const int first_new_pos = committed - n_new; + const float * new_features = ctx_features + (size_t) (ctx_len - n_new) * fc_in; + ggml_tensor * projected = nullptr; + if (!dspark_project_context_columns(backend, d, new_features, n_new, first_new_pos, + &projected) || !projected) { + return false; + } + const int keep = std::max(0, std::min(P.valid, n_swa - n_new)); + const int valid_next = keep + n_new; + const size_t row = sizeof(float) * (size_t) head_dim; + const uint8_t * src = (const uint8_t *) P.win[P.win_cur]->data; + uint8_t * dst = (uint8_t *) P.win[1 - P.win_cur]->data; + const uint8_t * fresh = (const uint8_t *) projected->data; + std::vector copies; + copies.reserve((size_t) w.n_layer * 2); + for (int il = 0; il < w.n_layer; ++il) { + const size_t layer = (size_t) il * (size_t) n_swa * row; + if (keep > 0) { + copies.push_back({src + layer + (size_t) (n_swa - keep) * row, + dst + layer + (size_t) (n_swa - valid_next) * row, + (size_t) keep * row}); + } + copies.push_back({fresh + (size_t) il * (size_t) n_new * row, + dst + layer + (size_t) (n_swa - n_new) * row, + (size_t) n_new * row}); + } + ggml_backend_cuda_copy_batch_async(backend, copies.data(), (int) copies.size()); + P.win_cur = 1 - P.win_cur; + P.valid = valid_next; P.end_pos = committed; - P.valid = 0; - P.host_kv.assign((size_t) w.n_layer * n_swa * head_dim, 0.0f); - *out_host_kv = &P.host_kv; + P.drafter = (const void *) &d; + if (P.valid != ctx_len) return false; + *out_dev_kv = P.win[P.win_cur]; + *out_host_kv = nullptr; return true; } - if (!ctx_features || ctx_len > n_swa) return false; - - int n_new = committed - P.end_pos; - const bool rebuild = - P.drafter != (const void *) &d || P.backend != backend || - P.end_pos < 0 || - n_new <= 0 || n_new > ctx_len || - std::min(n_swa, P.valid + n_new) != ctx_len || - P.host_kv.size() != (size_t) w.n_layer * n_swa * head_dim; - if (rebuild) { - P.host_kv.assign((size_t) w.n_layer * n_swa * head_dim, 0.0f); - P.valid = 0; - n_new = ctx_len; - } - - const int first_new_pos = committed - n_new; - const float * new_features = - ctx_features + (size_t) (ctx_len - n_new) * fc_in; - if (!dspark_project_context_columns( - backend, d, new_features, n_new, first_new_pos, P.projected)) { - return false; - } - - const int keep = std::max(0, std::min(P.valid, n_swa - n_new)); - const int drop = P.valid - keep; - for (int il = 0; il < w.n_layer; ++il) { - float * layer_cache = P.host_kv.data() + - (size_t) il * n_swa * head_dim; - if (keep > 0 && drop > 0) { - std::memmove( - layer_cache, - layer_cache + (size_t) drop * head_dim, - sizeof(float) * (size_t) keep * head_dim); - } - const float * layer_new = P.projected.data() + - (size_t) il * n_new * head_dim; - std::memcpy( - layer_cache + (size_t) keep * head_dim, - layer_new, - sizeof(float) * (size_t) n_new * head_dim); - } - P.valid = keep + n_new; + // No committed context yet: an all-zero window, which the mask hides. P.end_pos = committed; - P.drafter = (const void *) &d; - if (P.valid != ctx_len) return false; + P.valid = 0; + P.host_kv.assign((size_t) w.n_layer * n_swa * head_dim, 0.0f); *out_host_kv = &P.host_kv; return true; } @@ -12722,8 +13569,13 @@ struct DsparkDraftCache { ggml_tensor * attn_mask = nullptr; ggml_tensor * out = nullptr; ggml_tensor * confidence_out = nullptr; + ggml_tensor * expert_lut = nullptr; // cluster split: global -> local expert, -1 elsewhere + ggml_tensor * expert_valid = nullptr; // cluster split: 1 for this rank's experts + std::vector host_expert_lut; + std::vector host_expert_valid; std::vector padded_ctx; std::vector host_attn_mask; + PinnedStage stage; // one forward's inputs and readbacks std::vector> dbg_taps; // HC scales are immutable weights: read from the backend once. std::vector> s_attn, s_ffn; @@ -12746,16 +13598,7 @@ void reset_deepseek4_dspark_runtime_cache() { } cache = DsparkDraftCache{}; - DsparkContextKvProjector & projector = g_dspark_ctx_kv; - if (projector.alloc) { - ggml_gallocr_free(projector.alloc); - projector.alloc = nullptr; - } - if (projector.ctx) { - ggml_free(projector.ctx); - projector.ctx = nullptr; - } - projector = DsparkContextKvProjector{}; + g_dspark_ctx_kv.release(); } static bool deepseek4_dspark_draft_forward_impl( @@ -12786,15 +13629,20 @@ static bool deepseek4_dspark_draft_forward_impl( : valid_ctx_len; const std::vector * cached_host_kv = nullptr; + ggml_tensor * cached_dev_kv = nullptr; if (context_kv_cache && upload_context && !dspark_update_context_kv_cache( backend, d, ctx_features, valid_ctx_len, committed, - &cached_host_kv)) { + &cached_host_kv, &cached_dev_kv)) { return false; } DsparkDraftCache & C = g_dspark_draft_cache; const bool DS4_DBG = std::getenv("LUCE_DS4_DSPARK_DEBUG") != nullptr; + // LUCE_CLUSTER_DRAFT_SPLIT: half the heads and this rank's experts, the + // partial sums all-reduced after the attention and after the FFN. + Ds4ClusterRuntime * split_rt = d.split_size > 1 ? (Ds4ClusterRuntime *) d.cluster_rt : nullptr; + const bool dspark_split = split_rt && split_rt->size() == d.split_size && d.expert_local > 0; // Context reuse is deliberately strict: never submit a graph with an // uninitialized or differently-shaped context tensor. The normal warm @@ -12878,6 +13726,16 @@ static bool deepseek4_dspark_draft_forward_impl( C.attn_mask = ggml_new_tensor_1d( ctx, GGML_TYPE_F32, graph_ctx_len + block); ggml_set_input(C.attn_mask); + C.expert_lut = nullptr; + C.expert_valid = nullptr; + if (dspark_split) { + C.expert_lut = ggml_new_tensor_3d(ctx, GGML_TYPE_I32, 1, w.n_expert, block); + ggml_set_input(C.expert_lut); + ggml_set_output(C.expert_lut); + C.expert_valid = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, w.n_expert, block); + ggml_set_input(C.expert_valid); + ggml_set_output(C.expert_valid); + } // main_x = main_norm(main_proj(ctx_features)). Shared across layers. ggml_tensor * main_x = nullptr; @@ -12896,53 +13754,54 @@ static bool deepseek4_dspark_draft_forward_impl( dbg_tap("main_x", main_x); } - // HC state: [n_embd, n_hc, block], init = block embeds replicated over streams. + // HC state: [n_embd*n_hc, block], init = block embeds replicated over + // streams. The HC kernels take a token axis, so every sub-block + // boundary is one batched launch per op over all block positions + // (the per-position arithmetic is unchanged: rows normalize, mix and + // post independently), instead of per-position ops joined by concats. ggml_tensor * noise3 = ggml_reshape_3d(ctx, C.inp_noise, n_embd, 1, block); - ggml_tensor * hc_cur = ggml_repeat_4d(ctx, noise3, n_embd, n_hc, block, 1); + ggml_tensor * hc_cur = ggml_reshape_2d( + ctx, ggml_repeat_4d(ctx, noise3, n_embd, n_hc, block, 1), + (int64_t) n_embd * n_hc, block); - auto hc_col = [&](ggml_tensor * hc, int p) -> ggml_tensor * { - // Contiguous [n_embd*n_hc] slab for block position p. - return ggml_view_1d(ctx, hc, (int64_t) n_embd * n_hc, - (size_t) p * hc->nb[2]); - }; // Staggered pre-mix (V4.1, see ds4_hc_collapse): each sub-block's // input collapses the HC copies with the previous sub-block's pre // coefficients; the first attention takes copy 0 (identity mix) and - // the head the last FFN's. prev_pre[p] is null before layer 0. + // the head the last FFN's. prev_pre is null before layer 0. const bool staggered = w.hc_staggered_pre; - std::vector prev_pre((size_t) block, nullptr); - auto hc_collapse = [&](ggml_tensor * hcf, ggml_tensor * pre) { - return ds4_build_hc_collapse( - ctx, ggml_reshape_2d(ctx, hcf, (int64_t) n_embd * n_hc, 1), pre, n_embd, n_hc); + ggml_tensor * prev_pre = nullptr; // [n_hc, block] + ggml_tensor * flash_mask = nullptr; // shared by the layers' flash attention + // HC pre of one sub-block: returns its input [n_embd, block] and the + // split [mix_dim, block] that its HC post consumes. + auto hc_pre_block = [&](ggml_tensor * fn, ggml_tensor * base_w, + const std::array & s, + ggml_tensor ** split_out) -> ggml_tensor * { + ggml_tensor * normed = ggml_rms_norm(ctx, hc_cur, hc_eps); + ggml_tensor * mix = ggml_mul_mat(ctx, fn, normed); // [mix_dim, block] + ggml_tensor * base = ggml_reshape_1d(ctx, base_w, mix_dim); + ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base, hc_cur, n_hc, + w.n_hc_sinkhorn_iter, s[0], s[1], s[2]); + const size_t row = block > 1 ? pre->nb[1] : ggml_nbytes(pre); + *split_out = ggml_view_2d(ctx, pre, mix_dim, block, row, + (size_t) n_embd * sizeof(float)); + if (staggered) { + ggml_tensor * in = ds4_build_hc_collapse(ctx, hc_cur, prev_pre, n_embd, n_hc); + prev_pre = ds4_hc_split_pre(ctx, *split_out, n_hc); + return in; + } + return ggml_cont(ctx, ggml_view_2d(ctx, pre, n_embd, block, row, 0)); }; - auto split_pre = [&](ggml_tensor * split) { - return ds4_hc_split_pre(ctx, split, n_hc); + auto hc_post_block = [&](ggml_tensor * block_out, ggml_tensor * split) { + ggml_tensor * hp = ggml_ds4_hc_post(ctx, hc_cur, block_out, split, n_hc); + return ggml_reshape_2d(ctx, hp, (int64_t) n_embd * n_hc, block); }; for (int il = 0; il < w.n_layer; il++) { const DeepSeek4Layer & L = w.layers[il]; - // ── HC pre (attention), per block position ────────────────── - std::vector split_attn(block), work_cols(block); - for (int p = 0; p < block; p++) { - ggml_tensor * hcf = hc_col(hc_cur, p); - ggml_tensor * normed = ggml_rms_norm(ctx, hcf, hc_eps); - ggml_tensor * mix = ggml_mul_mat(ctx, L.hc_attn_fn, normed); - mix = ggml_reshape_1d(ctx, mix, mix_dim); - ggml_tensor * base = ggml_reshape_1d(ctx, L.hc_attn_base, mix_dim); - ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base, hcf, n_hc, - w.n_hc_sinkhorn_iter, - C.s_attn[il][0], C.s_attn[il][1], C.s_attn[il][2]); - split_attn[p] = ggml_view_1d(ctx, pre, mix_dim, (size_t) n_embd * sizeof(float)); - if (staggered) { - work_cols[p] = hc_collapse(hcf, prev_pre[(size_t) p]); - prev_pre[(size_t) p] = split_pre(split_attn[p]); - } else { - work_cols[p] = ggml_reshape_2d(ctx, ggml_view_1d(ctx, pre, n_embd, 0), n_embd, 1); - } - } - ggml_tensor * attn_in = work_cols[0]; - for (int p = 1; p < block; p++) attn_in = ggml_concat(ctx, attn_in, work_cols[p], 1); + // ── HC pre (attention) ────────────────────────────────────── + ggml_tensor * split_attn = nullptr; + ggml_tensor * attn_in = hc_pre_block(L.hc_attn_fn, L.hc_attn_base, C.s_attn[il], &split_attn); ggml_tensor * attn_normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); ggml_tensor * layer_ctx_kv = nullptr; if (context_kv_cache && graph_ctx_len > 0) { @@ -12951,85 +13810,64 @@ static bool deepseek4_dspark_draft_forward_impl( C.inp_ctx_kv->nb[1], (size_t) il * C.inp_ctx_kv->nb[2]); } + const int split_heads = dspark_split ? w.n_head / d.split_size : 0; ggml_tensor * attn_out = build_dspark_attention( ctx, attn_normed, main_x, layer_ctx_kv, w, L, graph_ctx_len, C.pos_block, C.neg_block, C.pos_ctx, - C.attn_mask, d.flash_attention); + C.attn_mask, d.flash_attention, + &flash_mask, + d.split_rank * split_heads, split_heads); + if (dspark_split) { + attn_out = ds4_cluster_allreduce_node(ctx, attn_out, *split_rt, /*inplace=*/true); + if (!attn_out) { ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; return false; } + } dbg_tap(std::string("attn_L") + std::to_string(il), attn_out); - // ── HC post (attention), per block position ───────────────── - ggml_tensor * hc_next = nullptr; - for (int p = 0; p < block; p++) { - ggml_tensor * bo = ggml_view_1d(ctx, attn_out, n_embd, (size_t) p * attn_out->nb[1]); - ggml_tensor * hp = ggml_ds4_hc_post(ctx, hc_col(hc_cur, p), bo, split_attn[p], n_hc); - hp = ggml_reshape_3d(ctx, hp, n_embd, n_hc, 1); - hc_next = hc_next ? ggml_concat(ctx, hc_next, hp, 2) : hp; - } - hc_cur = ggml_cont(ctx, hc_next); - - // ── HC pre (FFN), per block position ──────────────────────── - std::vector split_ffn(block), fwork(block); - for (int p = 0; p < block; p++) { - ggml_tensor * hcf = hc_col(hc_cur, p); - ggml_tensor * normed = ggml_rms_norm(ctx, hcf, hc_eps); - ggml_tensor * mix = ggml_mul_mat(ctx, L.hc_ffn_fn, normed); - mix = ggml_reshape_1d(ctx, mix, mix_dim); - ggml_tensor * base = ggml_reshape_1d(ctx, L.hc_ffn_base, mix_dim); - ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base, hcf, n_hc, - w.n_hc_sinkhorn_iter, - C.s_ffn[il][0], C.s_ffn[il][1], C.s_ffn[il][2]); - split_ffn[p] = ggml_view_1d(ctx, pre, mix_dim, (size_t) n_embd * sizeof(float)); - if (staggered) { - fwork[p] = hc_collapse(hcf, prev_pre[(size_t) p]); - prev_pre[(size_t) p] = split_pre(split_ffn[p]); - } else { - fwork[p] = ggml_reshape_2d(ctx, ggml_view_1d(ctx, pre, n_embd, 0), n_embd, 1); + // ── HC post (attention) ───────────────────────────────────── + hc_cur = hc_post_block(attn_out, split_attn); + + // ── HC pre (FFN) ──────────────────────────────────────────── + ggml_tensor * split_ffn = nullptr; + ggml_tensor * ffn_in = hc_pre_block(L.hc_ffn_fn, L.hc_ffn_base, C.s_ffn[il], &split_ffn); + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); + ggml_tensor * ffn_out = nullptr; + if (dspark_split) { + const int shexp_ff = L.ffn_gate_shexp ? (int) L.ffn_gate_shexp->ne[1] : 0; + const int shexp_count = shexp_ff / d.split_size; + ffn_out = build_dspark_moe_split(ctx, ffn_normed, w, L, block, C.expert_lut, + C.expert_valid, d.split_rank * shexp_count, + shexp_count); + if (ffn_out) { + ffn_out = ds4_cluster_allreduce_node(ctx, ffn_out, *split_rt, /*inplace=*/true); } + } else { + ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, block); } - ggml_tensor * ffn_in = fwork[0]; - for (int p = 1; p < block; p++) ffn_in = ggml_concat(ctx, ffn_in, fwork[p], 1); - ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); - ggml_tensor * ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, block); if (!ffn_out) { ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; return false; } dbg_tap(std::string("ffn_L") + std::to_string(il), ffn_out); // ── HC post (FFN) ─────────────────────────────────────────── - hc_next = nullptr; - for (int p = 0; p < block; p++) { - ggml_tensor * bo = ggml_view_1d(ctx, ffn_out, n_embd, (size_t) p * ffn_out->nb[1]); - ggml_tensor * hp = ggml_ds4_hc_post(ctx, hc_col(hc_cur, p), bo, split_ffn[p], n_hc); - hp = ggml_reshape_3d(ctx, hp, n_embd, n_hc, 1); - hc_next = hc_next ? ggml_concat(ctx, hc_next, hp, 2) : hp; - } - hc_cur = ggml_cont(ctx, hc_next); + hc_cur = hc_post_block(ggml_is_contiguous(ffn_out) ? ffn_out : ggml_cont(ctx, ffn_out), + split_ffn); dbg_tap(std::string("hcL") + std::to_string(il), hc_cur); } - // ── Tail: hc_head collapse -> out_norm, per block position ────── + // ── Tail: hc_head collapse -> out_norm ────────────────────────── // The tied lm_head consumes the normalized state. The confidence head // was trained on the HC-collapsed state before this output RMSNorm, so // keep both instead of reusing the normalized state for confidence. - ggml_tensor * out = nullptr; - ggml_tensor * confidence_out = nullptr; - for (int p = 0; p < block; p++) { - ggml_tensor * hcf = hc_col(hc_cur, p); - ggml_tensor * final_2d = nullptr; - if (staggered) { - final_2d = hc_collapse(hcf, prev_pre[(size_t) p]); - } else { - ggml_tensor * onorm = ggml_rms_norm(ctx, hcf, hc_eps); - ggml_tensor * omix = ggml_mul_mat(ctx, w.output_hc_fn, onorm); - omix = ggml_reshape_1d(ctx, omix, n_hc); - ggml_tensor * obase = ggml_reshape_1d(ctx, w.output_hc_base, n_hc); - ggml_tensor * final_embd = ggml_ds4_hc_out(ctx, omix, obase, hcf, n_hc, C.s_out); - final_2d = ggml_reshape_2d(ctx, final_embd, n_embd, 1); - } - ggml_tensor * hidden_p = build_rms_norm(ctx, final_2d, w.out_norm, w.rms_eps); - out = out ? ggml_concat(ctx, out, hidden_p, 1) : hidden_p; - confidence_out = confidence_out - ? ggml_concat(ctx, confidence_out, final_2d, 1) - : final_2d; - } + ggml_tensor * final_2d = nullptr; + if (staggered) { + final_2d = ds4_build_hc_collapse(ctx, hc_cur, prev_pre, n_embd, n_hc); + } else { + ggml_tensor * onorm = ggml_rms_norm(ctx, hc_cur, hc_eps); + ggml_tensor * omix = ggml_mul_mat(ctx, w.output_hc_fn, onorm); // [n_hc, block] + ggml_tensor * obase = ggml_reshape_1d(ctx, w.output_hc_base, n_hc); + ggml_tensor * final_embd = ggml_ds4_hc_out(ctx, omix, obase, hc_cur, n_hc, C.s_out); + final_2d = ggml_reshape_2d(ctx, final_embd, n_embd, block); + } + ggml_tensor * out = build_rms_norm(ctx, final_2d, w.out_norm, w.rms_eps); + ggml_tensor * confidence_out = final_2d; ggml_set_output(out); ggml_set_output(confidence_out); ggml_build_forward_expand(gf, out); @@ -13059,10 +13897,33 @@ static bool deepseek4_dspark_draft_forward_impl( } // ── Set inputs + compute (cached graph) ───────────────────────────── - ggml_backend_tensor_set(C.inp_noise, noise_embed, 0, sizeof(float) * (size_t) n_embd * block); + // Every upload and readback of the forward goes through pinned slices, + // so it costs one wait instead of a round trip per tensor. + const bool pinned_io = C.stage.reserve(backend, + (C.inp_noise ? ggml_nbytes(C.inp_noise) : 0) + + (C.inp_ctx ? ggml_nbytes(C.inp_ctx) : 0) + + (C.inp_ctx_kv ? ggml_nbytes(C.inp_ctx_kv) : 0) + + (C.pos_ctx ? ggml_nbytes(C.pos_ctx) : 0) + + (C.attn_mask ? ggml_nbytes(C.attn_mask) : 0) + + (C.expert_lut ? ggml_nbytes(C.expert_lut) + ggml_nbytes(C.expert_valid) : 0) + + 2 * (C.out ? ggml_nbytes(C.out) : 0) + 8192); + C.stage.reset(); + // Without the pinned block the stage falls back to blocking copies. + const auto set_in = [&](ggml_tensor * t, const void * data, size_t bytes) { + C.stage.set(backend, t, data, 0, bytes); + }; + set_in(C.inp_noise, noise_embed, sizeof(float) * (size_t) n_embd * block); if (upload_context) { if (graph_ctx_len > 0) { - if (context_kv_cache) { + if (context_kv_cache && cached_dev_kv) { + // The device window is already in the input's layout + // (graph_ctx_len == n_swa, valid rows end-aligned). + if (graph_ctx_len != w.n_swa || + ggml_nbytes(cached_dev_kv) != ggml_nbytes(C.inp_ctx_kv)) return false; + const ggml_cuda_copy_desc to_input{cached_dev_kv->data, C.inp_ctx_kv->data, + ggml_nbytes(cached_dev_kv)}; + ggml_backend_cuda_copy_batch_async(backend, &to_input, 1); + } else if (context_kv_cache) { if (!cached_host_kv) return false; const int head_dim = w.head_dim; const int source_stride = w.n_swa * head_dim; @@ -13081,9 +13942,7 @@ static bool deepseek4_dspark_draft_forward_impl( } upload = C.padded_ctx.data(); } - ggml_backend_tensor_set( - C.inp_ctx_kv, upload, 0, - sizeof(float) * (size_t) w.n_layer * + set_in(C.inp_ctx_kv, upload, sizeof(float) * (size_t) w.n_layer * graph_ctx_len * head_dim); } else { const float * upload = ctx_features; @@ -13099,16 +13958,12 @@ static bool deepseek4_dspark_draft_forward_impl( upload = C.padded_ctx.data(); } if (!upload) return false; - ggml_backend_tensor_set( - C.inp_ctx, upload, 0, - sizeof(float) * (size_t) fc_in * graph_ctx_len); + set_in(C.inp_ctx, upload, sizeof(float) * (size_t) fc_in * graph_ctx_len); std::vector pc(graph_ctx_len); for (int i = 0; i < graph_ctx_len; i++) { pc[i] = committed - graph_ctx_len + i; } - ggml_backend_tensor_set( - C.pos_ctx, pc.data(), 0, - sizeof(int32_t) * graph_ctx_len); + set_in(C.pos_ctx, pc.data(), sizeof(int32_t) * graph_ctx_len); } } C.host_attn_mask.assign((size_t) graph_ctx_len + block, 0.0f); @@ -13117,15 +13972,29 @@ static bool deepseek4_dspark_draft_forward_impl( C.host_attn_mask[(size_t)i] = -std::numeric_limits::infinity(); } - ggml_backend_tensor_set( - C.attn_mask, C.host_attn_mask.data(), 0, - sizeof(float) * C.host_attn_mask.size()); + set_in(C.attn_mask, C.host_attn_mask.data(), sizeof(float) * C.host_attn_mask.size()); C.valid_ctx_len = valid_ctx_len; } std::vector pb(block), nb(block); for (int i = 0; i < block; i++) { pb[i] = committed + i; nb[i] = -(committed + i); } - ggml_backend_tensor_set(C.pos_block, pb.data(), 0, sizeof(int32_t) * block); - ggml_backend_tensor_set(C.neg_block, nb.data(), 0, sizeof(int32_t) * block); + set_in(C.pos_block, pb.data(), sizeof(int32_t) * block); + set_in(C.neg_block, nb.data(), sizeof(int32_t) * block); + if (C.expert_lut && C.expert_valid) { + const int n_expert = w.n_expert; + if (C.host_expert_lut.size() != (size_t) n_expert * block) { + C.host_expert_lut.resize((size_t) n_expert * block); + C.host_expert_valid.resize((size_t) n_expert * block); + for (int t = 0; t < block; ++t) { + for (int e = 0; e < n_expert; ++e) { + const bool local = e >= d.expert_first && e < d.expert_first + d.expert_local; + C.host_expert_lut[(size_t) t * n_expert + e] = local ? e - d.expert_first : -1; + C.host_expert_valid[(size_t) t * n_expert + e] = local ? 1.0f : 0.0f; + } + } + } + set_in(C.expert_lut, C.host_expert_lut.data(), sizeof(int32_t) * C.host_expert_lut.size()); + set_in(C.expert_valid, C.host_expert_valid.data(), sizeof(float) * C.host_expert_valid.size()); + } // This cache owns a single immutable graph: topology, tensor addresses, // and shapes remain fixed until the cache is explicitly rebuilt above. @@ -13141,23 +14010,31 @@ static bool deepseek4_dspark_draft_forward_impl( /*disable_graphs=*/false, /*mmvq_max_ncols=*/0, /*skip_property_check=*/force_graph_replay); - const ggml_status st = out_hidden - ? ggml_backend_graph_compute(backend, C.gf) - : ggml_backend_graph_compute_async(backend, C.gf); + const ggml_status st = ggml_backend_graph_compute_async(backend, C.gf); if (st != GGML_STATUS_SUCCESS) { + ggml_backend_synchronize(backend); // Invalidate: a failed compute leaves no reusable state guarantees. ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; C.ctx_len = -1; return false; } - if (!out_hidden) return true; + if (!out_hidden) { + // The async probe path keeps the pinned inputs until it waits. + if (pinned_io) ggml_backend_synchronize(backend); + return true; + } + const size_t hidden_bytes = sizeof(float) * (size_t) n_embd * block; out_hidden->resize((size_t) n_embd * block); - ggml_backend_tensor_get( - C.out, out_hidden->data(), 0, sizeof(float) * out_hidden->size()); + if (confidence_hidden) confidence_hidden->resize((size_t) n_embd * block); + const void * hidden_src = C.stage.get(backend, C.out, 0, hidden_bytes); + const void * conf_src = confidence_hidden + ? C.stage.get(backend, C.confidence_out, 0, hidden_bytes) : nullptr; + ggml_backend_synchronize(backend); + if (hidden_src) std::memcpy(out_hidden->data(), hidden_src, hidden_bytes); + else ggml_backend_tensor_get(C.out, out_hidden->data(), 0, hidden_bytes); if (confidence_hidden) { - confidence_hidden->resize((size_t) n_embd * block); - ggml_backend_tensor_get(C.confidence_out, confidence_hidden->data(), 0, - sizeof(float) * confidence_hidden->size()); + if (conf_src) std::memcpy(confidence_hidden->data(), conf_src, hidden_bytes); + else ggml_backend_tensor_get(C.confidence_out, confidence_hidden->data(), 0, hidden_bytes); } if (DS4_DBG) { diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index f2bd8fccf..31f420c6a 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -32,6 +32,7 @@ #include "ggml-backend.h" #include "internal.h" +#include "cluster/cluster_config.h" #include "common/layer_split_utils.h" #include "common/paged_attention_config.h" #include "common/prefill_attention_mode.h" @@ -76,6 +77,12 @@ class MoeStreamedExpertCache; struct MoeExpertCacheOptions; struct DeepSeek4StepTelemetry { + // Expert-parallel cluster (server/src/cluster): host time in the + // per-layer all-reduce helper, payload bytes handed to the communicator, + // and time blocked on the control channel (decisions/drafts from rank 0). + uint64_t cluster_allreduce_us = 0; + uint64_t cluster_allreduce_bytes = 0; + uint64_t cluster_ctrl_wait_us = 0; uint64_t total_us = 0; uint64_t embed_us = 0; uint64_t hc_pre_attn_us = 0; @@ -248,6 +255,11 @@ struct DeepSeek4Weights { int head_dim = 512; // = value_dim for DS4 int n_rot = 64; // partial RoPE rotation dims int n_out_group = 8; // grouped output projection + // Cluster rank slices held instead of the full tensors (TargetLoadPlan + // slice_*): query heads and shared-expert intermediate units kept, 0 = + // the full tensors. The model's n_head / n_ff_exp stay the global ones. + int cluster_heads_local = 0; + int cluster_shexp_ff_local = 0; // Low-rank attention dimensions int n_lora_q = 1024; // Q low-rank bottleneck @@ -489,6 +501,9 @@ struct DeepSeek4LayerMajorBand { // Optional per-band Engram keys: read at the band's first Engram layer, // reused at the next (they depend only on the band's tokens). std::vector * engram_keys = nullptr; + // Cluster ranks: set once the band's keys (every Engram layer's rows) + // came from rank 0, so later Engram layers of the band skip the broadcast. + uint8_t * engram_keys_shared = nullptr; // Pipelined prefill (LUCE_DS4_PREFILL_PIPELINE): which set of prefill HC // graphs this band uses (bands alternate, so a band's deferred FFN join // keeps its own tensors), and the pass's pipeline state. @@ -507,7 +522,17 @@ struct DeepSeek4LayerMajorTail { std::vector pre; // [rows][n_hc] }; + +// Expert-parallel cluster runtime (deepseek4_cluster.h). Non-null only when +// the backend runs as one rank of a cluster; the forward then masks routes to +// other ranks' experts and all-reduces the routed partial. +struct Ds4ClusterRuntime; + struct DeepSeek4Cache { + // Set by DeepSeek4Backend when it runs as a cluster rank; owned by the + // backend, never freed here. nullptr keeps every forward path unchanged. + Ds4ClusterRuntime * cluster_rt = nullptr; + int cur_pos = 0; int max_ctx = 0; int n_layer = 0; @@ -588,6 +613,8 @@ struct DeepSeek4BackendConfig { std::string expert_placement_path; // three-tier expert ownership (JSON) std::string router_bias_path; // f32 [n_layer][n_expert] selection bias delta std::string protected_experts_path; // {"layer": [expert ids]} (JSON) + // Expert-parallel cluster rank (--cluster-*); size 0 = a single box. + cluster::ClusterConfig cluster; }; // ─── Function declarations ────────────────────────────────────────────── @@ -831,6 +858,11 @@ bool deepseek4_step( struct Ds4VerifyWindowRows { std::vector> kv; // [layer] -> [batch token][row bytes] std::vector> score; + // The fused verify keeps the rows on the device instead: [layer] -> the + // verify graph's own row tensors per batch token, valid until that graph + // slot runs again (a rollback restores them before the next verify). + std::vector> kv_dev; + std::vector> score_dev; }; struct Ds4VerifyHooks { diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index 009ed1626..cf9dd2b65 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -475,6 +475,17 @@ static bool should_split_ds4_dense_tensor(const char * name, int mask) { return false; } +// A rank's slice of a dense tensor (TargetLoadPlan slice_*): `count` rows +// (or columns, or elements of a vector) starting at `first`, out of the +// file's full rows of `full_ne0` elements. +struct DS4TensorSlice { + enum Kind { None, Rows, Columns } kind = None; + int64_t first = 0; + int64_t count = 0; + int64_t full_ne0 = 0; + int64_t full_ne1 = 0; +}; + struct DS4TensorAlloc { ggml_tensor * tensor = nullptr; size_t tensor_offset = 0; @@ -483,7 +494,20 @@ struct DS4TensorAlloc { size_t buffer_offset = 0; bool upload_to_backend = true; bool dense_split = false; + DS4TensorSlice slice; }; + +// Reshape `t` to its slice, in place (metadata only, before allocation). +static void ds4_apply_slice_shape(ggml_tensor * t, const DS4TensorSlice & s) { + if (s.kind == DS4TensorSlice::Rows) { + t->ne[1] = s.count; + } else if (s.kind == DS4TensorSlice::Columns) { + t->ne[0] = s.count; + t->nb[1] = ggml_row_size(t->type, t->ne[0]); + } + t->nb[2] = t->nb[1] * t->ne[1]; + t->nb[3] = t->nb[2] * t->ne[2]; +} } // namespace // ─── Source-layer resolution ──────────────────────────────────────────── @@ -2252,6 +2276,78 @@ bool load_deepseek4_gguf_partial(const std::string & path, dense_tp_mask, 1.0f - strix_fraction, strix_fraction); } + // Cluster rank slices: whole output groups of heads, block-aligned + // shared-expert units, or nothing (the full tensors are loaded). + int slice_head_first = 0, slice_heads = 0, slice_group_first = 0, slice_groups = 0; + int slice_ff_first = 0, slice_ff = 0; + if (plan.slice_ranks > 1 && plan.slice_attention && + n_head % (uint32_t) plan.slice_ranks == 0 && n_out_group % (uint32_t) plan.slice_ranks == 0) { + slice_heads = (int) n_head / plan.slice_ranks; + slice_head_first = plan.slice_rank * slice_heads; + slice_groups = (int) n_out_group / plan.slice_ranks; + slice_group_first = plan.slice_rank * slice_groups; + } + if (plan.slice_ranks > 1 && plan.slice_shared_expert && + n_ff_exp % (uint32_t) plan.slice_ranks == 0 && + (n_ff_exp / (uint32_t) plan.slice_ranks) % 256 == 0) { + slice_ff = (int) n_ff_exp / plan.slice_ranks; + slice_ff_first = plan.slice_rank * slice_ff; + } + // A slice the quant blocks cannot follow fails the load: keeping that one + // tensor whole would no longer match its sliced partner. + std::string slice_error; + const auto slice_for = [&](const char * name, const ggml_tensor * t) { + DS4TensorSlice s; + int layer_id = -1; + if (!parse_block_tensor_name(name, layer_id)) return s; + const auto ends_with = [&](const char * suffix) { + const size_t n = std::strlen(name), m = std::strlen(suffix); + return n >= m && std::strcmp(name + n - m, suffix) == 0; + }; + s.full_ne0 = t->ne[0]; + s.full_ne1 = t->ne[1]; + if (slice_heads > 0 && !std::strstr(name, ".indexer.")) { + if (ends_with(".attn_q_b.weight") && t->ne[1] == (int64_t) n_head * head_dim) { + s.kind = DS4TensorSlice::Rows; + s.first = (int64_t) slice_head_first * head_dim; + s.count = (int64_t) slice_heads * head_dim; + } else if (ends_with(".attn_sinks.weight") && t->ne[0] == (int64_t) n_head) { + s.kind = DS4TensorSlice::Columns; + s.first = slice_head_first; + s.count = slice_heads; + } else if (ends_with(".attn_output_a.weight") && + t->ne[1] == (int64_t) n_lora_o * n_out_group) { + s.kind = DS4TensorSlice::Rows; + s.first = (int64_t) slice_group_first * n_lora_o; + s.count = (int64_t) slice_groups * n_lora_o; + } else if (ends_with(".attn_output_b.weight") && + t->ne[0] == (int64_t) n_lora_o * n_out_group) { + s.kind = DS4TensorSlice::Columns; + s.first = (int64_t) slice_group_first * n_lora_o; + s.count = (int64_t) slice_groups * n_lora_o; + } + } + if (s.kind == DS4TensorSlice::None && slice_ff > 0) { + if ((ends_with(".ffn_gate_shexp.weight") || ends_with(".ffn_up_shexp.weight")) && + t->ne[1] == (int64_t) n_ff_exp) { + s.kind = DS4TensorSlice::Rows; + s.first = slice_ff_first; + s.count = slice_ff; + } else if (ends_with(".ffn_down_shexp.weight") && t->ne[0] == (int64_t) n_ff_exp) { + s.kind = DS4TensorSlice::Columns; + s.first = slice_ff_first; + s.count = slice_ff; + } + } + if (s.kind == DS4TensorSlice::Columns && ggml_blck_size(t->type) > 1 && + (s.first % ggml_blck_size(t->type) != 0 || s.count % ggml_blck_size(t->type) != 0)) { + if (slice_error.empty()) slice_error = name; + s.kind = DS4TensorSlice::None; + } + return s; + }; + size_t sliced_bytes_saved = 0; + std::vector allocs; allocs.reserve(n_tensors); size_t total_buf_size = 0; @@ -2275,6 +2371,14 @@ bool load_deepseek4_gguf_partial(const std::string & path, a.upload_to_backend = upload_to_backend; a.dense_split = upload_to_backend && split_buft && should_split_ds4_dense_tensor(tname, dense_tp_mask); + if (upload_to_backend && !a.dense_split) { + a.slice = slice_for(tname, t); + if (a.slice.kind != DS4TensorSlice::None) { + const size_t full_bytes = ggml_nbytes(t); + ds4_apply_slice_shape(t, a.slice); + sliced_bytes_saved += full_bytes - ggml_nbytes(t); + } + } if (upload_to_backend) { if (a.dense_split) { split_total_buf_size = align_up_size( @@ -2294,6 +2398,14 @@ bool load_deepseek4_gguf_partial(const std::string & path, } } + if (!slice_error.empty()) { + set_last_error("cluster slice of " + slice_error + " is not aligned to its quant blocks; " + + "this model cannot be sliced over " + std::to_string(plan.slice_ranks) + " ranks"); + gguf_free(gctx); + if (meta_ctx) ggml_free(meta_ctx); + return false; + } + // ── Allocate GPU buffer ───────────────────────────────────────────── ggml_backend_buffer_t buf = nullptr; ggml_backend_buffer_t split_buf = nullptr; @@ -2385,6 +2497,27 @@ bool load_deepseek4_gguf_partial(const std::string & path, #endif for (auto & a : allocs) { if (!a.upload_to_backend) continue; + if (a.slice.kind != DS4TensorSlice::None) { + // A rank slice: its rows are one contiguous range of the file; + // its columns are a range of every row, gathered on the host. + const uint8_t * full = (const uint8_t *) mmap.addr + a.file_offset; + const size_t full_row = ggml_row_size(a.tensor->type, a.slice.full_ne0); + if (a.slice.kind == DS4TensorSlice::Rows) { + ggml_backend_tensor_set(a.tensor, full + (size_t) a.slice.first * full_row, 0, + ggml_nbytes(a.tensor)); + } else { + const size_t row = ggml_row_size(a.tensor->type, a.slice.count); + const size_t skip = ggml_row_size(a.tensor->type, a.slice.first); + const int64_t rows = a.slice.full_ne1 * a.tensor->ne[2] * a.tensor->ne[3]; + std::vector gathered((size_t) rows * row); + for (int64_t r = 0; r < rows; ++r) { + std::memcpy(gathered.data() + (size_t) r * row, + full + (size_t) r * full_row + skip, row); + } + ggml_backend_tensor_set(a.tensor, gathered.data(), 0, gathered.size()); + } + continue; + } if (read_file && !a.dense_split) { spans.push_back({a.tensor, 0, a.file_offset, a.file_size}); continue; @@ -2484,6 +2617,17 @@ bool load_deepseek4_gguf_partial(const std::string & path, out.embedder.row_bytes = a.file_size / (size_t)n_vocab; } + if (slice_heads > 0 || slice_ff > 0) { + out.cluster_heads_local = slice_heads; + out.cluster_shexp_ff_local = slice_ff; + std::fprintf(stderr, "[deepseek4-cluster] rank %d/%d dense slices: heads %d..%d of %u, " + "shared expert units %d..%d of %u (%.2f GiB not loaded)\n", + plan.slice_rank, plan.slice_ranks, slice_head_first, + slice_head_first + slice_heads - 1, n_head, slice_ff_first, + slice_ff_first + slice_ff - 1, n_ff_exp, + sliced_bytes_saved / 1073741824.0); + } + // ── Bind tensors to weight struct fields ──────────────────────────── for (auto & a : allocs) { const char * name = ggml_get_name(a.tensor); diff --git a/server/src/internal.h b/server/src/internal.h index dfd52d875..519d6ade2 100644 --- a/server/src/internal.h +++ b/server/src/internal.h @@ -272,6 +272,14 @@ struct TargetLoadPlan { bool metadata_only = false; // parse tensor descriptors/scales without GPU allocation bool expert_metadata_only = false; // keep only routed expert tensor metadata; upload nothing bool load_ds4_image_bias = false; + // DS4 cluster rank slices: keep only this rank's attention output groups + // (attn_q_b rows, attn_sinks, attn_output_a groups, attn_output_b + // columns) and its shared-expert intermediate slice (gate/up rows, down + // columns) of slice_ranks equal parts. + int slice_rank = 0; + int slice_ranks = 1; + bool slice_attention = false; + bool slice_shared_expert = false; }; // Load a Q4_K_M target model from a GGUF file on disk. diff --git a/server/src/server/launch_profiles.h b/server/src/server/launch_profiles.h index 16cc1f878..77f9dfdcb 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -246,6 +246,61 @@ inline const std::vector & launch_profiles() { return profiles; } +// Environment defaults of expert-parallel cluster mode (--cluster-size > 1): +// the measured two-box DeepSeek V4.1 configuration. They are installed before +// a launch profile's defaults, so they win over a profile but never over a +// variable that is already set. Setting a switch to 0 turns it off; +// LUCE_CLUSTER_FAST_REDUCE_STALL_S=0 turns the stall watchdog off, and the +// cache slot counts and LUCE_DS4_ENGRAM_PRECOMPUTE take a size. +inline const std::vector & cluster_launch_env() { + static const std::vector env = { + // Decode exchange: the raw-RDMA fast reduce, zero-copy rows on a bf16 + // wire, no slot-reuse waits; the hybrid exchange reduces the iGPU's + // cold partial with the R9700's hot + shared partial in one export. + {"LUCE_CLUSTER_FAST_REDUCE", "1"}, + {"LUCE_CLUSTER_FAST_REDUCE_STALL_S", "15"}, + {"LUCE_CLUSTER_HYBRID_EXCHANGE", "1"}, + // Prefill-sized reductions stay on RCCL: 8-bit blocks, or bf16 with + // LUCE_CLUSTER_PREFILL_Q8=0. + {"LUCE_CLUSTER_PREFILL_Q8", "1"}, + {"LUCE_CLUSTER_PREFILL_BF16", "1"}, + // Work split: attention heads, the greedy verify head and the DSpark + // draft head are halved across the ranks; each rank loads only its + // slices of the dense weights. + {"LUCE_CLUSTER_ATTENTION_PARALLEL", "1"}, + {"LUCE_CLUSTER_VOCAB_SPLIT", "1"}, + {"LUCE_CLUSTER_DRAFT_VOCAB_SPLIT", "1"}, + {"LUCE_CLUSTER_DRAFT_SPLIT", "1"}, + {"LUCE_CLUSTER_WORKER_SKIP_DRAFT", "1"}, + // DS4.1 verify step: batched verify lanes, the drafter's context K/V + // window, GPU argmax, the drafter swapped out for long prompts (a + // cluster rank runs without the prefix cache, so no cached prefix has + // to match the swap's chunks), and the Engram projection run inside + // the exchange wait. + {"LUCE_DS4_GPU_ARGMAX_VERIFY", "1"}, + {"LUCE_DS4_LANE_BATCH", "1"}, + {"LUCE_DS4_DRAFT_CONTEXT_KV_CACHE", "1"}, + {"LUCE_DS4_DRAFT_SWAP", "1"}, + {"LUCE_DS4_SCHED_DEDICATED_COPIES", "1"}, + {"LUCE_DS4_HC_BOUNDARY_SPLIT", "1"}, + {"LUCE_DS4_ENGRAM_PRECOMPUTE", "2560"}, + {"LUCE_DS4_TP_CAPTURE_CACHE_SLOTS", "64"}, + {"LUCE_DS4_TP_FUSED_CACHE_SLOTS", "64"}, + // Expert products: q8_1 inputs handed to the expert owners, grouped + // MUL_MAT_ID (in prefill on both owners: the R9700's hot experts too), + // the ROCmFP2 prefetch loop and one peer-copy kernel per layer. + {"LUCE_MOE_Q8_HANDOFF", "1"}, + {"LUCE_MMID_GROUPED", "1"}, + {"LUCE_DS4_GROUPED_HYBRID_PREFILL", "2"}, + {"LUCE_MMID_GROUPED_TYPES", "79"}, + {"LUCE_CUDA_MMVQ_MOE_FP2_PREFETCH", "1"}, + {"GGML_CUDA_PEER_COPY_KERNEL", "1"}, + // The fast reduce reads pinned host rows from the GPU. + {"HSA_FORCE_FINE_GRAIN_PCIE", "1"}, + }; + return env; +} + inline const LaunchProfile * find_launch_profile(const std::string & name) { for (const LaunchProfile & profile : launch_profiles()) { if (name == profile.name) return &profile; diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 905c31edb..1d76bc53b 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -16,6 +16,9 @@ #include "launch_profiles.h" #include "model_card.h" #include "common/backend_factory.h" +#include "cluster/cluster_config.h" +#include "cluster/cluster_head_backend.h" +#include "cluster/cluster_worker_main.h" #include "common/chain_rollback_policy.h" #include "common/gguf_inspect.h" #include "common/layer_split_utils.h" @@ -186,6 +189,19 @@ static void print_usage(const char * prog) { " --ds4-prefill DeepSeek4 prefill: exact, dense, or sparse\n" " (default: exact; dense/sparse are experimental\n" " and may change generated tokens)\n" + " --cluster-rank This box's rank in [0, N); rank 0 serves HTTP\n" + " --cluster-size Number of ranks (2..8)\n" + " --cluster-head Control endpoint; rank 0 binds, workers connect\n" + " --cluster-ifname RoCE interface -> NCCL_SOCKET_IFNAME\n" + " --cluster-ib-hca RDMA device -> NCCL_IB_HCA\n" + " --cluster-gid-index RoCE v2 GID index -> NCCL_IB_GID_INDEX\n" + " --cluster-expert-placement uniform|balanced|\n" + " Rank of every routed expert (same file on every rank)\n" + " --cluster-shared-expert replicate|shard (default: shard)\n" + " --cluster-allreduce-dtype f32|bf16|auto (default: auto = f32)\n" + " --cluster-timeout-ms Collective/control watchdog (default: 30000)\n" + " --cluster-replicate-hot Replicate each layer's k hottest experts on every rank\n" + " --cluster-selftest Bootstrap, run the collective self-test, exit\n" " --fa-window Flash-attention sliding window (default: 0=full).\n" " WARNING: >0 drops system prompt / tool definitions\n" " from attention at long contexts. Use 0 for tools.\n" @@ -590,6 +606,74 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, : std::strcmp(flag, "--ds4-protected-experts") == 0 ? bargs.ds4_protected_experts : bargs.expert_placement; dst = path; + } else if (std::strncmp(argv[i], "--cluster-", 10) == 0) { + // Expert-parallel cluster flags. + const char * flag = argv[i]; + const auto int_value = [&](int & out) { + if (i + 1 >= argc) return false; + const char * v = argv[++i]; + const char * end = v + std::strlen(v); + const auto parsed = std::from_chars(v, end, out); + if (parsed.ec != std::errc{} || parsed.ptr != end) { + std::fprintf(stderr, "[server] %s expects an integer, got '%s'\n", flag, v); + return false; + } + return true; + }; + std::string err; + if (std::strcmp(flag, "--cluster-rank") == 0) { + if (!int_value(bargs.cluster.rank)) return 2; + } else if (std::strcmp(flag, "--cluster-size") == 0) { + if (!int_value(bargs.cluster.size)) return 2; + } else if (std::strcmp(flag, "--cluster-head") == 0 && i + 1 < argc) { + if (!luce::cluster::parse_host_port(argv[++i], bargs.cluster.head_host, + bargs.cluster.head_port, &err)) { + std::fprintf(stderr, "[server] bad --cluster-head value: %s\n", err.c_str()); + return 2; + } + } else if (std::strcmp(flag, "--cluster-ifname") == 0 && i + 1 < argc) { + bargs.cluster.ifname = argv[++i]; + } else if (std::strcmp(flag, "--cluster-ib-hca") == 0 && i + 1 < argc) { + bargs.cluster.ib_hca = argv[++i]; + } else if (std::strcmp(flag, "--cluster-gid-index") == 0) { + if (!int_value(bargs.cluster.gid_index) || bargs.cluster.gid_index < 0) { + std::fprintf(stderr, "[server] --cluster-gid-index must be non-negative\n"); + return 2; + } + } else if (std::strcmp(flag, "--cluster-expert-placement") == 0 && i + 1 < argc) { + if (!luce::cluster::parse_placement_source(argv[++i], bargs.cluster.placement_source, + bargs.cluster.placement_file, &err)) { + std::fprintf(stderr, "[server] bad --cluster-expert-placement: %s\n", err.c_str()); + return 2; + } + } else if (std::strcmp(flag, "--cluster-shared-expert") == 0 && i + 1 < argc) { + if (!luce::cluster::parse_shared_expert_mode(argv[++i], bargs.cluster.shared_expert, &err)) { + std::fprintf(stderr, "[server] bad --cluster-shared-expert: %s\n", err.c_str()); + return 2; + } + } else if (std::strcmp(flag, "--cluster-allreduce-dtype") == 0 && i + 1 < argc) { + if (!luce::cluster::parse_allreduce_dtype(argv[++i], bargs.cluster.allreduce_dtype, &err)) { + std::fprintf(stderr, "[server] bad --cluster-allreduce-dtype: %s\n", err.c_str()); + return 2; + } + } else if (std::strcmp(flag, "--cluster-timeout-ms") == 0) { + int v = 0; + if (!int_value(v) || v <= 0) { + std::fprintf(stderr, "[server] --cluster-timeout-ms must be positive\n"); + return 2; + } + bargs.cluster.timeout_ms = (uint32_t) v; + } else if (std::strcmp(flag, "--cluster-replicate-hot") == 0) { + if (!int_value(bargs.cluster.replicate_hot) || bargs.cluster.replicate_hot < 0) { + std::fprintf(stderr, "[server] --cluster-replicate-hot must be non-negative\n"); + return 2; + } + } else if (std::strcmp(flag, "--cluster-selftest") == 0) { + bargs.cluster.selftest = true; + } else { + std::fprintf(stderr, "[server] unknown or incomplete cluster flag %s\n", flag); + return 2; + } } else if (std::strcmp(argv[i], "--ds4-prefill") == 0 && i + 1 < argc) { const char * mode = argv[++i]; bargs.ds4_prefill_mode_set = true; @@ -1174,27 +1258,46 @@ static bool apply_expert_device(const DevicePlacement & expert, return true; } -// Install a profile's environment defaults. Variables that are already set -// keep their value, and the log names them. -static void apply_launch_profile_env(const luce::server::LaunchProfile & profile) { - if (profile.env.empty()) return; +// Variables an earlier set of defaults installed (cluster mode before a +// launch profile): a later set leaves them alone without calling them explicit. +static std::set & installed_env_defaults() { + static std::set names; + return names; +} + +// Install a set of environment defaults. Variables that are already set keep +// their value, and the log names the ones the operator set. +static void apply_env_defaults(const std::string & owner, + const std::vector & defaults) { + if (defaults.empty()) return; int applied = 0; std::string kept_env; - for (const luce::server::LaunchProfileEnv & env : profile.env) { + for (const luce::server::LaunchProfileEnv & env : defaults) { if (std::getenv(env.name)) { - kept_env += std::string(kept_env.empty() ? "" : ", ") + env.name; + if (!installed_env_defaults().count(env.name)) { + kept_env += std::string(kept_env.empty() ? "" : ", ") + env.name; + } continue; } set_environment_variable(env.name, env.value, false); + installed_env_defaults().insert(env.name); ++applied; } - std::fprintf(stderr, "[server] profile %s: %d environment defaults applied%s%s\n", - profile.name, applied, + std::fprintf(stderr, "[server] %s: %d environment defaults applied%s%s\n", + owner.c_str(), applied, kept_env.empty() ? "" : "; kept explicit ", kept_env.c_str()); } +static void apply_launch_profile_env(const luce::server::LaunchProfile & profile) { + apply_env_defaults(std::string("profile ") + profile.name, profile.env); +} + static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_model) { + // Cluster mode's defaults go in before a profile's, so they win over it. + if (model.bargs.cluster.enabled()) { + apply_env_defaults("cluster mode", luce::server::cluster_launch_env()); + } if (model.profile) apply_launch_profile_env(*model.profile); if (!model.adaptive_experts_tau.empty()) set_environment_variable("LUCE_ADAPTIVE_K_TAU", model.adaptive_experts_tau.c_str(), false); @@ -1272,6 +1375,42 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod bargs.routing_stats_requested = sconfig.freq_tracking || !sconfig.collect_routing_path.empty(); + // ── Expert-parallel cluster: structural validation, RCCL environment. + // The request descriptor the head broadcasts carries no prefix-cache + // restore points and snapshots are not shipped between ranks, so every + // snapshot cache stays off on a cluster rank. + const luce::cluster::ClusterConfig cluster_cfg = bargs.cluster; + if (cluster_cfg.enabled() || cluster_cfg.rank >= 0) { + if (!cluster_cfg.enabled()) { + std::fprintf(stderr, "[server] --cluster-rank requires --cluster-size\n"); + return 2; + } + const std::string cluster_error = cluster_cfg.validate(); + if (!cluster_error.empty()) { + std::fprintf(stderr, "[server] %s\n", cluster_error.c_str()); + return 2; + } + sconfig.prefix_cache_cap = 0; + sconfig.prefill_cache_cap = 0; + sconfig.disk_cache_dir.clear(); + sconfig.disk_cache_policy.mode = DiskPrefixCacheMode::Off; + const int n_env = luce::cluster::export_rccl_environment(cluster_cfg); + std::fprintf(stderr, + "cluster: rank %d/%d head=%s:%d iface=%s hca=%s placement=%s " + "(gid=%d shared_expert=%s allreduce=%s timeout=%u ms; %d NCCL vars exported)\n", + cluster_cfg.rank, cluster_cfg.size, + cluster_cfg.head_host.c_str(), cluster_cfg.head_port, + cluster_cfg.ifname.empty() ? "(env)" : cluster_cfg.ifname.c_str(), + cluster_cfg.ib_hca.empty() ? "(env)" : cluster_cfg.ib_hca.c_str(), + cluster_cfg.placement_source == luce::cluster::PlacementSource::File + ? cluster_cfg.placement_file.c_str() + : luce::cluster::placement_source_name(cluster_cfg.placement_source), + cluster_cfg.gid_index, + luce::cluster::shared_expert_mode_name(cluster_cfg.shared_expert), + luce::cluster::allreduce_dtype_name(cluster_cfg.allreduce_dtype), + cluster_cfg.timeout_ms, n_env); + } + BackendAdmissionContext backend_admission; backend_admission.pflash_enabled = sconfig.pflash_mode != ServerConfig::PflashMode::OFF; @@ -1552,6 +1691,22 @@ static int load_model(ModelOptions & model, LoadedModel & loaded, bool multi_mod } return 1; } + if (cluster_cfg.is_worker()) { + // Ranks >= 1 serve no HTTP: they replay what rank 0 broadcasts. + return luce::cluster::run_cluster_worker(cluster_cfg, std::move(backend_owner), + backend_model.path, + backend_placement.target.gpu); + } + if (cluster_cfg.is_head()) { + auto head = std::make_unique( + std::move(backend_owner), cluster_cfg, backend_model.path, + backend_placement.target.gpu); + if (!head->init()) { + std::fprintf(stderr, "[server] cluster head init failed\n"); + return 1; + } + backend_owner = std::move(head); + } ModelBackend * backend = backend_owner.get(); // --max-ctx auto: serve the context the backend fitted to its devices. if (backend_placement.target.fit_ctx && backend->context_limit() > 0) { @@ -2294,6 +2449,9 @@ int main(int argc, char ** argv) { auto model = std::make_unique(); const int ret = load_model(option, *model, load_balancing); if (ret != 0) return ret; + // A cluster worker (rank >= 1) has already run to completion inside + // load_model and serves no HTTP: its exit code is the process's. + if (!model->server) return ret; servers.push_back(model->server.get()); loaded.push_back(std::move(model)); } diff --git a/server/test/unit/test_cluster_unit.cpp b/server/test/unit/test_cluster_unit.cpp new file mode 100644 index 000000000..418c745ef --- /dev/null +++ b/server/test/unit/test_cluster_unit.cpp @@ -0,0 +1,1182 @@ +// GPU-free unit tests for the expert-parallel cluster runtime. +// +// Covers the control-channel wire protocol (cluster_protocol.h), the CLI +// config helpers (cluster_config.h), the N-rank expert placement +// (cluster_expert_placement.h) and the rank-0-decides hooks +// (cluster_decision_hooks.h) driven over a real loopback control channel. +// None of this needs ggml, HIP or RCCL: the target links only the cluster +// sources plus moe_hybrid_placement / moe_hybrid_routing_stats / sampler. +// +// Build: cmake --build . --target test_cluster_unit +// Run: ./test_cluster_unit +// +// main() comes from test/test_unit_main.cpp (CMake adds it, like every +// target in _new_cppunit_test_targets). + +#include "CppUnitTestFramework.hpp" + +#include "cluster/cluster_config.h" +#include "cluster/cluster_decision_hooks.h" +#include "cluster/cluster_expert_placement.h" +#include "cluster/cluster_protocol.h" +#include "common/moe_hybrid_placement.h" +#include "common/moe_hybrid_routing_stats.h" + +#if !defined(_WIN32) +#include "cluster/cluster_control.h" +#include +#include +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace CppUnitTestFramework; +using namespace luce::cluster; + +namespace { + +// DeepSeek V4 Flash shape used throughout the placement tests. +constexpr int kLayers = 43; +constexpr int kExperts = 256; +constexpr int kTopK = 6; +constexpr int kSlots = 8; + +luce::common::MoeHybridConfig ds4_cfg() { + luce::common::MoeHybridConfig cfg; + cfg.n_embd = 4096; + cfg.n_expert = kExperts; + cfg.n_expert_used = kTopK; + cfg.n_ff_exp = 2048; + cfg.n_layer = kLayers; + cfg.first_moe_layer = 0; + return cfg; +} + +// Helpers that assert (REQUIRE/CHECK) must be fixture members: the framework's +// HandleAssert is a CommonFixture method, so free functions cannot call it. +struct ClusterUnitFixture : CommonFixture { + using CommonFixture::CommonFixture; + + // Hotness that is adversarial to e % N: every fourth expert is hot, so + // uniform piles all heat onto rank 0 for N = 2 and N = 4. + luce::common::MoeHybridRoutingStats skewed_stats() { + luce::common::MoeHybridRoutingStats s; + REQUIRE(s.init(kLayers, kExperts, kTopK)); + for (int l = 0; l < kLayers; l++) { + uint64_t total = 0; + for (int e = 0; e < kExperts; e++) { + const uint64_t c = (e % 4 == 0) ? 1000u + (uint64_t) (kExperts - e) : 1u + (uint64_t) (e % 7); + s.counts[(size_t) l * kExperts + e] = c; + total += c; + } + s.layer_totals[(size_t) l] = total; + } + return s; + } + + // Round-trips `msg` through encode/decode and checks that every strict + // prefix of the payload is rejected and that trailing bytes fail parse_frame. + template + void check_roundtrip_and_truncation(const Msg & in, Msg & out) { + Frame f = make_frame(in); + REQUIRE(f.type == Msg::type()); + REQUIRE(parse_frame(f, out)); + + for (size_t cut = 0; cut < f.payload.size(); cut++) { + Frame t; + t.type = f.type; + t.payload.assign(f.payload.begin(), f.payload.begin() + (std::ptrdiff_t) cut); + Msg scratch; + CHECK_FALSE(parse_frame(t, scratch)); + } + Frame g = f; + g.payload.push_back(0xAA); + Msg scratch; + CHECK_FALSE(parse_frame(g, scratch)); + + Frame wrong = f; + wrong.type = (f.type == MsgType::Heartbeat) ? MsgType::Shutdown : MsgType::Heartbeat; + CHECK_FALSE(parse_frame(wrong, scratch)); + } +}; + +// Per-layer per-rank load from stats over owned experts. +std::vector rank_loads(const ClusterExpertPlacement & p, + const luce::common::MoeHybridRoutingStats & s, int layer) { + std::vector load((size_t) p.n_ranks, 0); + for (int e = 0; e < p.n_expert; e++) { + const int32_t o = p.owner_of(layer, e); + if (o >= 0) load[(size_t) o] += s.count(layer, e); + } + return load; +} + +double max_over_mean(const std::vector & load) { + const uint64_t mx = *std::max_element(load.begin(), load.end()); + const double mean = (double) std::accumulate(load.begin(), load.end(), (uint64_t) 0) / + (double) load.size(); + return mean > 0 ? (double) mx / mean : 0.0; +} + +} // namespace + +// ─── Protocol: bytes and header ───────────────────────────────────────── + +TEST_CASE(ClusterUnitFixture, ByteWriterReaderRoundTrip) { + ByteWriter w; + w.u8(0xAB); + w.u16(0xBEEF); + w.u32(0xDEADBEEFu); + w.u64(0x0123456789ABCDEFull); + w.i32(-17); + w.i64(-9000000000ll); + w.f32(1.5f); + w.str("hello"); + w.str(""); + const uint8_t raw[3] = {1, 2, 3}; + w.bytes(raw, 3); + w.vec_i32({}); + w.vec_i32({-1, 0, 1}); + w.vec_i64({-5, 5}); + w.vec_u64({7}); + + ByteReader r(w.data()); + uint8_t a; uint16_t b; uint32_t c; uint64_t d; int32_t e; int64_t f; float g; + std::string s1, s2; + std::vector bytes; + std::vector v0, v1; + std::vector v2; + std::vector v3; + REQUIRE(r.u8(a)); CHECK_EQUAL(a, 0xAB); + REQUIRE(r.u16(b)); CHECK_EQUAL(b, 0xBEEF); + REQUIRE(r.u32(c)); CHECK_EQUAL(c, 0xDEADBEEFu); + REQUIRE(r.u64(d)); CHECK_EQUAL(d, 0x0123456789ABCDEFull); + REQUIRE(r.i32(e)); CHECK_EQUAL(e, -17); + REQUIRE(r.i64(f)); CHECK_EQUAL(f, -9000000000ll); + REQUIRE(r.f32(g)); CHECK_EQUAL(g, 1.5f); + REQUIRE(r.str(s1)); CHECK_EQUAL(s1, std::string("hello")); + REQUIRE(r.str(s2)); CHECK(s2.empty()); + REQUIRE(r.bytes(bytes)); REQUIRE_EQUAL(bytes.size(), (size_t) 3); CHECK_EQUAL(bytes[2], 3); + REQUIRE(r.vec_i32(v0)); CHECK(v0.empty()); + REQUIRE(r.vec_i32(v1)); REQUIRE_EQUAL(v1.size(), (size_t) 3); CHECK_EQUAL(v1[0], -1); + REQUIRE(r.vec_i64(v2)); REQUIRE_EQUAL(v2.size(), (size_t) 2); CHECK_EQUAL(v2[0], -5); + REQUIRE(r.vec_u64(v3)); REQUIRE_EQUAL(v3.size(), (size_t) 1); CHECK_EQUAL(v3[0], 7u); + CHECK(r.at_end()); + CHECK_FALSE(r.failed()); + + // Reading past the end fails and latches. + uint32_t extra = 0; + CHECK_FALSE(r.u32(extra)); + CHECK(r.failed()); +} + +TEST_CASE(ClusterUnitFixture, ByteReaderRejectsOversizedVectors) { + ByteWriter w; + w.u32(1000); // declares 1000 elements, none follow + ByteReader r(w.data()); + std::vector v; + CHECK_FALSE(r.vec_i32(v)); + + ByteWriter w2; + w2.u32(10); + w2.u8(1); + ByteReader r2(w2.data()); + std::string s; + CHECK_FALSE(r2.str(s, /*max_len=*/4)); // exceeds the caller's cap +} + +TEST_CASE(ClusterUnitFixture, FrameHeaderEncodeDecode) { + FrameHeader h; + h.type = (uint16_t) MsgType::Decision; + h.payload_len = 17; + uint8_t buf[kFrameHeaderBytes]; + encode_header(h, buf); + // Little-endian magic "LBCL" == 0x4C43424C. + CHECK_EQUAL(buf[0], 0x4C); CHECK_EQUAL(buf[1], 0x42); CHECK_EQUAL(buf[2], 0x43); CHECK_EQUAL(buf[3], 0x4C); + + FrameHeader out; + std::string err; + REQUIRE(decode_header(buf, out, &err)); + CHECK_EQUAL(out.magic, kFrameMagic); + CHECK_EQUAL(out.version, kProtocolVersion); + CHECK_EQUAL(out.type, (uint16_t) MsgType::Decision); + CHECK_EQUAL(out.payload_len, 17u); + + uint8_t bad_magic[kFrameHeaderBytes]; + std::memcpy(bad_magic, buf, sizeof(buf)); + bad_magic[0] ^= 0xFF; + CHECK_FALSE(decode_header(bad_magic, out, &err)); + CHECK_FALSE(err.empty()); + + FrameHeader v = h; + v.version = kProtocolVersion + 1; + encode_header(v, buf); + CHECK_FALSE(decode_header(buf, out, &err)); + + FrameHeader big = h; + big.payload_len = (uint32_t) kMaxPayloadBytes + 1; + encode_header(big, buf); + CHECK_FALSE(decode_header(buf, out, &err)); + + FrameHeader max_ok = h; + max_ok.payload_len = (uint32_t) kMaxPayloadBytes; + encode_header(max_ok, buf); + CHECK(decode_header(buf, out, &err)); +} + +TEST_CASE(ClusterUnitFixture, Fnv1a64KnownVectors) { + // Reference values of the 64-bit FNV-1a function. + CHECK_EQUAL(fnv1a64(nullptr, 0), 0xcbf29ce484222325ull); + CHECK_EQUAL(fnv1a64("a", 1), 0xaf63dc4c8601ec8cull); + CHECK_EQUAL(fnv1a64("foobar", 6), 0x85944171f73967e8ull); + // Chaining with the previous hash as seed equals hashing the concatenation. + const uint64_t h1 = fnv1a64("foo", 3); + CHECK_EQUAL(fnv1a64("bar", 3, h1), fnv1a64("foobar", 6)); +} + +TEST_CASE(ClusterUnitFixture, MsgTypeNamesAreDistinct) { + const MsgType types[] = { + MsgType::Hello, MsgType::Welcome, MsgType::Heartbeat, MsgType::Abort, + MsgType::Request, MsgType::Decision, MsgType::Draft, MsgType::Accept, + MsgType::RequestReport, MsgType::BackendOp, MsgType::RequestEnd, MsgType::Shutdown, + }; + std::set names; + for (MsgType t : types) { + const char * n = msg_type_name(t); + REQUIRE_NOT_NULL(n); + CHECK(std::string(n) != "Unknown"); + names.insert(n); + } + CHECK_EQUAL(names.size(), sizeof(types) / sizeof(types[0])); +} + +// ─── Protocol: messages ───────────────────────────────────────────────── + +TEST_CASE(ClusterUnitFixture, HelloWelcomeRoundTrip) { + HelloMsg in; + in.rank = 3; + in.size = 4; + in.build_sha = "v1.2.3-7-gabcdef"; + in.model_sha = std::string(64, 'f'); + in.placement_hash = 0x1122334455667788ull; + in.hostname = "halo-3"; + HelloMsg out; + check_roundtrip_and_truncation(in, out); + CHECK_EQUAL(out.rank, 3); + CHECK_EQUAL(out.size, 4); + CHECK_EQUAL(out.protocol_version, kProtocolVersion); + CHECK_EQUAL(out.build_sha, in.build_sha); + CHECK_EQUAL(out.model_sha, in.model_sha); + CHECK_EQUAL(out.placement_hash, in.placement_hash); + CHECK_EQUAL(out.hostname, in.hostname); + + HelloMsg empty; // defaults with empty strings still round-trip + HelloMsg empty_out; + check_roundtrip_and_truncation(empty, empty_out); + CHECK_EQUAL(empty_out.rank, -1); + CHECK(empty_out.build_sha.empty()); + + WelcomeMsg w; + w.size = 4; + for (size_t i = 0; i < kRcclUniqueIdBytes; i++) w.rccl_unique_id[i] = (uint8_t) (i * 7); + w.placement_hash = 42; + w.timeout_ms = 30000; + w.allreduce_dtype = 2; + w.shared_expert = 1; + w.rank_hostnames = {"halo-0", "halo-1", "", "halo-3"}; + WelcomeMsg wo; + check_roundtrip_and_truncation(w, wo); + CHECK_EQUAL(wo.size, 4); + CHECK(wo.rccl_unique_id == w.rccl_unique_id); + CHECK_EQUAL(wo.placement_hash, 42u); + CHECK_EQUAL(wo.timeout_ms, 30000u); + CHECK_EQUAL(wo.allreduce_dtype, 2); + CHECK_EQUAL(wo.shared_expert, 1); + REQUIRE_EQUAL(wo.rank_hostnames.size(), (size_t) 4); + CHECK(wo.rank_hostnames[2].empty()); + CHECK_EQUAL(wo.rank_hostnames[3], std::string("halo-3")); + + WelcomeMsg w_empty; + WelcomeMsg w_empty_out; + check_roundtrip_and_truncation(w_empty, w_empty_out); + CHECK(w_empty_out.rank_hostnames.empty()); +} + +TEST_CASE(ClusterUnitFixture, HeartbeatAbortShutdownRoundTrip) { + HeartbeatMsg hb; + hb.rank = 2; + hb.mono_us = 123456789012ull; + HeartbeatMsg hbo; + check_roundtrip_and_truncation(hb, hbo); + CHECK_EQUAL(hbo.rank, 2); + CHECK_EQUAL(hbo.mono_us, hb.mono_us); + + AbortMsg ab; + ab.rank = 1; + ab.code = 3; + ab.request_id = 99; + ab.reason = "worker generate failed while head succeeded"; + AbortMsg abo; + check_roundtrip_and_truncation(ab, abo); + CHECK_EQUAL(abo.rank, 1); + CHECK_EQUAL(abo.code, 3); + CHECK_EQUAL(abo.request_id, 99u); + CHECK_EQUAL(abo.reason, ab.reason); + + AbortMsg ab_empty; + AbortMsg ab_empty_out; + check_roundtrip_and_truncation(ab_empty, ab_empty_out); + CHECK(ab_empty_out.reason.empty()); + + ShutdownMsg sd; + sd.reason = 7; + ShutdownMsg sdo; + check_roundtrip_and_truncation(sd, sdo); + CHECK_EQUAL(sdo.reason, 7); +} + +TEST_CASE(ClusterUnitFixture, RequestRoundTripEmptyAndLarge) { + RequestMsg in; + in.request_id = 7; + in.n_gen = 512; + in.max_ctx = 32768; + in.temperature = 0.7f; + in.top_p = 0.95f; + in.top_k = 40; + in.min_p = 0.05f; + in.repeat_penalty = 1.1f; + in.seed = 0xDEADBEEFCAFEBABEull; + in.decode_mode = DecodeMode::Speculative; + in.force_ar = true; + in.restore_slot = 5; + in.kv_offset = 1024; + in.snapshot_slot = 3; + in.snapshot_pos = 2048; + RequestMsg out; + + SECTION("empty prompt and stop tokens") { + check_roundtrip_and_truncation(in, out); + CHECK(out.prompt_tokens.empty()); + CHECK(out.stop_token_ids.empty()); + CHECK_EQUAL(out.request_id, 7u); + CHECK_EQUAL(out.n_gen, 512); + CHECK_EQUAL(out.max_ctx, 32768); + CHECK_EQUAL(out.temperature, 0.7f); + CHECK_EQUAL(out.top_p, 0.95f); + CHECK_EQUAL(out.top_k, 40); + CHECK_EQUAL(out.min_p, 0.05f); + CHECK_EQUAL(out.repeat_penalty, 1.1f); + CHECK_EQUAL(out.seed, in.seed); + CHECK(out.decode_mode == DecodeMode::Speculative); + CHECK(out.force_ar); + // A forced AR retry decodes AR on every rank, whatever the mode. + CHECK(request_decodes_ar(out)); + out.force_ar = false; + CHECK_FALSE(request_decodes_ar(out)); + // Restore and inline save are independent halves; a request that + // resumes from one slot and check-points into another must survive + // the wire with both, or the workers' slots drift from the head's. + CHECK_EQUAL(out.restore_slot, 5); + CHECK_EQUAL(out.snapshot_slot, 3); + CHECK_EQUAL(out.snapshot_pos, 2048); + CHECK_EQUAL(out.kv_offset, 1024); + } + SECTION("128k prompt") { + in.prompt_tokens.resize(131072); + for (size_t i = 0; i < in.prompt_tokens.size(); i++) in.prompt_tokens[i] = (int32_t) (i * 31 % 129280); + in.stop_token_ids = {1718, 37947, 32}; + Frame f = make_frame(in); + REQUIRE(parse_frame(f, out)); + CHECK(out.prompt_tokens == in.prompt_tokens); + CHECK(out.stop_token_ids == in.stop_token_ids); + CHECK(f.payload.size() < kMaxPayloadBytes); + // Truncating inside the token vector must fail cleanly. + Frame t = f; + t.payload.resize(t.payload.size() / 2); + RequestMsg scratch; + CHECK_FALSE(parse_frame(t, scratch)); + } + SECTION("defaults") { + RequestMsg d; + RequestMsg d_out; + check_roundtrip_and_truncation(d, d_out); + CHECK(d_out.decode_mode == DecodeMode::Autoregressive); + CHECK_FALSE(d_out.force_ar); + CHECK(request_decodes_ar(d_out)); + CHECK_EQUAL(d_out.snapshot_slot, -1); + } +} + +TEST_CASE(ClusterUnitFixture, DecisionDraftAcceptRoundTrip) { + DecisionMsg d; + d.request_id = 5; + d.step = 4095; + d.token = 129279; + d.flags = kDecisionEos | kDecisionFinal; + DecisionMsg d_out; + check_roundtrip_and_truncation(d, d_out); + CHECK_EQUAL(d_out.request_id, 5u); + CHECK_EQUAL(d_out.step, 4095u); + CHECK_EQUAL(d_out.token, 129279); + CHECK_EQUAL(d_out.flags, (uint8_t) (kDecisionEos | kDecisionFinal)); + + DecisionMsg neg; + neg.token = -1; + neg.flags = kDecisionCancel; + DecisionMsg neg_out; + check_roundtrip_and_truncation(neg, neg_out); + CHECK_EQUAL(neg_out.token, -1); + + DraftMsg dr; + dr.request_id = 5; + dr.step = 9; + dr.pos = 2047; + dr.tokens = {11, 22, 33, 44, 55}; + DraftMsg dr_out; + check_roundtrip_and_truncation(dr, dr_out); + CHECK_EQUAL(dr_out.pos, 2047); + CHECK(dr_out.tokens == dr.tokens); + + DraftMsg dr_empty; + DraftMsg dr_empty_out; + check_roundtrip_and_truncation(dr_empty, dr_empty_out); + CHECK(dr_empty_out.tokens.empty()); + + AcceptMsg ac; + ac.request_id = 5; + ac.step = 9; + ac.accept = 3; + ac.bonus = 777; + ac.flags = kDecisionStop; + AcceptMsg ac_out; + check_roundtrip_and_truncation(ac, ac_out); + CHECK_EQUAL(ac_out.accept, 3); + CHECK_EQUAL(ac_out.bonus, 777); + CHECK_EQUAL(ac_out.flags, (uint8_t) kDecisionStop); +} + +TEST_CASE(ClusterUnitFixture, ReportBackendOpRequestEndRoundTrip) { + RequestReportMsg rr; + rr.request_id = 1; + rr.rank = 1; + rr.steps = 512; + rr.compute_us = 20'000'000; + rr.allreduce_calls = 43 * 512; + rr.allreduce_bytes = 16384ull * 43 * 512; + rr.allreduce_wait_us = 1'500'000; + rr.ctrl_wait_us = 400'000; + rr.peak_device_bytes = 90ull << 30; + RequestReportMsg rr_out; + check_roundtrip_and_truncation(rr, rr_out); + CHECK_EQUAL(rr_out.steps, 512u); + CHECK_EQUAL(rr_out.allreduce_bytes, rr.allreduce_bytes); + CHECK_EQUAL(rr_out.peak_device_bytes, rr.peak_device_bytes); + + BackendOpMsg op; + op.request_id = 0; + op.kind = BackendOpKind::Park; + op.args = {2}; + BackendOpMsg op_out; + check_roundtrip_and_truncation(op, op_out); + CHECK(op_out.kind == BackendOpKind::Park); + REQUIRE_EQUAL(op_out.args.size(), (size_t) 1); + CHECK_EQUAL(op_out.args[0], (int64_t) 2); + + BackendOpMsg op_noargs; + op_noargs.kind = BackendOpKind::FreeDrafter; + BackendOpMsg op_noargs_out; + check_roundtrip_and_truncation(op_noargs, op_noargs_out); + CHECK(op_noargs_out.args.empty()); + + RequestEndMsg re; + re.request_id = 1; + re.steps = 12; + re.status = 6; + RequestEndMsg re_out; + check_roundtrip_and_truncation(re, re_out); + CHECK_EQUAL(re_out.steps, 12u); + CHECK_EQUAL(re_out.status, 6); +} + +TEST_CASE(ClusterUnitFixture, DecisionFlagHelpers) { + CHECK_EQUAL(make_decision_flags(false, false, false, false, false), (uint8_t) 0); + CHECK_EQUAL(make_decision_flags(true, false, false, false, false), (uint8_t) kDecisionEos); + CHECK_EQUAL(make_decision_flags(false, true, true, true, true), + (uint8_t) (kDecisionStop | kDecisionCancel | kDecisionBudget | kDecisionFinal)); + CHECK_FALSE(decision_terminates(0)); + CHECK_FALSE(decision_terminates(kDecisionBudget)); + CHECK(decision_terminates(kDecisionEos)); + CHECK(decision_terminates(kDecisionCancel)); + CHECK(decision_terminates(kDecisionFinal)); +} + +// ─── Config ───────────────────────────────────────────────────────────── + +TEST_CASE(ClusterUnitFixture, ClusterConfigValidate) { + ClusterConfig c; + CHECK_FALSE(c.enabled()); + CHECK(c.validate().empty()); // "not a cluster run" is valid + + c.size = 2; + c.rank = 0; + c.head_host = "10.0.0.1"; + CHECK(c.validate().empty()); + CHECK(c.is_head()); + CHECK_FALSE(c.is_worker()); + + c.rank = 1; + CHECK(c.validate().empty()); + CHECK(c.is_worker()); + + ClusterConfig bad = c; + bad.rank = 2; // == size + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.rank = -1; + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.size = 1; + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.size = kClusterMaxSize + 1; + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.head_host.clear(); + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.head_port = 0; + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.replicate_hot = -1; + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.timeout_ms = 0; + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.gid_index = -2; + CHECK_FALSE(bad.validate().empty()); + bad = c; bad.placement_source = PlacementSource::File; bad.placement_file.clear(); + CHECK_FALSE(bad.validate().empty()); + bad.placement_file = "placement.json"; + CHECK(bad.validate().empty()); + + ClusterConfig eight = c; + eight.size = kClusterMaxSize; + eight.rank = kClusterMaxSize - 1; + CHECK(eight.validate().empty()); +} + +TEST_CASE(ClusterUnitFixture, ClusterConfigParsers) { + std::string host; + int port = 0; + std::string err; + REQUIRE(parse_host_port("10.0.0.1:9400", host, port, &err)); + CHECK_EQUAL(host, std::string("10.0.0.1")); + CHECK_EQUAL(port, 9400); + REQUIRE(parse_host_port("halo-0:1", host, port, &err)); + CHECK_EQUAL(host, std::string("halo-0")); + CHECK_EQUAL(port, 1); + CHECK_FALSE(parse_host_port("", host, port, &err)); + // A bare host takes the default control port. + REQUIRE(parse_host_port("10.0.0.1", host, port, &err)); + CHECK_EQUAL(port, kClusterDefaultPort); + CHECK_FALSE(parse_host_port(":9400", host, port, &err)); + CHECK_FALSE(parse_host_port("10.0.0.1:", host, port, &err)); + CHECK_FALSE(parse_host_port("10.0.0.1:abc", host, port, &err)); + CHECK_FALSE(parse_host_port("10.0.0.1:70000", host, port, &err)); + CHECK_FALSE(parse_host_port("10.0.0.1:0", host, port, &err)); + + SharedExpertMode se; + REQUIRE(parse_shared_expert_mode("replicate", se, &err)); CHECK(se == SharedExpertMode::Replicate); + REQUIRE(parse_shared_expert_mode("shard", se, &err)); CHECK(se == SharedExpertMode::Shard); + CHECK_FALSE(parse_shared_expert_mode("rank0", se, &err)); + CHECK_FALSE(parse_shared_expert_mode("bogus", se, &err)); + CHECK_FALSE(parse_shared_expert_mode("", se, &err)); + CHECK_EQUAL(std::string(shared_expert_mode_name(SharedExpertMode::Shard)), std::string("shard")); + + AllreduceDType dt; + REQUIRE(parse_allreduce_dtype("f32", dt, &err)); CHECK(dt == AllreduceDType::F32); + REQUIRE(parse_allreduce_dtype("bf16", dt, &err)); CHECK(dt == AllreduceDType::BF16); + REQUIRE(parse_allreduce_dtype("auto", dt, &err)); CHECK(dt == AllreduceDType::Auto); + CHECK_FALSE(parse_allreduce_dtype("fp16", dt, &err)); + CHECK_EQUAL(std::string(allreduce_dtype_name(AllreduceDType::BF16)), std::string("bf16")); + + PlacementSource ps; + std::string file; + REQUIRE(parse_placement_source("uniform", ps, file, &err)); + CHECK(ps == PlacementSource::Uniform); + CHECK(file.empty()); + REQUIRE(parse_placement_source("balanced", ps, file, &err)); + CHECK(ps == PlacementSource::Balanced); + REQUIRE(parse_placement_source("/tmp/placement.json", ps, file, &err)); + CHECK(ps == PlacementSource::File); + CHECK_EQUAL(file, std::string("/tmp/placement.json")); + CHECK_FALSE(parse_placement_source("random", ps, file, &err)); + CHECK_FALSE(parse_placement_source("placement.yaml", ps, file, &err)); + CHECK_FALSE(parse_placement_source("", ps, file, &err)); + CHECK_EQUAL(std::string(placement_source_name(PlacementSource::Balanced)), std::string("balanced")); +} + +// ─── Expert placement ─────────────────────────────────────────────────── + +TEST_CASE(ClusterUnitFixture, UniformPlacementExactlyOneOwnerPerSlot) { + for (int n : {2, 3, 4}) { + ClusterExpertPlacement p; + std::string err; + REQUIRE(ClusterExpertPlacement::build_uniform(n, ds4_cfg(), 0, nullptr, p, &err)); + REQUIRE(p.valid(&err)); + CHECK(p.matches(ds4_cfg())); + CHECK_EQUAL(p.n_ranks, n); + CHECK_EQUAL(p.n_layer, kLayers); + CHECK_EQUAL(p.n_expert, kExperts); + CHECK_EQUAL(p.owner.size(), (size_t) kLayers * kExperts); + + // owner_of follows e % N and nothing is replicated. + for (int l = 0; l < kLayers; l++) { + for (int e = 0; e < kExperts; e++) { + CHECK_EQUAL(p.owner_of(l, e), (int32_t) (e % n)); + CHECK_FALSE(p.is_replicated(l, e)); + } + } + + // Exactly one rank evaluates every (layer, expert, slot). + std::vector> validity((size_t) n); + for (int r = 0; r < n; r++) { + p.slot_validity(r, kSlots, validity[(size_t) r]); + REQUIRE_EQUAL(validity[(size_t) r].size(), (size_t) kLayers * kSlots * kExperts); + } + for (int l = 0; l < kLayers; l++) { + for (int s = 0; s < kSlots; s++) { + for (int e = 0; e < kExperts; e++) { + int owners = 0; + for (int r = 0; r < n; r++) { + owners += validity[(size_t) r][((size_t) l * kSlots + s) * kExperts + e] ? 1 : 0; + } + REQUIRE_EQUAL(owners, 1); + const int32_t so = p.slot_owner(l, e, s); + REQUIRE(so >= 0 && so < n); + REQUIRE(validity[(size_t) so][((size_t) l * kSlots + s) * kExperts + e] != 0); + } + } + } + + // Shard sizes within +-1 of n_expert / n per layer, and in total. + for (int l = 0; l < kLayers; l++) { + int sum = 0; + for (int r = 0; r < n; r++) { + const int c = p.owned_count(l, r); + CHECK(std::abs(c - kExperts / n) <= 1); + sum += c; + } + CHECK_EQUAL(sum, kExperts); + } + int total = 0; + for (int r = 0; r < n; r++) total += p.total_owned(r); + CHECK_EQUAL(total, kLayers * kExperts); + } +} + +TEST_CASE(ClusterUnitFixture, ReplicateHotMarksHottestAndKeepsOneOwnerPerSlot) { + const auto stats = skewed_stats(); + const int n = 4; + const int k = 8; + ClusterExpertPlacement p; + std::string err; + REQUIRE(ClusterExpertPlacement::build_uniform(n, ds4_cfg(), k, &stats, p, &err)); + REQUIRE(p.valid(&err)); + CHECK_EQUAL(p.replicate_hot, k); + + for (int l = 0; l < kLayers; l++) { + const std::vector hottest = stats.hot_experts(l, k); + REQUIRE_EQUAL(hottest.size(), (size_t) k); + int replicated = 0; + for (int e = 0; e < kExperts; e++) { + if (p.is_replicated(l, e)) { + replicated++; + CHECK(std::find(hottest.begin(), hottest.end(), e) != hottest.end()); + } + } + CHECK_EQUAL(replicated, k); + + // Replicated experts rotate over ranks by slot; owned ones do not. + for (int e : hottest) { + std::set seen; + for (int s = 0; s < kSlots; s++) { + const int32_t so = p.slot_owner(l, e, s); + CHECK_EQUAL(so, (int32_t) ((e + s) % n)); + seen.insert(so); + } + CHECK_EQUAL(seen.size(), (size_t) n); + } + } + + // One evaluator per (layer, expert, slot) still holds with replication. + std::vector> validity((size_t) n); + for (int r = 0; r < n; r++) p.slot_validity(r, kSlots, validity[(size_t) r]); + for (int l = 0; l < kLayers; l++) { + for (int s = 0; s < kSlots; s++) { + for (int e = 0; e < kExperts; e++) { + int owners = 0; + for (int r = 0; r < n; r++) { + owners += validity[(size_t) r][((size_t) l * kSlots + s) * kExperts + e] ? 1 : 0; + } + REQUIRE_EQUAL(owners, 1); + } + } + } + + // Resident sets: owned + all replicated, sorted ascending. + for (int r = 0; r < n; r++) { + const std::vector res = p.resident_experts(0, r); + CHECK(std::is_sorted(res.begin(), res.end())); + CHECK_EQUAL((int) res.size(), p.owned_count(0, r) + k); + } +} + +TEST_CASE(ClusterUnitFixture, BalancedPlacementBeatsUniformOnSkewedStats) { + const auto stats = skewed_stats(); + for (int n : {2, 4}) { + ClusterExpertPlacement uni, bal; + std::string err; + REQUIRE(ClusterExpertPlacement::build_uniform(n, ds4_cfg(), 0, nullptr, uni, &err)); + REQUIRE(ClusterExpertPlacement::build_balanced(n, stats, 0, bal, &err)); + REQUIRE(bal.valid(&err)); + CHECK(bal.matches(ds4_cfg())); + + double worst_uni = 0.0, worst_bal = 0.0; + for (int l = 0; l < kLayers; l++) { + const double u = max_over_mean(rank_loads(uni, stats, l)); + const double b = max_over_mean(rank_loads(bal, stats, l)); + worst_uni = std::max(worst_uni, u); + worst_bal = std::max(worst_bal, b); + CHECK(b <= u); + // Memory balance constraint: shard sizes stay within +-1. + for (int r = 0; r < n; r++) { + CHECK(std::abs(bal.owned_count(l, r) - kExperts / n) <= 1); + } + } + // Uniform sends every hot expert (e % 4 == 0) to rank 0. + CHECK(worst_uni > 1.5); + CHECK(worst_bal < worst_uni); + CHECK(worst_bal < 1.1); + CHECK(bal.hash() != uni.hash()); + } +} + +TEST_CASE(ClusterUnitFixture, PlacementJsonRoundTripAndHash) { + const auto stats = skewed_stats(); + ClusterExpertPlacement p; + std::string err; + REQUIRE(ClusterExpertPlacement::build_uniform(4, ds4_cfg(), 2, &stats, p, &err)); + const uint64_t h0 = p.hash(); + CHECK(h0 != 0u); + + const long long stamp = (long long) std::chrono::steady_clock::now().time_since_epoch().count(); + const std::filesystem::path path = + std::filesystem::temp_directory_path() / + ("cluster_placement_test_" + std::to_string(stamp) + ".json"); + REQUIRE(p.save_json(path.string(), &err)); + ClusterExpertPlacement q; + REQUIRE(ClusterExpertPlacement::load_json(path.string(), q, &err)); + std::filesystem::remove(path); + CHECK_EQUAL(q.n_ranks, p.n_ranks); + CHECK_EQUAL(q.n_layer, p.n_layer); + CHECK_EQUAL(q.n_expert, p.n_expert); + CHECK_EQUAL(q.n_expert_used, p.n_expert_used); + CHECK_EQUAL(q.replicate_hot, p.replicate_hot); + CHECK(q.owner == p.owner); + CHECK_EQUAL(q.hash(), h0); + + // Any ownership change changes the hash; the same content hashes equally. + ClusterExpertPlacement r = p; + CHECK_EQUAL(r.hash(), h0); + const size_t idx = (size_t) 5 * kExperts + 17; + r.owner[idx] = (r.owner[idx] + 1) % 4; + CHECK(r.hash() != h0); + r.owner[idx] = p.owner[idx]; + CHECK_EQUAL(r.hash(), h0); + + // Loading garbage fails cleanly. + ClusterExpertPlacement bad; + CHECK_FALSE(ClusterExpertPlacement::load_json("/nonexistent/dir/placement.json", bad, &err)); + CHECK_FALSE(err.empty()); + CHECK_FALSE(p.describe(&stats).empty()); +} + +TEST_CASE(ClusterUnitFixture, ToRankPlacementIsValidAndSorted) { + const auto stats = skewed_stats(); + for (int n : {2, 3, 4}) { + ClusterExpertPlacement p; + std::string err; + REQUIRE(ClusterExpertPlacement::build_uniform(n, ds4_cfg(), 4, &stats, p, &err)); + for (int r = 0; r < n; r++) { + luce::common::MoeHybridPlacement mp; + REQUIRE(p.to_rank_placement(r, mp, &err)); + REQUIRE(mp.valid(&err)); + CHECK(mp.matches(ds4_cfg())); + CHECK_EQUAL(mp.n_layer, kLayers); + CHECK_EQUAL(mp.n_expert, kExperts); + REQUIRE_EQUAL(mp.hot_expert_ids.size(), (size_t) kLayers); + REQUIRE_EQUAL(mp.hot_counts.size(), (size_t) kLayers); + int total = 0; + for (int l = 0; l < kLayers; l++) { + const auto & ids = mp.hot_expert_ids[(size_t) l]; + CHECK(std::is_sorted(ids.begin(), ids.end())); + CHECK(std::adjacent_find(ids.begin(), ids.end()) == ids.end()); + CHECK_EQUAL((int) ids.size(), mp.hot_counts[(size_t) l]); + CHECK(ids == p.resident_experts(l, r)); + for (int32_t e : ids) CHECK(mp.is_hot(l, e)); + total += (int) ids.size(); + } + CHECK_EQUAL(mp.total_hot, total); + } + luce::common::MoeHybridPlacement out_of_range; + CHECK_FALSE(p.to_rank_placement(n, out_of_range, &err)); + CHECK_FALSE(err.empty()); + } +} + +TEST_CASE(ClusterUnitFixture, PlacementBuildersRejectBadInput) { + ClusterExpertPlacement p; + std::string err; + CHECK_FALSE(ClusterExpertPlacement::build_uniform(0, ds4_cfg(), 0, nullptr, p, &err)); + CHECK_FALSE(ClusterExpertPlacement::build_uniform(2, ds4_cfg(), -1, nullptr, p, &err)); + // n_ranks == 1 is a legal degenerate placement (loopback comm, size-1 runs). + ClusterExpertPlacement single; + CHECK(ClusterExpertPlacement::build_uniform(1, ds4_cfg(), 0, nullptr, single, &err)); + CHECK_EQUAL(single.owned_count(0, 0), kExperts); + luce::common::MoeHybridConfig empty; + CHECK_FALSE(ClusterExpertPlacement::build_uniform(2, empty, 0, nullptr, p, &err)); + // replicate_hot without stats cannot pick "hottest": either refused or zero replicated. + ClusterExpertPlacement q; + if (ClusterExpertPlacement::build_uniform(2, ds4_cfg(), 3, nullptr, q, &err)) { + int replicated = 0; + for (int32_t o : q.owner) replicated += (o == kReplicated) ? 1 : 0; + CHECK_EQUAL(replicated, 0); + } + // replicate_hot >= n_expert leaves nothing to shard. + const auto stats = skewed_stats(); + CHECK_FALSE(ClusterExpertPlacement::build_balanced(2, stats, kExperts, p, &err)); + luce::common::MoeHybridRoutingStats none; + CHECK_FALSE(ClusterExpertPlacement::build_balanced(2, none, 0, p, &err)); +} + +// ─── Decision hooks ───────────────────────────────────────────────────── + +TEST_CASE(ClusterUnitFixture, LocalHooksAreTransparent) { + LocalHooks h; + CHECK(h.is_head()); + CHECK_FALSE(h.is_cluster()); + int32_t token = 123; + uint8_t flags = kDecisionEos; + luce::common::SamplerCfg cfg; + std::string err; + const float logits[3] = {0.f, 1.f, 0.f}; + CHECK(h.decide_next_token(1, 0, logits, 3, cfg, false, token, flags, &err)); + CHECK_EQUAL(token, 123); + CHECK_EQUAL(flags, (uint8_t) kDecisionEos); + std::vector draft = {1, 2}; + CHECK(h.decide_draft(1, 0, 0, draft, &err)); + CHECK_EQUAL(draft.size(), (size_t) 2); + int32_t accept = 2, bonus = 9; + CHECK(h.decide_accept(1, 0, accept, bonus, flags, &err)); + CHECK_EQUAL(accept, 2); + CHECK_EQUAL(bonus, 9); + CHECK(h.backend_op(BackendOpMsg{}, &err)); + h.set_current_request(77); + CHECK_EQUAL(h.current_request(), 77u); + CHECK_EQUAL(h.counters().decisions, 0u); +} + +#if !defined(_WIN32) + +namespace { + +// Head and worker on 127.0.0.1. The head picks a port from a small range +// (no ephemeral-port query on ClusterHeadControl); the worker retries. +struct LoopbackPair { + ClusterHeadControl head; + ClusterWorkerControl worker; + HelloMsg head_id; + HelloMsg worker_id; + WelcomeMsg welcome; + std::vector hellos; + int port = 0; + std::string head_err, worker_err; + bool ok = false; + + LoopbackPair(const std::string & build = "test-build", uint64_t placement = 0x1234) { + head_id.rank = 0; head_id.size = 2; head_id.build_sha = build; + head_id.model_sha = "model"; head_id.placement_hash = placement; head_id.hostname = "head"; + worker_id = head_id; worker_id.rank = 1; worker_id.hostname = "worker"; + + ControlTimeouts t; + t.connect_ms = 1000; + t.connect_retries = 20; + t.recv_ms = 5000; + head.set_timeouts(t); + worker.set_timeouts(t); + + const int base = 39400 + (int) (getpid() % 500); + for (int attempt = 0; attempt < 20 && port == 0; attempt++) { + if (head.listen("127.0.0.1", base + attempt, &head_err)) port = base + attempt; + } + } + + // Runs the handshake with the worker on a second thread. + bool handshake(const HelloMsg & wid) { + if (port == 0) return false; + WelcomeMsg send; + send.size = 2; + send.timeout_ms = 5000; + send.rank_hostnames = {"head", "worker"}; + std::atomic w_ok{false}; + std::thread wt([&] { + w_ok = worker.connect_and_handshake("127.0.0.1", port, wid, welcome, &worker_err); + }); + const bool h_ok = head.accept_workers(1, head_id, 10000, hellos, &head_err); + bool sent = false; + if (h_ok) sent = head.broadcast(make_frame(send), &head_err); + wt.join(); + ok = h_ok && sent && w_ok; + return ok; + } +}; + +} // namespace + +TEST_CASE(ClusterUnitFixture, ControlHandshakeAcceptsMatchingIdentity) { + LoopbackPair lp; + if (lp.port == 0) SKIP("no loopback port available"); + REQUIRE(lp.handshake(lp.worker_id)); + REQUIRE_EQUAL(lp.hellos.size(), (size_t) 1); + CHECK_EQUAL(lp.hellos[0].rank, 1); + CHECK_EQUAL(lp.hellos[0].hostname, std::string("worker")); + CHECK_EQUAL(lp.welcome.size, 2); + CHECK_EQUAL(lp.welcome.timeout_ms, 5000u); + CHECK_EQUAL(lp.head.n_workers(), 1); + CHECK(lp.head.all_alive()); + CHECK(lp.worker.alive()); + lp.worker.close(); + lp.head.close(); +} + +TEST_CASE(ClusterUnitFixture, ControlHandshakeRefusesMismatch) { + SECTION("placement hash") { + LoopbackPair lp; + if (lp.port == 0) SKIP("no loopback port available"); + HelloMsg wid = lp.worker_id; + wid.placement_hash ^= 1; + CHECK_FALSE(lp.handshake(wid)); + CHECK_FALSE(lp.head_err.empty()); + } + SECTION("build sha") { + LoopbackPair lp; + if (lp.port == 0) SKIP("no loopback port available"); + HelloMsg wid = lp.worker_id; + wid.build_sha = "other"; + CHECK_FALSE(lp.handshake(wid)); + } + SECTION("rank out of range") { + LoopbackPair lp; + if (lp.port == 0) SKIP("no loopback port available"); + HelloMsg wid = lp.worker_id; + wid.rank = 5; + CHECK_FALSE(lp.handshake(wid)); + } +} + +TEST_CASE(ClusterUnitFixture, HeadAndWorkerHooksOverLoopback) { + LoopbackPair lp; + if (lp.port == 0) SKIP("no loopback port available"); + REQUIRE(lp.handshake(lp.worker_id)); + + HeadHooks head(lp.head); + WorkerHooksConfig wc; + wc.timeout_ms = 5000; + WorkerHooks worker(lp.worker, wc); + CHECK(head.is_head()); CHECK(head.is_cluster()); + CHECK_FALSE(worker.is_head()); CHECK(worker.is_cluster()); + + const uint64_t req = 42; + head.set_current_request(req); + worker.set_current_request(req); + luce::common::SamplerCfg scfg; + const float logits[4] = {0.1f, 0.9f, 0.2f, 0.3f}; + + SECTION("decision, draft, accept and backend op") { + std::string w_err; + int32_t w_tok = -1; uint8_t w_flags = 0; + std::vector w_draft; + int32_t w_accept = -1, w_bonus = -1; uint8_t w_aflags = 0; + bool w_dec = false, w_dr = false, w_ac = false; + std::thread wt([&] { + w_dec = worker.decide_next_token(req, 0, nullptr, 0, scfg, false, w_tok, w_flags, &w_err); + w_dr = worker.decide_draft(req, 1, 100, w_draft, &w_err); + w_ac = worker.decide_accept(req, 1, w_accept, w_bonus, w_aflags, &w_err); + }); + + std::string h_err; + int32_t tok = 7; uint8_t flags = make_decision_flags(true, false, false, false, true); + REQUIRE(head.decide_next_token(req, 0, logits, 4, scfg, true, tok, flags, &h_err)); + CHECK_EQUAL(tok, 7); // caller's token is passed through untouched + std::vector draft = {7, 8, 9, 10}; + REQUIRE(head.decide_draft(req, 1, 100, draft, &h_err)); + int32_t accept = 3, bonus = 11; uint8_t aflags = 0; + REQUIRE(head.decide_accept(req, 1, accept, bonus, aflags, &h_err)); + wt.join(); + + CHECK(w_dec); CHECK_EQUAL(w_tok, 7); CHECK_EQUAL(w_flags, flags); + CHECK(w_dr); CHECK(w_draft == draft); + CHECK(w_ac); CHECK_EQUAL(w_accept, 3); CHECK_EQUAL(w_bonus, 11); + CHECK_EQUAL(head.counters().decisions, 1u); + CHECK_EQUAL(head.counters().drafts, 1u); + CHECK_EQUAL(head.counters().accepts, 1u); + CHECK_EQUAL(worker.counters().decisions, 1u); + + // Backend op reaches the worker as a plain frame (the worker main + // loop executes it; WorkerHooks::backend_op is a no-op). + BackendOpMsg op; + op.kind = BackendOpKind::SnapshotFree; + op.args = {5}; + REQUIRE(head.backend_op(op, &h_err)); + Frame f; + REQUIRE(lp.worker.recv(f, 5000, &w_err)); + BackendOpMsg got; + REQUIRE(parse_frame(f, got)); + CHECK(got.kind == BackendOpKind::SnapshotFree); + CHECK(worker.backend_op(got, &w_err)); + CHECK_EQUAL(head.counters().backend_ops, 1u); + + // Rank 0 samples in its decode loop: without a token the head + // broadcasts nothing. + int32_t none = -1; uint8_t sflags = 0; + h_err.clear(); + CHECK_FALSE(head.decide_next_token(req, 1, logits, 4, scfg, false, none, sflags, &h_err)); + CHECK_FALSE(h_err.empty()); + CHECK_EQUAL(head.counters().decisions, 1u); + } + + SECTION("worker rejects wrong request or step") { + std::string w_err; + int32_t w_tok = -1; uint8_t w_flags = 0; + bool w_ok = true; + std::thread wt([&] { + w_ok = worker.decide_next_token(req, 3, nullptr, 0, scfg, false, w_tok, w_flags, &w_err); + }); + std::string h_err; + int32_t tok = 1; uint8_t flags = 0; + REQUIRE(head.decide_next_token(req, 2, logits, 4, scfg, false, tok, flags, &h_err)); // step 2 != 3 + wt.join(); + CHECK_FALSE(w_ok); + CHECK_FALSE(w_err.empty()); + } + + SECTION("unexpected RequestEnd is parked, Abort is fatal") { + std::string w_err; + int32_t w_tok = -1; uint8_t w_flags = 0; + bool w_ok = true; + std::thread wt([&] { + w_ok = worker.decide_next_token(req, 0, nullptr, 0, scfg, false, w_tok, w_flags, &w_err); + }); + RequestEndMsg end; + end.request_id = req; + end.status = 6; + std::string h_err; + REQUIRE(lp.head.broadcast(make_frame(end), &h_err)); + wt.join(); + CHECK_FALSE(w_ok); + Frame parked; + REQUIRE(worker.take_pending_frame(parked)); + CHECK(parked.type == MsgType::RequestEnd); + CHECK_FALSE(worker.take_pending_frame(parked)); + + bool w_ok2 = true; + std::thread wt2([&] { + std::vector d; + w_ok2 = worker.decide_draft(req, 0, 0, d, &w_err); + }); + AbortMsg abort; + abort.rank = 0; + abort.code = 1; + abort.reason = "test abort"; + REQUIRE(lp.head.broadcast(make_frame(abort), &h_err)); + wt2.join(); + CHECK_FALSE(w_ok2); + CHECK(worker.aborted()); + CHECK_EQUAL(worker.abort_reason(), std::string("test abort")); + // Once aborted every further call fails immediately without blocking. + int32_t a = 0, b = 0; uint8_t fl = 0; + CHECK_FALSE(worker.decide_accept(req, 0, a, b, fl, &w_err)); + } + + lp.worker.close(); + lp.head.close(); +} + +// An empty draft block is the head's end marker for a request only it +// can see ending (client cancel). It must survive the wire as an empty +// vector, and it must clear whatever the worker had drafted locally. +TEST_CASE(ClusterUnitFixture, EmptyDraftBlockIsTheEndMarker) { + LoopbackPair lp; + if (lp.port == 0) SKIP("no loopback port available"); + REQUIRE(lp.handshake(lp.worker_id)); + + HeadHooks head(lp.head); + WorkerHooksConfig wc; wc.timeout_ms = 5000; + WorkerHooks worker(lp.worker, wc); + + const uint64_t req = 9; + head.set_current_request(req); + worker.set_current_request(req); + + std::string w_err; + // The worker drafted its own block; the head's end marker must replace it. + std::vector w_draft = {5, 6, 7}; + bool w_ok = false; + std::thread wt([&] { w_ok = worker.decide_draft(req, 3, 128, w_draft, &w_err); }); + + std::string h_err; + std::vector end_marker; + REQUIRE(head.decide_draft(req, 3, 128, end_marker, &h_err)); + wt.join(); + + CHECK(w_ok); + CHECK(w_draft.empty()); + CHECK_EQUAL(head.counters().drafts, 1u); + CHECK_EQUAL(worker.counters().drafts, 1u); +} + +// The end marker also has to round-trip through the frame codec on its own: +// a zero-length token vector is a valid Draft, not a malformed one. +TEST_CASE(ClusterUnitFixture, EmptyDraftMsgRoundTrip) { + DraftMsg msg; + msg.request_id = 77; + msg.step = 12; + msg.pos = 4096; + const Frame f = make_frame(msg); + DraftMsg got; + REQUIRE(parse_frame(f, got)); + CHECK_EQUAL(got.request_id, 77u); + CHECK_EQUAL(got.step, 12u); + CHECK_EQUAL(got.pos, 4096); + CHECK(got.tokens.empty()); +} + +TEST_CASE(ClusterUnitFixture, HeadGatherFailsWhenWorkerGone) { + LoopbackPair lp; + if (lp.port == 0) SKIP("no loopback port available"); + REQUIRE(lp.handshake(lp.worker_id)); + lp.worker.close(); + + std::vector frames; + std::string err; + // A gather with no live worker must return false within the deadline. + CHECK_FALSE(lp.head.gather(MsgType::RequestReport, 500, frames, &err)); + CHECK_FALSE(err.empty()); + lp.head.close(); +} + +#endif // !_WIN32 diff --git a/server/test/unit/test_deepseek4_unit.cpp b/server/test/unit/test_deepseek4_unit.cpp index fa9d60f19..539603b88 100644 --- a/server/test/unit/test_deepseek4_unit.cpp +++ b/server/test/unit/test_deepseek4_unit.cpp @@ -6367,6 +6367,181 @@ static void test_ds4_flash_attention_segmented_kv_gpu(bool benchmark = false) { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +// Lane rows (ggml_flash_attn_ext_set_ds4_lane_rows): one split-KV op over +// several single-token lanes of one raw ring must equal, byte for byte, one +// op per lane over that lane's own ring state (the old ring with the rows of +// lanes 0..t written), with per-lane masks and selections. +static bool run_ds4_lane_rows_case(ggml_backend_t backend, int lanes, int first_row, + int compressed_rows, bool identity, uint32_t seed) { + constexpr int head_dim = 512; + constexpr int n_heads = 64; + constexpr int raw_rows = 128; + const int selected = identity ? compressed_rows : 512; + const int n_kv = raw_rows + compressed_rows; + ggml_context * ctx = make_test_context(8u << 20); + if (!ctx) return false; + + ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_dim, n_heads, lanes); + ggml_tensor * ring = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, head_dim, raw_rows); + ggml_tensor * lane_kv = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, head_dim, lanes); + ggml_tensor * ring_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, lanes); + ggml_tensor * comp = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, head_dim, compressed_rows); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, n_kv, lanes); + ggml_tensor * topk = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, selected, lanes); + ggml_tensor * sinks = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_heads); + ggml_tensor * comp_head = ggml_view_2d(ctx, comp, head_dim, compressed_rows - 1, comp->nb[1], 0); + ggml_tensor * comp_tail = ggml_view_2d(ctx, comp, head_dim, 1, comp->nb[1], + (size_t) (compressed_rows - 1) * comp->nb[1]); + const float scale = 1.0f / std::sqrt((float) head_dim); + auto make_op = [&](ggml_tensor * q_lanes, ggml_tensor * k, ggml_tensor * m, ggml_tensor * sel) { + ggml_tensor * op = ggml_flash_attn_ext( + ctx, ggml_permute(ctx, q_lanes, 0, 2, 1, 3), k, k, m, scale, 0.0f, 0.0f); + ggml_flash_attn_ext_add_sinks(op, sinks); + ggml_flash_attn_ext_set_prec(op, GGML_PREC_F32); + ggml_flash_attn_ext_set_ds4_sparse(op, raw_rows, raw_rows, -selected, 32); + ggml_flash_attn_ext_set_ds4_kv_segments(op, comp_head, comp_tail); + ggml_flash_attn_ext_set_ds4_indexer_topk(op, sel); + return op; + }; + ggml_tensor * batched = make_op(q, ring, mask, topk); + ggml_flash_attn_ext_set_ds4_lane_rows(batched, ring_rows, lane_kv); + ggml_set_output(batched); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 256, false); + ggml_build_forward_expand(graph, batched); + std::vector lane_q(lanes), lane_ring(lanes), lane_mask(lanes), + lane_topk(lanes), lane_ref(lanes); + bool supported = ggml_backend_supports_op(backend, batched); + for (int t = 0; t < lanes; ++t) { + lane_q[t] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_dim, n_heads, 1); + lane_ring[t] = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, head_dim, raw_rows); + lane_mask[t] = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, n_kv, 1); + lane_topk[t] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, selected, 1); + lane_ref[t] = make_op(lane_q[t], lane_ring[t], lane_mask[t], lane_topk[t]); + ggml_set_output(lane_ref[t]); + supported = supported && ggml_backend_supports_op(backend, lane_ref[t]); + ggml_build_forward_expand(graph, lane_ref[t]); + } + if (!supported) { + std::fprintf(stderr, " [lane rows unsupported]"); + ggml_free(ctx); + return false; + } + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + bool ok = ggml_gallocr_alloc_graph(alloc, graph); + if (ok) { + TestLcg rng(seed); + std::vector q_data((size_t) head_dim * n_heads * lanes); + for (float & v : q_data) v = 0.5f * rng.next(); + std::vector ring_data((size_t) head_dim * raw_rows); + for (ggml_fp16_t & v : ring_data) v = ggml_fp32_to_fp16(0.5f * rng.next()); + std::vector lane_data((size_t) head_dim * lanes); + for (ggml_fp16_t & v : lane_data) v = ggml_fp32_to_fp16(0.5f * rng.next()); + std::vector comp_data((size_t) head_dim * compressed_rows); + for (ggml_fp16_t & v : comp_data) v = ggml_fp32_to_fp16(0.5f * rng.next()); + std::vector sink_data(n_heads); + for (float & v : sink_data) v = rng.next(); + std::vector rows(lanes); + for (int j = 0; j < lanes; ++j) rows[j] = (first_row + j) % raw_rows; + // Before the ring wraps, rows past the lane's own are unwritten. + const bool wrapped = first_row + lanes > raw_rows || first_row >= raw_rows - 1; + std::vector mask_data((size_t) n_kv * lanes, ggml_fp32_to_fp16(0.0f)); + std::vector topk_data((size_t) selected * lanes); + for (int t = 0; t < lanes; ++t) { + ggml_fp16_t * m = mask_data.data() + (size_t) t * n_kv; + if (!wrapped) { + for (int r = rows[t] + 1; r < raw_rows; ++r) m[r] = ggml_fp32_to_fp16(-1.0e30f); + } + const int visible = compressed_rows - (lanes - 1 - t); + for (int c = visible; c < compressed_rows; ++c) { + m[raw_rows + c] = ggml_fp32_to_fp16(-1.0e30f); + } + int32_t * s = topk_data.data() + (size_t) t * selected; + for (int i = 0; i < selected; ++i) { + s[i] = identity ? i : (int32_t) (((int64_t) i * 7919 + t * 131) % compressed_rows); + } + } + ggml_backend_tensor_set(q, q_data.data(), 0, q_data.size() * sizeof(float)); + ggml_backend_tensor_set(ring, ring_data.data(), 0, ring_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(lane_kv, lane_data.data(), 0, lane_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(ring_rows, rows.data(), 0, rows.size() * sizeof(int32_t)); + ggml_backend_tensor_set(comp, comp_data.data(), 0, comp_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, mask_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(topk, topk_data.data(), 0, topk_data.size() * sizeof(int32_t)); + ggml_backend_tensor_set(sinks, sink_data.data(), 0, sink_data.size() * sizeof(float)); + for (int t = 0; t < lanes; ++t) { + ggml_backend_tensor_set(lane_q[t], q_data.data() + (size_t) t * head_dim * n_heads, 0, + (size_t) head_dim * n_heads * sizeof(float)); + std::vector own = ring_data; + for (int j = 0; j <= t; ++j) { + std::memcpy(own.data() + (size_t) rows[j] * head_dim, + lane_data.data() + (size_t) j * head_dim, head_dim * sizeof(ggml_fp16_t)); + } + ggml_backend_tensor_set(lane_ring[t], own.data(), 0, own.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(lane_mask[t], mask_data.data() + (size_t) t * n_kv, 0, + (size_t) n_kv * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(lane_topk[t], topk_data.data() + (size_t) t * selected, 0, + (size_t) selected * sizeof(int32_t)); + } + ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + if (ok) { + std::vector all((size_t) head_dim * n_heads * lanes); + ggml_backend_tensor_get(batched, all.data(), 0, all.size() * sizeof(float)); + for (int t = 0; t < lanes && ok; ++t) { + std::vector one((size_t) head_dim * n_heads); + ggml_backend_tensor_get(lane_ref[t], one.data(), 0, one.size() * sizeof(float)); + ok = std::memcmp(one.data(), all.data() + (size_t) t * head_dim * n_heads, + one.size() * sizeof(float)) == 0; + if (!ok) { + std::fprintf(stderr, " [lanes=%d first=%d comp=%d %s: lane %d differs]", + lanes, first_row, compressed_rows, + identity ? "identity" : "top-k", t); + } + } + } + } + ggml_gallocr_free(alloc); + ggml_free(ctx); + return ok; +} + +static void test_ds4_flash_attention_lane_rows_gpu() { + std::fprintf(stderr, " test_ds4_flash_attention_lane_rows_gpu ..."); +#if !defined(GGML_USE_HIP) + std::fprintf(stderr, " skipped (HIP-only contract)\n"); + return; +#endif + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, " skipped (no GPU backend)\n"); + return; + } + ScopedCudaGraphOverrides eager( + /*disable_graphs=*/true, /*mmvq_max_ncols=*/0, /*skip_property_check=*/false); + ScopedEnvVar split_kv_guard("GGML_CUDA_MLA_SPLIT_KV"); + ScopedEnvVar split_count_guard("GGML_CUDA_MLA_SPLIT_KV_COUNT"); + setenv("GGML_CUDA_MLA_SPLIT_KV", "1", 1); + int cases = 0; + int passed = 0; + for (const char * splits : {"4", "8"}) { + setenv("GGML_CUDA_MLA_SPLIT_KV_COUNT", splits, 1); + for (int lanes : {2, 5}) { + for (int first_row : {3, 125, 126}) { + for (bool identity : {true, false}) { + const int compressed_rows = identity ? 384 : 1408; + ++cases; + passed += run_ds4_lane_rows_case( + backend, lanes, first_row, compressed_rows, identity, + 0x1234567u + (uint32_t) (cases * 977)) ? 1 : 0; + } + } + } + } + std::fprintf(stderr, " %d/%d cases", passed, cases); + TEST_ASSERT_MSG(passed == cases, "lane-row attention differs from per-lane attention"); + ggml_backend_free(backend); + std::fprintf(stderr, passed == cases ? " ok\n" : " FAIL\n"); +} + static void test_ds4_flash_attention_ratio4_maskless_gpu() { std::fprintf(stderr, " test_ds4_flash_attention_ratio4_maskless_gpu ..."); @@ -9099,6 +9274,10 @@ int main(int argc, char ** argv) { test_ds4_flash_attention_segmented_kv_gpu(); return g_failures == 0 ? 0 : 1; } + if (argc == 2 && std::strcmp(argv[1], "--test-lane-rows") == 0) { + test_ds4_flash_attention_lane_rows_gpu(); + return g_failures == 0 ? 0 : 1; + } if (argc == 2 && std::strcmp(argv[1], "--bench-segmented-kv") == 0) { test_ds4_flash_attention_segmented_kv_gpu(/*benchmark=*/true); return g_failures == 0 ? 0 : 1; @@ -9249,6 +9428,7 @@ int main(int argc, char ** argv) { test_ds4_topk_block_radix_gpu(ncols); } test_ds4_flash_attention_segmented_kv_gpu(); + test_ds4_flash_attention_lane_rows_gpu(); test_ds4_flash_attention_ratio4_maskless_gpu(); test_ds4_flash_attention_sparse_value_skip_gpu(); test_ds4_indexer_score_m32_gpu(); diff --git a/server/test/unit/test_ds4_fused_ops_cuda.cpp b/server/test/unit/test_ds4_fused_ops_cuda.cpp index 2863ca4b9..d2c5da07f 100644 --- a/server/test/unit/test_ds4_fused_ops_cuda.cpp +++ b/server/test/unit/test_ds4_fused_ops_cuda.cpp @@ -3,7 +3,11 @@ // top_k, protected routes, get_rows, sum_rows, clamp, div, scale: the same // expert ids in the same order and bit-identical weights; // - ggml_ds4_hc_collapse vs mul, permute, cont, sum_rows: bit-identical. -// So LUCE_DS4_FUSE_ROUTER / LUCE_DS4_FUSE_COLLAPSE never change a token. +// - the HC sub-block boundary the CUDA backend runs as one launch +// (rms_norm, F16 fn mix, hc_pre, hc_post, collapse, rms_norm * w) vs the +// same chain kept unfused: bit-identical streams and normalized input. +// So LUCE_DS4_FUSE_ROUTER / LUCE_DS4_FUSE_COLLAPSE / the boundary fusion +// never change a token. // // Exit codes: 0 pass, 1 mismatch, 77 no GPU backend. #include "ggml-alloc.h" @@ -11,6 +15,7 @@ #include "ggml-cuda.h" #include "ggml.h" +#include #include #include #include @@ -173,6 +178,739 @@ bool run_collapse(ggml_backend_t backend, const CollapseCase & c, uint32_t seed) return ok; } +struct BoundaryCase { + int n_embd; + int tokens; + bool pre_read; // something outside the chain reads hc_pre's output +}; + +// One boundary chain over the given inputs. `keep` marks the first rms_norm +// as a graph output, which keeps the backend from fusing the chain. +struct BoundaryChain { + ggml_tensor * pre = nullptr; + ggml_tensor * post = nullptr; + ggml_tensor * out = nullptr; + ggml_tensor * proj = nullptr; // a Q8_0 projection of out (reads its q8_1 form) +}; + +BoundaryChain build_boundary(ggml_context * ctx, ggml_tensor * R, ggml_tensor * block_out, + ggml_tensor * fn, ggml_tensor * base, ggml_tensor * w, ggml_tensor * wq, + int n_embd, int n_hc, int tokens, bool keep, bool pre_read) { + const int mix_dim = 2 * n_hc + n_hc * n_hc; + BoundaryChain c; + ggml_tensor * n1 = ggml_rms_norm(ctx, R, 1e-6f); + if (keep) ggml_set_output(n1); + ggml_tensor * mix = ggml_reshape_2d(ctx, ggml_mul_mat(ctx, fn, n1), mix_dim, tokens); + c.pre = ggml_ds4_hc_pre(ctx, mix, base, R, n_hc, 20, 0.5f, 0.75f, 1.25f); + ggml_tensor * split = ggml_view_2d(ctx, c.pre, mix_dim, tokens, + tokens > 1 ? c.pre->nb[1] : ggml_nbytes(c.pre), + (size_t) n_embd * sizeof(float)); + c.post = ggml_ds4_hc_post(ctx, R, block_out, split, n_hc); // 1-D at one token, as in the graph + ggml_tensor * cpre = ggml_view_2d(ctx, split, n_hc, tokens, split->nb[1], 0); + ggml_tensor * coll = ggml_ds4_hc_collapse(ctx, c.post, cpre, n_hc); + c.out = ggml_mul(ctx, ggml_rms_norm(ctx, coll, 1e-6f), w); + // The reference reads out through a view: no producer hands it a q8_1 + // form, so MMVQ quantizes out itself. + c.proj = ggml_mul_mat(ctx, wq, keep ? ggml_view_2d(ctx, c.out, c.out->ne[0], c.out->ne[1], c.out->nb[1], 0) + : c.out); + if (pre_read) ggml_set_output(c.pre); + ggml_set_output(c.post); + ggml_set_output(c.out); + ggml_set_output(c.proj); + return c; +} + +bool run_boundary(ggml_backend_t backend, const BoundaryCase & c, uint32_t seed) { + const int n_hc = 4; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + const int64_t ncols = (int64_t) c.n_embd * n_hc; + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 64 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + ggml_tensor * R = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ncols, c.tokens); + ggml_tensor * block_out = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, c.n_embd, c.tokens); + ggml_tensor * fn = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, ncols, mix_dim); + ggml_tensor * base = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mix_dim); + ggml_tensor * w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, c.n_embd); + const int n_proj = 64; + ggml_tensor * wq = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, c.n_embd, n_proj); + for (ggml_tensor * t : {R, block_out, fn, base, w, wq}) ggml_set_input(t); + + const BoundaryChain ref = build_boundary(ctx, R, block_out, fn, base, w, wq, c.n_embd, n_hc, c.tokens, true, c.pre_read); + const BoundaryChain fused = build_boundary(ctx, R, block_out, fn, base, w, wq, c.n_embd, n_hc, c.tokens, false, c.pre_read); + ggml_cgraph * gf = ggml_new_graph(ctx); + for (ggml_tensor * t : {ref.post, ref.out, ref.proj, fused.post, fused.out, fused.proj}) ggml_build_forward_expand(gf, t); + if (c.pre_read) { + ggml_build_forward_expand(gf, ref.pre); + ggml_build_forward_expand(gf, fused.pre); + } + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_R((size_t) ncols * c.tokens), h_bo((size_t) c.n_embd * c.tokens); + std::vector h_base(mix_dim), h_w(c.n_embd); + std::vector h_fn((size_t) ncols * mix_dim); + for (float & v : h_R) v = normal(rng) * 4.0f; + for (float & v : h_bo) v = normal(rng) * 2.0f; + for (float & v : h_base) v = normal(rng) * 0.5f; + for (float & v : h_w) v = 1.0f + 0.1f * normal(rng); + for (ggml_fp16_t & v : h_fn) v = ggml_fp32_to_fp16(normal(rng) * 0.02f); + std::vector h_wf((size_t) c.n_embd * n_proj); + for (float & v : h_wf) v = normal(rng) * 0.05f; + std::vector h_wq(ggml_nbytes(wq)); + ggml_quantize_chunk(GGML_TYPE_Q8_0, h_wf.data(), h_wq.data(), 0, n_proj, c.n_embd, nullptr); + ggml_backend_tensor_set(wq, h_wq.data(), 0, h_wq.size()); + ggml_backend_tensor_set(R, h_R.data(), 0, sizeof(float) * h_R.size()); + ggml_backend_tensor_set(block_out, h_bo.data(), 0, sizeof(float) * h_bo.size()); + ggml_backend_tensor_set(base, h_base.data(), 0, sizeof(float) * h_base.size()); + ggml_backend_tensor_set(w, h_w.data(), 0, sizeof(float) * h_w.size()); + ggml_backend_tensor_set(fn, h_fn.data(), 0, sizeof(ggml_fp16_t) * h_fn.size()); + + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + auto same = [&](ggml_tensor * a, ggml_tensor * b, const char * what) { + const size_t n = (size_t) ggml_nelements(a); + std::vector va(n), vb(n); + ggml_backend_tensor_get(a, va.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(b, vb.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n; ++i) { + if (std::memcmp(&va[i], &vb[i], sizeof(float)) != 0) { + std::fprintf(stderr, "boundary %s mismatch n_embd=%d tokens=%d pre_read=%d at %zu: ref %.9g fused %.9g\n", + what, c.n_embd, c.tokens, c.pre_read ? 1 : 0, i, va[i], vb[i]); + return false; + } + } + return true; + }; + if (ok && c.pre_read) ok = same(ref.pre, fused.pre, "pre"); + if (ok) ok = same(ref.post, fused.post, "post") && same(ref.out, fused.out, "out") && + same(ref.proj, fused.proj, "proj"); + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +// rms_norm * w read by a quantized matmul: the fused norm writes the q8_1 +// form MMVQ would quantize; the reference matmul reads through a view. +bool run_norm_q8(ggml_backend_t backend, int ncols, int tokens, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + const int n_proj = 96; + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ncols, tokens); + ggml_tensor * w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ncols); + ggml_tensor * wq = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, ncols, n_proj); + for (ggml_tensor * t : {x, w, wq}) ggml_set_input(t); + ggml_tensor * ref_n = ggml_mul(ctx, ggml_rms_norm(ctx, x, 1e-6f), w); + ggml_tensor * ref = ggml_mul_mat(ctx, wq, ggml_view_2d(ctx, ref_n, ncols, tokens, ref_n->nb[1], 0)); + ggml_tensor * fus_n = ggml_mul(ctx, ggml_rms_norm(ctx, x, 1e-6f), w); + ggml_tensor * fus = ggml_mul_mat(ctx, wq, fus_n); + for (ggml_tensor * t : {ref_n, ref, fus_n, fus}) ggml_set_output(t); + ggml_cgraph * gf = ggml_new_graph(ctx); + for (ggml_tensor * t : {ref_n, ref, fus_n, fus}) ggml_build_forward_expand(gf, t); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_x((size_t) ncols * tokens), h_w(ncols), h_wf((size_t) ncols * n_proj); + for (float & v : h_x) v = normal(rng) * 3.0f; + for (float & v : h_w) v = 1.0f + 0.2f * normal(rng); + for (float & v : h_wf) v = normal(rng) * 0.05f; + std::vector h_wq(ggml_nbytes(wq)); + ggml_quantize_chunk(GGML_TYPE_Q8_0, h_wf.data(), h_wq.data(), 0, n_proj, ncols, nullptr); + ggml_backend_tensor_set(x, h_x.data(), 0, sizeof(float) * h_x.size()); + ggml_backend_tensor_set(w, h_w.data(), 0, sizeof(float) * h_w.size()); + ggml_backend_tensor_set(wq, h_wq.data(), 0, h_wq.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + for (auto pr : {std::make_pair(ref_n, fus_n), std::make_pair(ref, fus)}) { + if (!ok) break; + const size_t n = (size_t) ggml_nelements(pr.first); + std::vector a(n), b(n); + ggml_backend_tensor_get(pr.first, a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(pr.second, b.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n; ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + std::fprintf(stderr, "norm q8 mismatch ncols=%d tokens=%d at %zu: ref %.9g fused %.9g\n", + ncols, tokens, i, a[i], b[i]); + ok = false; + break; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +// ggml_moe_combine_masked vs the owner remap's mul(weights, get_rows(valid, +// ids)), mul(experts, masked) and repeat_back. +bool run_masked_combine(ggml_backend_t backend, int n_embd, int n_used, int n_expert, int tokens, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + ggml_tensor * experts = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_used, tokens); + ggml_tensor * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_used, tokens); + ggml_tensor * valid = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, n_expert, tokens, 1); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_used, tokens); + for (ggml_tensor * t : {experts, weights, valid, ids}) ggml_set_input(t); + ggml_tensor * v = ggml_reshape_2d(ctx, ggml_get_rows(ctx, valid, ids), n_used, tokens); + ggml_tensor * masked = ggml_mul(ctx, weights, v); + ggml_tensor * prod = ggml_mul(ctx, experts, ggml_reshape_3d(ctx, masked, 1, n_used, tokens)); + ggml_tensor * ref = ggml_reshape_2d(ctx, + ggml_repeat_back(ctx, prod, ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, tokens)), n_embd, tokens); + ggml_tensor * fused = ggml_moe_combine_masked(ctx, experts, weights, valid, ids); + ggml_set_output(ref); + ggml_set_output(fused); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, ref); + ggml_build_forward_expand(gf, fused); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::uniform_real_distribution uni(0.01f, 1.0f); + std::uniform_int_distribution pick(0, n_expert - 1); + std::vector h_e((size_t) n_embd * n_used * tokens), h_w((size_t) n_used * tokens); + std::vector h_v((size_t) n_expert * tokens); + std::vector h_ids((size_t) n_used * tokens); + for (float & x : h_e) x = normal(rng); + for (float & x : h_w) x = uni(rng); + for (size_t i = 0; i < h_v.size(); ++i) h_v[i] = (i % 3 == 0) ? 0.0f : 1.0f; + for (int32_t & x : h_ids) x = pick(rng); + ggml_backend_tensor_set(experts, h_e.data(), 0, sizeof(float) * h_e.size()); + ggml_backend_tensor_set(weights, h_w.data(), 0, sizeof(float) * h_w.size()); + ggml_backend_tensor_set(valid, h_v.data(), 0, sizeof(float) * h_v.size()); + ggml_backend_tensor_set(ids, h_ids.data(), 0, sizeof(int32_t) * h_ids.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + const size_t n = (size_t) n_embd * tokens; + std::vector a(n), b(n); + if (ok) { + ggml_backend_tensor_get(ref, a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(fused, b.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n && ok; ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0 && !(a[i] == 0.0f && b[i] == 0.0f)) { + std::fprintf(stderr, "masked combine mismatch n_embd=%d n_used=%d tokens=%d at %zu: ref %.9g fused %.9g\n", + n_embd, n_used, tokens, i, a[i], b[i]); + ok = false; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +// swiglu_ds4 read by a routed (MUL_MAT_ID) Q8_0 down projection: the GLU +// launch writes the q8_1 form MMVQ would quantize; the reference reads +// through a view. +bool run_glu_q8(ggml_backend_t backend, int n_ff, int n_used, int tokens, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + const int n_out = 64, n_expert = 8; + ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, n_used, tokens); + ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, n_used, tokens); + ggml_tensor * wq = ggml_new_tensor_3d(ctx, GGML_TYPE_Q8_0, n_ff, n_out, n_expert); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_used, tokens); + for (ggml_tensor * t : {gate, up, wq, ids}) ggml_set_input(t); + ggml_tensor * ref_g = ggml_swiglu_ds4_split(ctx, gate, up, 7.0f); + ggml_tensor * ref = ggml_mul_mat_id(ctx, wq, + ggml_view_3d(ctx, ref_g, n_ff, n_used, tokens, ref_g->nb[1], ref_g->nb[2], 0), ids); + ggml_tensor * fus_g = ggml_swiglu_ds4_split(ctx, gate, up, 7.0f); + ggml_tensor * fus = ggml_mul_mat_id(ctx, wq, fus_g, ids); + for (ggml_tensor * t : {ref_g, ref, fus_g, fus}) ggml_set_output(t); + ggml_cgraph * gf = ggml_new_graph(ctx); + for (ggml_tensor * t : {ref_g, ref, fus_g, fus}) ggml_build_forward_expand(gf, t); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::uniform_int_distribution pick(0, n_expert - 1); + std::vector h_g((size_t) n_ff * n_used * tokens), h_u(h_g.size()); + for (float & v : h_g) v = normal(rng) * 4.0f; + for (float & v : h_u) v = normal(rng) * 4.0f; + std::vector h_wf((size_t) n_ff * n_out * n_expert); + for (float & v : h_wf) v = normal(rng) * 0.05f; + std::vector h_wq(ggml_nbytes(wq)); + ggml_quantize_chunk(GGML_TYPE_Q8_0, h_wf.data(), h_wq.data(), 0, (int64_t) n_out * n_expert, n_ff, nullptr); + std::vector h_ids((size_t) n_used * tokens); + for (int32_t & v : h_ids) v = pick(rng); + ggml_backend_tensor_set(gate, h_g.data(), 0, sizeof(float) * h_g.size()); + ggml_backend_tensor_set(up, h_u.data(), 0, sizeof(float) * h_u.size()); + ggml_backend_tensor_set(wq, h_wq.data(), 0, h_wq.size()); + ggml_backend_tensor_set(ids, h_ids.data(), 0, sizeof(int32_t) * h_ids.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + for (auto pr : {std::make_pair(ref_g, fus_g), std::make_pair(ref, fus)}) { + if (!ok) break; + const size_t n = (size_t) ggml_nelements(pr.first); + std::vector a(n), b(n); + ggml_backend_tensor_get(pr.first, a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(pr.second, b.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n; ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + std::fprintf(stderr, "glu q8 mismatch n_ff=%d n_used=%d tokens=%d at %zu: ref %.9g fused %.9g\n", + n_ff, n_used, tokens, i, a[i], b[i]); + ok = false; + break; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +// A prefill band's HC mix: RMS_NORM read only by a product against 24 F16 +// rows runs as ggml_cuda_ds4_rms_norm_mix, F32 throughout. Checked against a +// double-precision host reference: the unfused GEMM path rounds its input to +// 16 bits on some GPUs, so it is no reference for an F32 kernel. A wrong row +// or token is off by O(1); float accumulation over K stays far inside 1e-4. +bool run_rms_norm_mix(ggml_backend_t backend, int k, int tokens, uint32_t seed) { + const int m = 24; + const float eps = 1e-6f; + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 8 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, tokens); + ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, k, m); + for (ggml_tensor * t : {x, w}) ggml_set_input(t); + ggml_tensor * fus = ggml_mul_mat(ctx, w, ggml_rms_norm(ctx, x, eps)); + ggml_set_output(fus); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, fus); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_x((size_t) k * tokens); + std::vector h_w((size_t) k * m); + for (float & v : h_x) v = normal(rng); + for (ggml_fp16_t & v : h_w) v = ggml_fp32_to_fp16(normal(rng) * 0.02f); + ggml_backend_tensor_set(x, h_x.data(), 0, sizeof(float) * h_x.size()); + ggml_backend_tensor_set(w, h_w.data(), 0, sizeof(ggml_fp16_t) * h_w.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + const size_t n = (size_t) m * tokens; + std::vector got(n); + if (ok) { + ggml_backend_tensor_get(fus, got.data(), 0, sizeof(float) * n); + std::vector want(n); + double peak = 0.0; + for (int t = 0; t < tokens; ++t) { + const float * xr = h_x.data() + (size_t) t * k; + double ss = 0.0; + for (int i = 0; i < k; ++i) ss += (double) xr[i] * xr[i]; + const double scale = 1.0 / std::sqrt(ss / k + eps); + for (int r = 0; r < m; ++r) { + double dot = 0.0; + for (int i = 0; i < k; ++i) dot += (double) ggml_fp16_to_fp32(h_w[(size_t) r * k + i]) * xr[i]; + want[(size_t) t * m + r] = dot * scale; + peak = std::max(peak, std::fabs(dot * scale)); + } + } + for (size_t i = 0; i < n; ++i) { + if (!(std::fabs(got[i] - want[i]) <= 1e-4 * peak)) { + std::fprintf(stderr, "rms norm mix mismatch k=%d tokens=%d at %zu: want %.9g got %.9g (peak %.6g)\n", + k, tokens, i, want[i], got[i], peak); + ok = false; + break; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +// The cluster shared expert's gate/up pair: Q8_0 row slices of wider weights +// at several columns (batch-invariant MMVQ), fused with SwiGLU-DS4 into one +// RDNA4 launch (ggml_cuda_mmvq_rdna4_glu_pair) vs the same products kept +// apart (a copy between each product and the GLU breaks the pattern). One +// column is a single-token decode step, which the generic fusion forms. +bool run_glu_pair(ggml_backend_t backend, int n_in, int n_full, int n_ff, int ff_begin, int ncols, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 48 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + ggml_tensor * gw = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, n_in, n_full); + ggml_tensor * uw = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, n_in, n_full); + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_in, ncols); + for (ggml_tensor * t : {gw, uw, x}) ggml_set_input(t); + auto slice = [&](ggml_tensor * w) { + return ggml_view_2d(ctx, w, n_in, n_ff, w->nb[1], (size_t) ff_begin * w->nb[1]); + }; + ggml_tensor * fus = ggml_swiglu_ds4_split(ctx, ggml_mul_mat(ctx, slice(gw), x), ggml_mul_mat(ctx, slice(uw), x), 7.0f); + ggml_tensor * ref = ggml_swiglu_ds4_split(ctx, ggml_cont(ctx, ggml_mul_mat(ctx, slice(gw), x)), + ggml_cont(ctx, ggml_mul_mat(ctx, slice(uw), x)), 7.0f); + for (ggml_tensor * t : {fus, ref}) ggml_set_output(t); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, fus); + ggml_build_forward_expand(gf, ref); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_w((size_t) n_in * n_full), h_x((size_t) n_in * ncols); + std::vector h_q(ggml_nbytes(gw)); + for (ggml_tensor * w : {gw, uw}) { + for (float & v : h_w) v = normal(rng) * 0.05f; + ggml_quantize_chunk(GGML_TYPE_Q8_0, h_w.data(), h_q.data(), 0, n_full, n_in, nullptr); + ggml_backend_tensor_set(w, h_q.data(), 0, h_q.size()); + } + for (float & v : h_x) v = normal(rng); + ggml_backend_tensor_set(x, h_x.data(), 0, sizeof(float) * h_x.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + const size_t n = (size_t) n_ff * ncols; + std::vector a(n), b(n); + if (ok) { + ggml_backend_tensor_get(ref, a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(fus, b.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n; ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + std::fprintf(stderr, "glu pair mismatch n_in=%d n_ff=%d begin=%d cols=%d at %zu: ref %.9g fused %.9g\n", + n_in, n_ff, ff_begin, ncols, i, a[i], b[i]); + ok = false; + break; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +// A rank's slice of a Q8_0 contraction: a column-range view flagged with +// ggml_backend_cuda_mul_mat_whole_row_lanes vs the whole rows times the slice +// zero-padded to full width (build_shared_ffn_slice's padded form). The lanes +// are a HIP layout, so there the two are bit-identical; other backends compute +// the range's product as usual, and the check is the value to rounding. +#if defined(GGML_USE_HIP) +constexpr bool k_col_slice_bit_exact = true; +#else +constexpr bool k_col_slice_bit_exact = false; +#endif + +bool run_col_slice(ggml_backend_t backend, int n_out, int n_full, int k_begin, int k_count, int ncols, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 48 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, n_full, n_out); + ggml_tensor * mid = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k_count, ncols); + for (ggml_tensor * t : {w, mid}) ggml_set_input(t); + ggml_tensor * view = ggml_view_2d(ctx, w, k_count, n_out, w->nb[1], ggml_row_size(GGML_TYPE_Q8_0, k_begin)); + ggml_tensor * fus = ggml_mul_mat(ctx, view, mid); + ggml_backend_cuda_mul_mat_whole_row_lanes(fus); + ggml_tensor * zeros = ggml_scale(ctx, mid, 0.0f); + ggml_tensor * padded = nullptr; + for (int off = 0; off < n_full; off += k_count) { + ggml_tensor * piece = off == k_begin ? mid : zeros; + padded = padded ? ggml_concat(ctx, padded, piece, 0) : piece; + } + ggml_tensor * ref = ggml_mul_mat(ctx, w, padded); + for (ggml_tensor * t : {fus, ref}) ggml_set_output(t); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, fus); + ggml_build_forward_expand(gf, ref); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_w((size_t) n_full * n_out), h_m((size_t) k_count * ncols); + for (float & v : h_w) v = normal(rng) * 0.05f; + for (float & v : h_m) v = normal(rng); + std::vector h_q(ggml_nbytes(w)); + ggml_quantize_chunk(GGML_TYPE_Q8_0, h_w.data(), h_q.data(), 0, n_out, n_full, nullptr); + ggml_backend_tensor_set(w, h_q.data(), 0, h_q.size()); + ggml_backend_tensor_set(mid, h_m.data(), 0, sizeof(float) * h_m.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + const size_t n = (size_t) n_out * ncols; + std::vector a(n), b(n); + if (ok) { + ggml_backend_tensor_get(ref, a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(fus, b.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n; ++i) { + const bool same = k_col_slice_bit_exact + ? std::memcmp(&a[i], &b[i], sizeof(float)) == 0 + : std::fabs(a[i] - b[i]) <= 1e-5f * fmaxf(1.0f, std::fabs(a[i])); + if (!same) { + std::fprintf(stderr, "col slice mismatch n_out=%d n_full=%d begin=%d count=%d cols=%d at %zu: ref %.9g fused %.9g\n", + n_out, n_full, k_begin, k_count, ncols, i, a[i], b[i]); + ok = false; + break; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +#if defined(GGML_USE_HIP) +// rms_norm * w -> tail rope -> F16 copy (the K/V latent of the flash lanes) +// in one launch vs the three; the reference keeps its rms_norm as an output. +bool run_norm_rope(ggml_backend_t backend, int ncols, int tokens, int pos0, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ncols, tokens); + ggml_tensor * w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ncols); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens); + for (ggml_tensor * t : {x, w, pos}) ggml_set_input(t); + ggml_tensor * outs[2][2]; + for (int k = 0; k < 2; ++k) { + ggml_tensor * nrm = ggml_rms_norm(ctx, x, 1e-6f); + if (k == 0) ggml_set_output(nrm); + ggml_tensor * m = ggml_mul(ctx, nrm, w); + ggml_tensor * r = ggml_rope_ext(ctx, ggml_reshape_3d(ctx, m, ncols, 1, tokens), pos, nullptr, + 64, GGML_ROPE_TYPE_NORMAL | GGML_ROPE_TYPE_TAIL, 65536, + 10000.0f, 0.25f, 1.0f, 1.0f, 32.0f, 1.0f); + ggml_tensor * r2 = ggml_reshape_2d(ctx, r, ncols, tokens); + outs[k][0] = r2; + outs[k][1] = ggml_cast(ctx, r2, GGML_TYPE_F16); + ggml_set_output(outs[k][0]); + ggml_set_output(outs[k][1]); + } + ggml_cgraph * gf = ggml_new_graph(ctx); + for (int k = 0; k < 2; ++k) { + ggml_build_forward_expand(gf, outs[k][1]); + ggml_build_forward_expand(gf, outs[k][0]); + } + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_x((size_t) ncols * tokens), h_w(ncols); + for (float & v : h_x) v = normal(rng) * 2.0f; + for (float & v : h_w) v = 1.0f + 0.3f * normal(rng); + std::vector h_pos(tokens); + for (int t = 0; t < tokens; ++t) h_pos[t] = pos0 + t; + ggml_backend_tensor_set(x, h_x.data(), 0, sizeof(float) * h_x.size()); + ggml_backend_tensor_set(w, h_w.data(), 0, sizeof(float) * h_w.size()); + ggml_backend_tensor_set(pos, h_pos.data(), 0, sizeof(int32_t) * h_pos.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + const size_t n = (size_t) ncols * tokens; + if (ok) { + std::vector a(n), b(n); + std::vector ah(n), bh(n); + ggml_backend_tensor_get(outs[0][0], a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(outs[1][0], b.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(outs[0][1], ah.data(), 0, sizeof(ggml_fp16_t) * n); + ggml_backend_tensor_get(outs[1][1], bh.data(), 0, sizeof(ggml_fp16_t) * n); + for (size_t i = 0; i < n && ok; ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0 || ah[i] != bh[i]) { + std::fprintf(stderr, "norm rope mismatch ncols=%d tokens=%d at %zu: ref %.9g fused %.9g\n", + ncols, tokens, i, a[i], b[i]); + ok = false; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} +#endif + +// cont(permute(x)) read only by a Q8_0 matmul: quantized straight from the +// view, vs the same graph with the cont read through a view (kept). +bool run_cont_mmvq(ggml_backend_t backend, int ne0, int ne1, int tokens, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + const int n_out = 96; + ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, ne0, tokens, ne1); // [ne0, T, groups] + ggml_tensor * wq = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, ne0 * ne1, n_out); + for (ggml_tensor * t : {x, wq}) ggml_set_input(t); + ggml_tensor * res[2]; + for (int k = 0; k < 2; ++k) { + ggml_tensor * c = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [ne0, groups, T] + ggml_tensor * r = ggml_reshape_2d(ctx, c, (int64_t) ne0 * ne1, tokens); + if (k == 0) r = ggml_view_2d(ctx, r, r->ne[0], r->ne[1], r->nb[1], 0); + res[k] = ggml_mul_mat(ctx, wq, r); + ggml_set_output(res[k]); + } + ggml_cgraph * gf = ggml_new_graph(ctx); + for (int k = 0; k < 2; ++k) ggml_build_forward_expand(gf, res[k]); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_x((size_t) ne0 * ne1 * tokens), h_wf((size_t) ne0 * ne1 * n_out); + for (float & v : h_x) v = normal(rng); + for (float & v : h_wf) v = normal(rng) * 0.05f; + std::vector h_wq(ggml_nbytes(wq)); + ggml_quantize_chunk(GGML_TYPE_Q8_0, h_wf.data(), h_wq.data(), 0, n_out, (int64_t) ne0 * ne1, nullptr); + ggml_backend_tensor_set(x, h_x.data(), 0, sizeof(float) * h_x.size()); + ggml_backend_tensor_set(wq, h_wq.data(), 0, h_wq.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + const size_t n = (size_t) n_out * tokens; + if (ok) { + std::vector a(n), b(n); + ggml_backend_tensor_get(res[0], a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(res[1], b.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n && ok; ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + std::fprintf(stderr, "cont mmvq mismatch ne0=%d ne1=%d tokens=%d at %zu: ref %.9g fused %.9g\n", + ne0, ne1, tokens, i, a[i], b[i]); + ok = false; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +// Masked single-token attention softmax with a sink column: scale, add the +// mask row, concat the sinks, soft_max, view vs one sink-column launch. +bool run_sink_mask(ggml_backend_t backend, int n_attn, int n_head, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + const float kq_scale = 0.0441941738f; + ggml_tensor * scores = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_attn, n_head); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_attn, 1); + ggml_tensor * sinks = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_head); + for (ggml_tensor * t : {scores, mask, sinks}) ggml_set_input(t); + ggml_tensor * s = ggml_add(ctx, ggml_scale(ctx, scores, kq_scale), mask); + ggml_tensor * sws = ggml_concat(ctx, s, ggml_reshape_2d(ctx, sinks, 1, n_head), 0); + ggml_tensor * pws = ggml_soft_max(ctx, sws); + ggml_tensor * ref = ggml_cont(ctx, ggml_view_2d(ctx, pws, n_attn, n_head, pws->nb[1], 0)); + ggml_tensor * fused = ggml_reshape_2d(ctx, ggml_soft_max_ext_sink_col_mask( + ctx, ggml_reshape_3d(ctx, scores, n_attn, 1, n_head), mask, sinks, kq_scale), n_attn, n_head); + ggml_set_output(ref); + ggml_set_output(fused); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, ref); + ggml_build_forward_expand(gf, fused); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_s((size_t) n_attn * n_head), h_m(n_attn), h_k(n_head); + for (float & v : h_s) v = normal(rng) * 20.0f; + for (int i = 0; i < n_attn; ++i) h_m[i] = (i % 7 == 3) ? -INFINITY : 0.0f; + for (float & v : h_k) v = normal(rng); + ggml_backend_tensor_set(scores, h_s.data(), 0, sizeof(float) * h_s.size()); + ggml_backend_tensor_set(mask, h_m.data(), 0, sizeof(float) * h_m.size()); + ggml_backend_tensor_set(sinks, h_k.data(), 0, sizeof(float) * h_k.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + const size_t n = (size_t) n_attn * n_head; + if (ok) { + std::vector a(n), b(n); + ggml_backend_tensor_get(ref, a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(fused, b.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n && ok; ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + std::fprintf(stderr, "sink mask mismatch n_attn=%d n_head=%d at %zu: ref %.9g fused %.9g\n", + n_attn, n_head, i, a[i], b[i]); + ok = false; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +// Inverse tail rope of the attention context read by the grouped output +// projection (batched Q8_0 MMVQ over a [D * hpg, T, groups] view): the rope +// writes the view's q8_1; the reference reads a copy of the rope output. +bool run_rope_q8(ggml_backend_t backend, int n_head, int hpg, int tokens, uint32_t seed) { + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * 32 + ggml_graph_overhead(); + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + const int D = 512, n_out = 64; + const int G = D * hpg, NG = n_head / hpg; + ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, D, n_head, tokens); + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens); + ggml_tensor * wq = ggml_new_tensor_3d(ctx, GGML_TYPE_Q8_0, G, n_out, NG); + for (ggml_tensor * t : {x, pos, wq}) ggml_set_input(t); + ggml_tensor * res[2]; + for (int k = 0; k < 2; ++k) { + ggml_tensor * r = ggml_rope_ext(ctx, x, pos, nullptr, 64, GGML_ROPE_TYPE_NORMAL | GGML_ROPE_TYPE_TAIL, + 65536, 10000.0f, 0.25f, 1.0f, 1.0f, 32.0f, 1.0f); + if (k == 0) r = ggml_cont(ctx, r); + ggml_tensor * v = ggml_permute(ctx, ggml_reshape_3d(ctx, r, G, NG, tokens), 0, 2, 1, 3); + res[k] = ggml_mul_mat(ctx, wq, v); + ggml_set_output(res[k]); + } + ggml_cgraph * gf = ggml_new_graph(ctx); + for (int k = 0; k < 2; ++k) ggml_build_forward_expand(gf, res[k]); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + std::mt19937 rng(seed); + std::normal_distribution normal(0.0f, 1.0f); + std::vector h_x((size_t) D * n_head * tokens), h_wf((size_t) G * n_out * NG); + for (float & v : h_x) v = normal(rng); + for (float & v : h_wf) v = normal(rng) * 0.05f; + std::vector h_wq(ggml_nbytes(wq)); + ggml_quantize_chunk(GGML_TYPE_Q8_0, h_wf.data(), h_wq.data(), 0, (int64_t) n_out * NG, G, nullptr); + std::vector h_pos(tokens); + for (int t = 0; t < tokens; ++t) h_pos[t] = -(1234 + t); + ggml_backend_tensor_set(x, h_x.data(), 0, sizeof(float) * h_x.size()); + ggml_backend_tensor_set(pos, h_pos.data(), 0, sizeof(int32_t) * h_pos.size()); + ggml_backend_tensor_set(wq, h_wq.data(), 0, h_wq.size()); + bool ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + const size_t n = (size_t) ggml_nelements(res[0]); + if (ok) { + std::vector a(n), b(n); + ggml_backend_tensor_get(res[0], a.data(), 0, sizeof(float) * n); + ggml_backend_tensor_get(res[1], b.data(), 0, sizeof(float) * n); + for (size_t i = 0; i < n && ok; ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + std::fprintf(stderr, "rope q8 mismatch n_head=%d hpg=%d tokens=%d at %zu: ref %.9g fused %.9g\n", + n_head, hpg, tokens, i, a[i], b[i]); + ok = false; + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + } // namespace int main() { @@ -210,6 +948,65 @@ int main() { if (!run_collapse(backend, c, seed++)) ++failures; } } + const BoundaryCase boundary_cases[] = { + {5120, 1, false}, {5120, 2, false}, {5120, 3, true}, {5120, 5, false}, + {5120, 5, true}, {5120, 8, false}, {4096, 4, false}, {1024, 5, true}, + }; + for (const BoundaryCase & c : boundary_cases) { + for (int rep = 0; rep < 2; ++rep) { + if (!run_boundary(backend, c, seed++)) ++failures; + } + } + const int norm_cases[][2] = {{1280, 1}, {1280, 3}, {1280, 5}, {512, 4}, {5120, 2}, {800, 5}}; + for (const auto & nc : norm_cases) { + if (!run_norm_q8(backend, nc[0], nc[1], seed++)) ++failures; + } + const int combine_cases[][4] = {{5120, 6, 384, 1}, {5120, 6, 384, 5}, {512, 8, 64, 3}, {2048, 2, 16, 7}}; + for (const auto & cc : combine_cases) { + if (!run_masked_combine(backend, cc[0], cc[1], cc[2], cc[3], seed++)) ++failures; + } + const int glu_cases[][3] = {{2304, 6, 5}, {2304, 6, 1}, {2304, 3, 3}, {512, 2, 4}}; + for (const auto & gc : glu_cases) { + if (!run_glu_q8(backend, gc[0], gc[1], gc[2], seed++)) ++failures; + } +#if defined(GGML_USE_HIP) + const int rope_cases[][3] = {{512, 5, 1000}, {512, 1, 77777}, {512, 3, 5}, {1024, 4, 123456}}; + for (const auto & rc : rope_cases) { + if (!run_norm_rope(backend, rc[0], rc[1], rc[2], seed++)) ++failures; + } +#else + std::fprintf(stderr, "norm rope ... skipped (HIP-only candidate)\n"); +#endif + const int cont_cases[][3] = {{1024, 4, 5}, {1024, 4, 1}, {512, 8, 3}, {1024, 2, 8}}; + for (const auto & cc : cont_cases) { + if (!run_cont_mmvq(backend, cc[0], cc[1], cc[2], seed++)) ++failures; + } + const int sink_cases[][2] = {{128, 32}, {129, 32}, {128, 64}, {640, 16}}; + for (const auto & sc : sink_cases) { + if (!run_sink_mask(backend, sc[0], sc[1], seed++)) ++failures; + } + const int rope_q8_cases[][3] = {{32, 8, 5}, {32, 8, 2}, {64, 8, 3}, {16, 4, 5}}; + for (const auto & rc : rope_q8_cases) { + if (!run_rope_q8(backend, rc[0], rc[1], rc[2], seed++)) ++failures; + } + { + // Verify-shaped products: batch-invariant MMVQ, as the DS4 verify runs them. + const bool prev = ggml_backend_cuda_set_mmvq_batch_invariant(true); + const int norm_mix_cases[][2] = {{20480, 64}, {20480, 130}, {1024, 67}}; + for (const auto & nc : norm_mix_cases) { + if (!run_rms_norm_mix(backend, nc[0], nc[1], seed++)) ++failures; + } + const int pair_cases[][5] = {{1024, 64, 32, 0, 5}, {1024, 64, 32, 32, 4}, {5120, 96, 48, 48, 2}, {1024, 64, 32, 32, 8}, + {1024, 64, 32, 32, 1}}; + for (const auto & pc : pair_cases) { + if (!run_glu_pair(backend, pc[0], pc[1], pc[2], pc[3], pc[4], seed++)) ++failures; + } + const int slice_cases[][5] = {{64, 2304, 0, 1152, 5}, {64, 2304, 1152, 1152, 4}, {96, 4096, 2048, 2048, 3}, {64, 2304, 1152, 1152, 1}}; + for (const auto & sc : slice_cases) { + if (!run_col_slice(backend, sc[0], sc[1], sc[2], sc[3], sc[4], seed++)) ++failures; + } + ggml_backend_cuda_set_mmvq_batch_invariant(prev); + } ggml_backend_free(backend); std::printf("ds4 fused ops: %d failures\n", failures); return failures == 0 ? 0 : 1; diff --git a/server/test/unit/test_launch_policy.cpp b/server/test/unit/test_launch_policy.cpp index ed07c7f33..9b2b1e9d3 100644 --- a/server/test/unit/test_launch_policy.cpp +++ b/server/test/unit/test_launch_policy.cpp @@ -309,6 +309,19 @@ struct LaunchPolicyFixture : CommonFixture { CHECK(find_launch_profile("missing") == nullptr); } + void test_cluster_env_is_well_formed() { + // Cluster-mode defaults are switches an operator can turn off with 0, + // so every one has a name, a non-empty value and appears once. + std::set names; + for (const LaunchProfileEnv & entry : cluster_launch_env()) { + CHECK(entry.name && *entry.name); + CHECK(entry.value && *entry.value); + CHECK(names.insert(entry.name).second); + } + CHECK(names.count("LUCE_CLUSTER_FAST_REDUCE") == 1); + CHECK(names.count("LUCE_CLUSTER_ATTENTION_PARALLEL") == 1); + } + void test_profile_flags_yield_to_explicit_flags() { const LaunchProfile * profile = find_launch_profile("ds4-r9700-strix"); CHECK(profile != nullptr); @@ -443,6 +456,7 @@ TEST_CASE(LaunchPolicyFixture, launch_policy_suite) { test_auto_placement_preserves_separate_vision_encoder(); test_draft_placement_precedence(); test_profiles_are_well_formed(); + test_cluster_env_is_well_formed(); test_profile_flags_yield_to_explicit_flags(); test_profile_replaces_documented_recipe(); test_ds41_profile_is_the_lucebox_recipe(); diff --git a/server/test/unit/test_rocmfp_mix_gateup_glu.cpp b/server/test/unit/test_rocmfp_mix_gateup_glu.cpp index 370e31f36..89758fac9 100644 --- a/server/test/unit/test_rocmfp_mix_gateup_glu.cpp +++ b/server/test/unit/test_rocmfp_mix_gateup_glu.cpp @@ -419,7 +419,7 @@ void RocmfpMixGateupGluFixture::check_fused_gateup_glu(bool fp3) { "half-registered/mismatched pairs are refused, and out-of-range ids zero\n"); } -// Prefill-sized batches of the MIX types take the F16 WMMA GEMM on RDNA3.5 +// Prefill-sized batches of the MIX types take the F16 WMMA GEMM on RDNA3.5 and RDNA4 // (mix-wmma-moe.cu; LUCE_MMID_TELEMETRY=1 logs path=mix_wmma) and MMQ // elsewhere. Either way the graph result must match the registry's exact // per-route dots, run in verify-sized chunks, within F16 rounding, and masked