From 3d44f159d6d051d2e305164b4f5e1a5a3b18c2bf Mon Sep 17 00:00:00 2001 From: mrciffa Date: Thu, 8 Oct 2026 13:08:39 +0200 Subject: [PATCH 01/40] perf(ggml-cuda): fused, bit-exact DeepSeek V4.1 decode kernels The verify step of DS4.1 runs many short kernels per layer; these cut the launch count without changing a bit of the output: - one launch for an HC sub-block boundary (norm, mix, Sinkhorn split, post, collapse and the next norm), and for the router's select and weights over the same logits; - RMS norm, rope and the F16 copy in one launch (HIP only: nvcc's fast math contracts rope_norm's rotation in a way another kernel cannot be guaranteed to repeat), and a rope that also writes the q8_1 input of the grouped output projection; - the masked sink softmax in one launch, and one split-KV flash launch for all single-token verify lanes of a layer; - a skinny-K single-column F16 matvec, and wide Q4_K products over many rows routed to MMQ on RDNA4; - one expert grouping per MoE layer for gate, up and down, and a q8_1 activation memo shared across reshapes; - column-range Q8_0 products that keep the whole row's lane mapping, so a slice of a weight matrix gives the same sums as the full product; - a split's batched peer copies as one copy kernel, an in-place cluster all-reduce node, and an argmax pair/pick that finds the greedy token of a vocabulary split across ranks exactly. Each fused kernel repeats the replaced kernels' arithmetic in the same order; test_ds4_fused_ops_cuda checks them bit for bit. Co-Authored-By: Claude Opus 5.5 --- .../llama.cpp/ggml/include/ggml-backend.h | 7 + .../deps/llama.cpp/ggml/include/ggml-cuda.h | 12 + server/deps/llama.cpp/ggml/include/ggml.h | 66 ++ .../deps/llama.cpp/ggml/src/ggml-backend.cpp | 16 +- .../llama.cpp/ggml/src/ggml-cuda/common.cuh | 12 + .../deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu | 9 + .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu | 715 +++++++++++++++++- .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh | 19 + .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 192 ++++- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 569 +++++++++++++- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu | 155 +++- .../llama.cpp/ggml/src/ggml-cuda/mmvf.cuh | 3 + .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 195 +++-- .../llama.cpp/ggml/src/ggml-cuda/mmvq.cuh | 5 + .../llama.cpp/ggml/src/ggml-cuda/moe-fused.cu | 66 ++ .../deps/llama.cpp/ggml/src/ggml-cuda/norm.cu | 68 +- .../llama.cpp/ggml/src/ggml-cuda/norm.cuh | 5 +- .../deps/llama.cpp/ggml/src/ggml-cuda/rope.cu | 202 +++++ .../llama.cpp/ggml/src/ggml-cuda/rope.cuh | 10 + .../llama.cpp/ggml/src/ggml-cuda/softmax.cu | 15 +- .../llama.cpp/ggml/src/ggml-cuda/unary.cu | 48 +- .../llama.cpp/ggml/src/ggml-cuda/unary.cuh | 3 +- server/deps/llama.cpp/ggml/src/ggml.c | 131 +++- server/test/unit/test_ds4_fused_ops_cuda.cpp | 585 +++++++++++++- 24 files changed, 2945 insertions(+), 163 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml-backend.h b/server/deps/llama.cpp/ggml/include/ggml-backend.h index e6bf95e5e8..1d4315b00d 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-backend.h +++ b/server/deps/llama.cpp/ggml/include/ggml-backend.h @@ -384,6 +384,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 0f5c1fae64..143591a9ee 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); @@ -181,6 +188,11 @@ GGML_BACKEND_API void ggml_backend_cuda_copy_batch_async(ggml_backend_t backend, // Graph copy runs issued as one batched launch on the calling thread (see // GGML_CUDA_DISABLE_COPY_BATCH); for tests and profiling. GGML_BACKEND_API size_t ggml_backend_cuda_get_copy_batch_run_count(void); +// Keeps the backend's GPU out of its idle power state: on starts one wave that +// sleeps in a loop on a stream of its own (at most max_ms), off releases it. +// For a process that leaves the GPU idle for milliseconds between launches and +// pays the wake-up (about 1 ms on RDNA4) at the next one. HIP only. +GGML_BACKEND_API void ggml_backend_cuda_keepalive(ggml_backend_t backend, bool on, int max_ms); GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index a0874b6e77..191df31cee 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); @@ -2567,6 +2578,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( @@ -2684,6 +2710,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 @@ -2745,11 +2772,33 @@ 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). + 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. @@ -2916,6 +2965,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 ae3baf9305..b381a75439 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp @@ -882,6 +882,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]; @@ -1521,7 +1522,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) { + if (sched->n_copies > 1 || sched->dedicated_copies) { ggml_set_input(tensor_copy); ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor } @@ -2369,7 +2370,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]; @@ -2954,6 +2959,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-cuda/common.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh index 232bc01594..2693e050d9 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 @@ -1491,6 +1494,15 @@ struct ggml_backend_cuda_context { }; std::vector luce_q8_memo; + // DS4 HC boundary launches (ds4-hc.cu) synchronize their blocks on this + // device counter; launches on this context's stream run in order, so the + // running count of blocks launched is each launch's target. + unsigned long long * ds4_hc_barrier = nullptr; + unsigned long long ds4_hc_barrier_blocks = 0; + // 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 3d7a698b9b..e3586763b6 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 0ee122822e..2ee53c53b7 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 @@ -490,11 +490,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 +557,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 +635,628 @@ 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; +} + +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 long long * barrier, const unsigned long long barrier_target, + 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); + + __syncthreads(); + if (tid == 0) { + __threadfence(); + atomicAdd(barrier, 1ull); + while (atomicAdd(barrier, 0ull) < barrier_target) { +#if defined(GGML_USE_HIP) + __builtin_amdgcn_s_sleep(1); +#endif + } + __threadfence(); + } + __syncthreads(); +#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 long long * barrier, const unsigned long long barrier_target, + 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); + + __syncthreads(); + if (tid == 0) { + __threadfence(); + atomicAdd(barrier, 1ull); + while (atomicAdd(barrier, 0ull) < barrier_target) { +#if defined(GGML_USE_HIP) + __builtin_amdgcn_s_sleep(1); +#endif + } + __threadfence(); + } + __syncthreads(); +#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); + } + } + } + } +} + +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); + // Launches on this context's stream run in order, so a running count of + // the blocks they launched is each launch's barrier target. + if (!ctx.ds4_hc_barrier) { + ggml_cuda_set_device(ctx.device); + CUDA_CHECK(cudaMalloc((void **) &ctx.ds4_hc_barrier, sizeof(unsigned long long))); + CUDA_CHECK(cudaMemsetAsync(ctx.ds4_hc_barrier, 0, sizeof(unsigned long long), ctx.stream())); + ctx.ds4_hc_barrier_blocks = 0; + } + ctx.ds4_hc_barrier_blocks += (unsigned long long) n_tokens; + 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. + static const bool split_mix = [] { + const char * v = getenv("LUCE_DS4_HC_BOUNDARY_SPLIT"); + return v && *v && strcmp(v, "0") != 0; + }(); + if (split_mix && n_tokens * 6 <= 48) { + 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, ctx.ds4_hc_barrier_blocks, \ + (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, ctx.ds4_hc_barrier_blocks, \ + (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]; 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 fc03b57cc1..51d4e84e4c 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,22 @@ // 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) + +// 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 a412b72f00..cf924cfb6d 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 7bb47b709a..a65c8af374 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 @@ -154,6 +154,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; } @@ -951,6 +956,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) { @@ -3035,6 +3050,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); @@ -3342,6 +3362,19 @@ 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); + // Q4_K at two or more columns over many rows (the DSpark drafter runs + // its dense projections at width five): on gfx1201 MMQ streams these + // at about twice the MMVQ rate (bench_dspark_mmv, width 5: attn_q_b + // 67 vs 137 us, attn_output_b 79 vs 117 us, fc 113 vs 217 us). MMVQ + // keeps the narrow shapes (under 2048 rows), where it is faster. + static const bool q4k_ncol_mmq = [] { + const char * e = std::getenv("LUCE_CUDA_Q4K_NCOL_MMQ"); + return !(e && e[0] == '0'); + }(); + if (q4k_ncol_mmq && use_mul_mat_vec_q && use_mul_mat_q && src0->type == GGML_TYPE_Q4_K && + GGML_CUDA_CC_IS_RDNA4(cc) && src1->ne[1] > 1 && src0->ne[1] >= 2048) { + use_mul_mat_vec_q = false; + } } // debug helpers @@ -3476,7 +3509,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"); @@ -3500,6 +3533,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 @@ -4283,8 +4324,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; @@ -4309,6 +4363,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)); @@ -4479,10 +4539,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); @@ -4516,6 +4590,64 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ return true; } +#if defined(GGML_USE_HIP) +static __global__ void ggml_cuda_keepalive_kernel(const volatile unsigned int * stop, unsigned int epoch, + unsigned long long max_ticks) { + const unsigned long long t0 = wall_clock64(); + while (*stop != epoch) { + __builtin_amdgcn_s_sleep(100); + if (wall_clock64() - t0 > max_ticks) break; + } +} +#endif + +void ggml_backend_cuda_keepalive(ggml_backend_t backend, bool on, int max_ms) { +#if defined(GGML_USE_HIP) + if (!ggml_backend_is_cuda(backend)) { + return; + } + struct state { + cudaStream_t stream = nullptr; + unsigned int * flag = nullptr; // pinned host word the kernel polls + unsigned int epoch = 0; + bool running = false; + unsigned long long ticks_per_ms = 0; + }; + static state states[GGML_CUDA_MAX_DEVICES]; + ggml_backend_cuda_context * ctx = (ggml_backend_cuda_context *) backend->context; + state & s = states[ctx->device]; + if (on) { + if (s.running) { + return; + } + ggml_cuda_set_device(ctx->device); + if (!s.stream) { + CUDA_CHECK(cudaStreamCreateWithFlags(&s.stream, cudaStreamNonBlocking)); + CUDA_CHECK(cudaMallocHost((void **) &s.flag, sizeof(unsigned int))); + *s.flag = 0; + int khz = 0; + if (hipDeviceGetAttribute(&khz, hipDeviceAttributeWallClockRate, ctx->device) != hipSuccess || khz <= 0) { + khz = 100000; + } + s.ticks_per_ms = (unsigned long long) khz; + } + s.epoch++; + ggml_cuda_keepalive_kernel<<<1, 32, 0, s.stream>>>(s.flag, s.epoch, + s.ticks_per_ms * (unsigned long long) std::max(1, max_ms)); + CUDA_CHECK(cudaGetLastError()); + s.running = true; + } else { + if (!s.running) { + return; + } + __atomic_store_n(s.flag, s.epoch, __ATOMIC_RELEASE); + s.running = false; + } +#else + GGML_UNUSED(backend); GGML_UNUSED(on); GGML_UNUSED(max_ms); +#endif +} + static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; @@ -5325,6 +5457,361 @@ 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. LUCE_DS4_FUSE_HC_BOUNDARY=0 +// keeps the six launches. +static bool ggml_cuda_ds4_hc_boundary_enabled() { + static const bool on = [] { + const char * v = getenv("LUCE_DS4_FUSE_HC_BOUNDARY"); + return !(v && v[0] == '0' && v[1] == '\0'); + }(); + return on; +} + +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]; +} + +// 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 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; + // The plain MMVQ branch of ggml_cuda_mul_mat, nothing else. + 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; + if (!ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX || + src0->type == GGML_TYPE_Q3_1_ROCMFP3_MIX || src0->type == GGML_TYPE_Q2_0_ROCMFP2 || + bad_padding_clear || ggml_backend_buft_is_cuda_split(src0->buffer->buft) || + ggml_mul_mat_is_grouped_src(mm) || src1->ne[1] > ggml_cuda_mmvq_max_ncols()) 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; + // The plain MMVQ branch of ggml_cuda_mul_mat, nothing else. + 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; + if (!ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX || + src0->type == GGML_TYPE_Q3_1_ROCMFP3_MIX || src0->type == GGML_TYPE_Q2_0_ROCMFP2 || + bad_padding_clear || ggml_backend_buft_is_cuda_split(src0->buffer->buft) || + ggml_mul_mat_is_grouped_src(mm) || src1->ne[1] > ggml_cuda_mmvq_max_ncols()) 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) { + if (!ggml_cuda_ds4_hc_boundary_enabled()) return -1; + const int n = cgraph->n_nodes; + // The next node at or after j that launches something; the skipped views + // are recorded so their single reader can be checked. + 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) { @@ -5539,6 +6026,50 @@ 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 > i) { + i = last; + continue; + } + } + const bool qwen4exp_rdna35 = ggml_cuda_qwen4exp_enabled() && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc); if (qwen4exp_rdna35) { @@ -5968,7 +6499,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; } @@ -6384,6 +6918,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()) { @@ -6425,7 +6970,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 @@ -7278,6 +7823,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; } @@ -7389,6 +7940,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/mmvf.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu index f6c2760972..8b8356f287 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,104 @@ 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). The +// generic kernel gives every row its own block, so a K=256 row is a single +// 512-byte trip and the launch is bound by block scheduling (~360 GB/s on +// gfx1201 for the DSpark Markov head, vocab 129280 x rank 256). Here every lane +// loads 16 bytes, K/8 lanes share a row and one wave covers several rows, which +// streams the matrix at ~540 GB/s. Sums accumulate in F32, so values can differ +// from the generic F16-accumulating path in the last bits. +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; + } + } +} + +static bool mmvf_skinny_enabled() { + static const bool enabled = [] { + const char * value = std::getenv("LUCE_CUDA_MMVF_SKINNY"); + return !(value && value[0] == '0'); + }(); + return enabled; +} + +// 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 (src0->type != GGML_TYPE_F16 || ids || !mmvf_skinny_enabled()) { + 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 device = ggml_cuda_get_device(); + 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; +} + 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 +795,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 a50f7c0218..9f68ce4d95 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 beb82d810e..fa94bf0081 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,18 @@ static bool mmid_grouped_type_ok(ggml_type type) { } } +// [TAG_MMID_GROUPED] LUCE_MMID_PREP_MEMO=0 recomputes the expert grouping for +// every MUL_MAT_ID (see the grouped path in ggml_cuda_mul_mat_vec_q). +static bool mmid_prep_memo_on() { + static const bool on = [] { + const char * e = std::getenv("LUCE_MMID_PREP_MEMO"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + return on; +} +// 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"); @@ -2597,7 +2609,8 @@ static __global__ void mul_mat_vec_q8_0_rdna4( 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) { 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"); @@ -2636,11 +2649,16 @@ static __global__ void mul_mat_vec_q8_0_rdna4( #pragma unroll for (int q = 0; q < nq; ++q) part[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]; @@ -2735,14 +2753,14 @@ 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); } template @@ -2752,12 +2770,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 @@ -2779,11 +2797,11 @@ static void mul_mat_vec_q8_0_rdna4_launch( 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_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); @@ -2831,7 +2849,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. @@ -2982,7 +3001,7 @@ 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, stream); return; } } @@ -3329,7 +3348,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 +3384,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 +3525,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. @@ -3576,13 +3604,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 +3625,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 +3661,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 +3721,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 +3732,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 && mmid_prep_memo_on() && 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 +3840,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 b8b3d721ec..af82a0b89d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cuh @@ -27,3 +27,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 60e5553b76..f09993b279 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,52 @@ 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); + 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 +1077,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 701be15077..cc7045179e 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 69ac62333a..3b5c3b698c 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 cfcd97613c..3bd2d006b6 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 304ef60aef..7dca1360c6 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 1ed928d289..69b0929714 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,14 @@ 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) { +#pragma clang fp contract(off) + return 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 +108,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 +466,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 ed7f3d02a6..ce1773163e 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 1864ce2dc4..39ca8dad9f 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 08a9dea24e..a7d3606a75 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/test/unit/test_ds4_fused_ops_cuda.cpp b/server/test/unit/test_ds4_fused_ops_cuda.cpp index 2863ca4b9e..4606e61599 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" @@ -173,6 +177,544 @@ 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; +} + +#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 +752,47 @@ 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; + } ggml_backend_free(backend); std::printf("ds4 fused ops: %d failures\n", failures); return failures == 0 ? 0 : 1; From 352519268d1dfaaa8e2bf90c51635a54d2865b76 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Thu, 8 Oct 2026 13:08:39 +0200 Subject: [PATCH 02/40] perf(moe): fused routed combine and q8_1 inputs for the hybrid expert owners - The owner route mask, the route weighting and the route sum run in one launch, and the clamped routed SwiGLU in one launch; the masked combine rounds its products and its sums separately, as the unfused ops do. - The shared-expert slice is queued ahead of the owner join. - Opt-in (LUCE_MOE_Q8_HANDOFF=1): the secondary expert owner receives its FFN input as the q8_1 rows its MMVQ reads, quantized once on the primary device. That moves about 3.6x fewer bytes over PCIe and the secondary skips its own quantize; the products are unchanged. Co-Authored-By: Claude Opus 5.5 --- server/src/common/moe_hybrid_ffn_eval.cpp | 128 ++++++++++++++++++--- server/src/common/moe_hybrid_ffn_eval.h | 6 + server/src/common/moe_hybrid_placement.cpp | 7 +- server/src/common/moe_hybrid_placement.h | 3 + server/src/common/moe_hybrid_storage.cpp | 24 +++- server/src/common/moe_hybrid_storage.h | 8 ++ server/src/common/moe_hybrid_types.h | 4 + 7 files changed, 164 insertions(+), 16 deletions(-) diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index 17a3fa30e2..408a3c1c7e 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,50 @@ 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_MASKED_COMBINE=0 keeps the separate route-mask get_rows and mul, +// the weighting mul and the route-sum repeat_back. +// 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 bool moe_masked_combine_enabled() { + static const bool on = [] { + const char * v = std::getenv("LUCE_MOE_MASKED_COMBINE"); + return !(v && v[0] == '0' && v[1] == '\0'); + }(); + return on; } static void align_moe_owner_routes( @@ -1062,7 +1126,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 +1242,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 +1308,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 +1448,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 moe_masked_combine_enabled() && 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 +1470,30 @@ 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 (on an R9700 behind PCIe 3.0 x4 + // the per-layer handoff grows ~0.25 us per KB), 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 +1882,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 +4492,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 097e9acec8..64bfcac46e 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 efb42a47c9..b577a84861 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 159399b2b3..709f4f8726 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 d46eac8cd0..798556cac0 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 8053113783..fd7896bbea 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 8205fd6282..8700c0a030 100644 --- a/server/src/common/moe_hybrid_types.h +++ b/server/src/common/moe_hybrid_types.h @@ -53,6 +53,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. From af2e85f4ab45f7ef3d65f87f2d55ca656b6a4649 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Thu, 8 Oct 2026 13:08:40 +0200 Subject: [PATCH 03/40] fix(spec): the verify-width controller prices widths it has not measured AdaptiveSpecWidth seeded every width with one device's cost curve and replaced an entry only when it ran that width, so a narrow width measured on faster hardware could price every wider one out for good. Unmeasured widths are now priced from the seed shape scaled by the nearest measured width, never below a measured narrower one; graph-build outliers are clipped out of the samples, and a width's first live sample is held to 1.25x its unmeasured price once another width has been measured. Co-Authored-By: Claude Opus 5.5 --- server/src/common/adaptive_spec_width.h | 129 ++++++++++++++++-- server/test/unit/test_adaptive_spec_width.cpp | 116 +++++++++++++++- 2 files changed, 227 insertions(+), 18 deletions(-) diff --git a/server/src/common/adaptive_spec_width.h b/server/src/common/adaptive_spec_width.h index 4c1c60ccd0..41e8d20eba 100644 --- a/server/src/common/adaptive_spec_width.h +++ b/server/src/common/adaptive_spec_width.h @@ -31,11 +31,16 @@ // are calibrated online: observe_confidence() tracks predicted against // observed acceptance per depth and each score is scaled by their ratio, // so an optimistic head stops buying widths the target does not pay for. -// Step costs are seeded by set_relative_costs() -// and refined online from the observed cost of the offered width; the -// first kCostWarmupSamples observations of a seeded width are ignored so a -// cold shape-specific graph build cannot poison a calibrated seed. A width -// that is not offered keeps its last cost estimate. +// Step costs are seeded by set_relative_costs() and measured online from +// the observed cost of the offered width; the first kCostWarmupSamples +// observations of a seeded width are ignored so a cold shape-specific +// graph build cannot poison the estimate. A seed is a shape, not a +// measurement of this machine: until a width is measured it is priced at +// its seed scaled by the measured/seed ratio of the nearest measured +// width, and never above a measured narrower width, so a wider shape the +// seed curve overprices gets measured (kCostWarmupSamples + 1 steps) +// instead of never being offered. Measured costs survive reset(), and +// measured_costs()/set_measured_costs() carry them across controllers. // // Survival at depths beyond the offered width is not frozen: a rejection // at depth d is a real zero sample for every deeper prefix, and a clean @@ -91,6 +96,12 @@ class AdaptiveSpecWidth { static constexpr float kCostAlpha = 0.20f; // Observations of a seeded width ignored before its cost starts tracking. static constexpr int kCostWarmupSamples = 4; + // A measured width's sample counts at most this multiple of its estimate: + // a step that also built a graph (a new position band, a new request) + // costs several times a steady step, and one such outlier must not make + // the width look expensive for the following dozen steps. A real cost + // rise still lands, at up to kCostAlpha * 25% per step. + static constexpr float kCostOutlierRatio = 1.25f; // Below this prefix survival a depth has no usable support: the ratio of // two near-zero estimates is noise, so the learned conditional acceptance // of the next depth is taken against this floor instead. @@ -124,6 +135,8 @@ class AdaptiveSpecWidth { accepted_candidates_ema_(initial_accepted_candidates_), prefix_survival_ema_((size_t) std::max(1, max_width_), kSurvivalPrior), + width_cost_seed_((size_t) std::max(1, max_width_) + 1, + std::numeric_limits::quiet_NaN()), width_cost_ema_((size_t) std::max(1, max_width_) + 1, std::numeric_limits::quiet_NaN()), width_cost_samples_((size_t) std::max(1, max_width_) + 1, 0), @@ -152,6 +165,21 @@ class AdaptiveSpecWidth { // positive leave the width unseeded. A backend can seed this from a short // calibration and then refine it online through observe(). void set_relative_costs(const std::vector & costs) { + const size_t n = std::min(costs.size(), width_cost_seed_.size()); + for (size_t width = 0; width < n; ++width) { + const float cost = costs[width]; + if (std::isfinite(cost) && cost > 0.0f) { + width_cost_seed_[width] = cost; + } + } + } + + // Measured total step costs indexed by width, in observe() units (NaN + // where a width is unmeasured). Costs belong to the machine, not to the + // text, so a backend that builds a controller per request can hand them + // from one request's controller to the next. + const std::vector & measured_costs() const { return width_cost_ema_; } + void set_measured_costs(const std::vector & costs) { const size_t n = std::min(costs.size(), width_cost_ema_.size()); for (size_t width = 0; width < n; ++width) { const float cost = costs[width]; @@ -182,8 +210,9 @@ class AdaptiveSpecWidth { // candidates 0..i-1 were accepted). With no model confidence, use the // target-observed prefix survival estimates learned by observe(); a // shorter vector is extended with the learned conditional acceptance of - // the depths it does not cover. Widths without a usable cost estimate - // are skipped; when none is usable the acceptance-only policy decides. + // the depths it does not cover. Widths are priced by width_cost(); those + // without a usable cost are skipped, and when none is usable the + // acceptance-only policy decides. int next_width_cost_aware( const std::vector & conditional_acceptance, int proposed_width) const { @@ -217,7 +246,7 @@ class AdaptiveSpecWidth { } expected_commits += survival; if (width < min_width_) continue; - const float cost = width_cost_ema_[(size_t) width]; + const float cost = width_cost(width); if (!std::isfinite(cost) || cost <= 0.0f) continue; const float utility = expected_commits / cost; utilities[(size_t) width] = utility; @@ -318,17 +347,33 @@ class AdaptiveSpecWidth { if (std::isfinite(observed_step_cost) && observed_step_cost > 0.0f) { float & estimate = width_cost_ema_[(size_t) offered]; int & samples = width_cost_samples_[(size_t) offered]; - if (!std::isfinite(estimate) || estimate <= 0.0f) { - // No calibrated seed for this width: adopt the first - // observation and track from the next one on. + const bool measured = std::isfinite(estimate) && estimate > 0.0f; + if (!measured && !has_seed(offered)) { + // No seed for this width: adopt the first observation and + // track from the next one on. estimate = observed_step_cost; samples = kCostWarmupSamples; } else if (++samples > kCostWarmupSamples) { // Shape-specific graph construction makes the first few - // observations of a seeded width cold outliers. Preserve the - // calibrated seed through warmup, then track the live cost. - estimate = (1.0f - kCostAlpha) * estimate + - kCostAlpha * observed_step_cost; + // observations of a width cold outliers (again after reset(): + // a new request builds its own shapes). Skip them, then + // measure the live cost. Every live sample counts at most + // kCostOutlierRatio times the width's price before it: its + // estimate, or for the first one the price width_cost() gave + // it unmeasured, once another width is measured here (a bare + // seed is in another machine's units). A rarely offered width + // still meets graph builds after its warmup (a new position + // band), and a first sample taken there would keep it + // overpriced, and so never offered again, for the process. + const float reference = measured ? estimate + : has_other_measured_cost(offered) ? width_cost(offered) + : std::numeric_limits::quiet_NaN(); + const float sample = std::isfinite(reference) && reference > 0.0f + ? std::min(observed_step_cost, kCostOutlierRatio * reference) + : observed_step_cost; + estimate = measured + ? (1.0f - kCostAlpha) * estimate + kCostAlpha * sample + : sample; } } @@ -433,6 +478,58 @@ class AdaptiveSpecWidth { return std::clamp(deeper / shallower, 0.0f, 1.0f); } + bool has_seed(int width) const { + const float seed = width_cost_seed_[(size_t) width]; + return std::isfinite(seed) && seed > 0.0f; + } + + bool has_measured_cost(int width) const { + const float cost = width_cost_ema_[(size_t) width]; + return std::isfinite(cost) && cost > 0.0f; + } + + bool has_other_measured_cost(int width) const { + for (int other = 1; other <= max_width_; ++other) { + if (other != width && has_measured_cost(other)) return true; + } + return false; + } + + // Total step cost of a width in observe() units, NaN when it has none. + // A measured width answers with its estimate. An unmeasured seeded width + // is priced at its seed scaled by the measured/seed ratio of the nearest + // measured seeded width (the narrower one on a tie), and no higher than + // the widest measured width below it: extra candidates count as free + // until their shape is measured, so a seed curve from another machine + // cannot keep a profitable width from ever being offered. + float width_cost(int width) const { + if (has_measured_cost(width)) return width_cost_ema_[(size_t) width]; + if (!has_seed(width)) return std::numeric_limits::quiet_NaN(); + float cost = width_cost_seed_[(size_t) width]; + for (int distance = 1; distance <= max_width_; ++distance) { + int scale_width = -1; + for (const int other : {width - distance, width + distance}) { + if (other >= 1 && other <= max_width_ && + has_seed(other) && has_measured_cost(other)) { + scale_width = other; + break; + } + } + if (scale_width >= 0) { + cost *= width_cost_ema_[(size_t) scale_width] / + width_cost_seed_[(size_t) scale_width]; + break; + } + } + for (int below = width - 1; below >= 1; --below) { + if (has_measured_cost(below)) { + cost = std::min(cost, width_cost_ema_[(size_t) below]); + break; + } + } + return cost; + } + bool max_width_guard_enabled() const { return max_width_probe_streak_ > 0 && max_width_ > min_width_; } @@ -452,6 +549,8 @@ class AdaptiveSpecWidth { float full_accept_probe_; float accepted_candidates_ema_; std::vector prefix_survival_ema_; + // Seed step-cost curve (a shape) and the measured costs, both by width. + std::vector width_cost_seed_; std::vector width_cost_ema_; std::vector width_cost_samples_; // Per-width expected value of the current cost-aware decision; kept as diff --git a/server/test/unit/test_adaptive_spec_width.cpp b/server/test/unit/test_adaptive_spec_width.cpp index 7e138f597c..6bf801b996 100644 --- a/server/test/unit/test_adaptive_spec_width.cpp +++ b/server/test/unit/test_adaptive_spec_width.cpp @@ -115,15 +115,74 @@ TEST_CASE(AdaptiveSpecWidthFixture, cost_aware_width_learns_prefix_survival) { CHECK(width.next_width_cost_aware({}) == 2); } -TEST_CASE(AdaptiveSpecWidthFixture, observed_cost_updates_only_offered_width) { +TEST_CASE(AdaptiveSpecWidthFixture, measured_cost_rescales_the_unmeasured_seeds) { AdaptiveSpecWidth width(4); width.set_relative_costs({0.0f, 0.0f, 70.0f, 100.0f, 130.0f}); CHECK(width.next_width_cost_aware({1.0f, 1.0f, 1.0f}) == 4); + // q4 measures at twice its seed. The unmeasured q2 and q3 are priced at + // twice their seeds too (140, 200), so q4 keeps its place instead of + // losing to seeds in other units (4/260 against 3/100). for (int i = 0; i < 5; ++i) { width.observe(4, 4, 260.0f); } - CHECK(width.next_width_cost_aware({1.0f, 1.0f, 1.0f}) == 3); + CHECK(near(width.measured_costs()[4], 260.0f)); + CHECK(!std::isfinite(width.measured_costs()[3])); + CHECK(width.next_width_cost_aware({1.0f, 1.0f, 1.0f}) == 4); +} + +TEST_CASE(AdaptiveSpecWidthFixture, seed_curve_is_a_shape_scaled_to_the_machine) { + AdaptiveSpecWidth width(5, 2); + width.set_relative_costs(kDs4StepCosts); + // q5 measures at 40% of its seed: this machine is faster than the + // seed's. The unmeasured q2 is priced at 40% of its seed too (30, not + // 75), so a weak head settles on q2: 1.5/30 beats 2.0/49.2. + for (int i = 0; i <= AdaptiveSpecWidth::kCostWarmupSamples; ++i) { + width.observe(5, 5, 49.2f); + } + REQUIRE(near(width.measured_costs()[5], 49.2f)); + CHECK(width.next_width_cost_aware({0.5f, 0.5f, 0.5f}) == 2); + CHECK(width.next_width_cost_aware({0.95f, 0.95f, 0.95f}) == 5); +} + +TEST_CASE(AdaptiveSpecWidthFixture, wider_width_is_measured_before_the_seed_curve_rules_it_out) { + // The seed curve prices q5 at 1.64x q2; on this machine it costs 1.37x. + AdaptiveSpecWidth width(5, 2); + width.set_relative_costs(kDs4StepCosts); + const std::vector head = {0.8f, 0.75f, 0.7f}; + for (int i = 0; i < 12; ++i) { + width.observe(2, 2, 38.0f); + } + // Scaled from q2, q5 would cost 62 and lose (3.2/62 < 1.8/38). Until + // q3..q5 are measured they cost no more than the measured q2. + CHECK(width.next_width_cost_aware(head) > 2); + for (int w = 3; w <= 5; ++w) { + for (int i = 0; i <= AdaptiveSpecWidth::kCostWarmupSamples; ++i) { + width.observe(w, w, 38.0f + 4.7f * (float) (w - 2)); + } + } + // At the measured prices the head decides: q5 commits 3.2 tokens for 52 + // against 1.8 for 38 at q2, and a weak head goes back to q2. + CHECK(width.next_width_cost_aware(head) == 5); + CHECK(width.next_width_cost_aware({0.25f, 0.25f, 0.25f}) == 2); +} + +TEST_CASE(AdaptiveSpecWidthFixture, measured_costs_carry_across_controllers) { + AdaptiveSpecWidth first(5, 2); + first.set_relative_costs(kDs4StepCosts); + for (int i = 0; i <= AdaptiveSpecWidth::kCostWarmupSamples; ++i) { + first.observe(5, 5, 50.0f); + } + first.reset(); + CHECK(near(first.measured_costs()[5], 50.0f)); + + AdaptiveSpecWidth second(5, 2); + second.set_relative_costs(kDs4StepCosts); + second.set_measured_costs(first.measured_costs()); + CHECK(near(second.measured_costs()[5], 50.0f)); + CHECK(!std::isfinite(second.measured_costs()[2])); + // The carried q5 price rescales the unmeasured q2 at once (30.5). + CHECK(second.next_width_cost_aware({0.5f, 0.5f, 0.5f}) == 2); } TEST_CASE(AdaptiveSpecWidthFixture, nearly_free_fifth_slot_stays_profitable) { @@ -258,22 +317,73 @@ TEST_CASE(AdaptiveSpecWidthFixture, unseeded_width_adopts_first_cost_then_tracks CHECK(width.next_width_cost_aware(certain) == 2); width.observe(3, 3, 1.0f); // adopted as the q3 estimate: 3/1 beats 2/1 CHECK(width.next_width_cost_aware(certain) == 3); - width.observe(3, 3, 100.0f); // no warmup hold without a seed: tracks now + // No warmup hold without a seed: it tracks from here, a lone spike as an + // outlier (at most 1.25x the estimate) and a sustained rise in full. + width.observe(3, 3, 100.0f); + CHECK(width.next_width_cost_aware(certain) == 3); + for (int step = 0; step < 12; ++step) width.observe(3, 3, 100.0f); CHECK(width.next_width_cost_aware(certain) == 2); } +TEST_CASE(AdaptiveSpecWidthFixture, graph_build_outlier_does_not_flip_the_width) { + AdaptiveSpecWidth width(5, 2); + width.set_relative_costs(kDs4StepCosts); + for (int i = 0; i <= AdaptiveSpecWidth::kCostWarmupSamples; ++i) { + width.observe(2, 2, 38.0f); + width.observe(3, 3, 41.0f); + width.observe(4, 4, 44.0f); + width.observe(5, 5, 47.0f); + } + const std::vector head = {0.97f, 0.96f, 0.95f}; + REQUIRE(width.next_width_cost_aware(head) == 5); + // One q5 step that also built a graph costs 3x: the estimate moves 5%. + width.observe(5, 5, 141.0f); + CHECK(width.measured_costs()[5] < 50.0f); + CHECK(width.next_width_cost_aware(head) == 5); +} + TEST_CASE(AdaptiveSpecWidthFixture, seeded_width_holds_its_seed_through_warmup) { AdaptiveSpecWidth width(3); width.set_relative_costs({0.0f, 0.0f, 1.0f, 1.0f}); const std::vector certain = {1.0f, 1.0f}; + for (int sample = 0; sample <= AdaptiveSpecWidth::kCostWarmupSamples; ++sample) { + width.observe(2, 2, 1.0f); // q2 measured at its seed + } for (int sample = 0; sample < AdaptiveSpecWidth::kCostWarmupSamples; ++sample) { width.observe(3, 3, 100.0f); // cold graph-build outliers CHECK(width.next_width_cost_aware(certain) == 3); } + // The first live sample counts at most 1.25x the price q3 had (q2's + // 1.0); a sustained cost then lands at up to 5% a step, and q3 loses to + // q2 once it costs over 1.5. width.observe(3, 3, 100.0f); + CHECK(near(width.measured_costs()[3], 1.25f)); + CHECK(width.next_width_cost_aware(certain) == 3); + for (int step = 0; step < 5; ++step) width.observe(3, 3, 100.0f); CHECK(width.next_width_cost_aware(certain) == 2); } +TEST_CASE(AdaptiveSpecWidthFixture, graph_build_on_a_first_live_sample_does_not_price_a_width_out) { + // Two-box code: q2, q4 and q5 are measured; q3 is rarely offered, and + // the step right after its warmup also built a graph (58.9 ms). + AdaptiveSpecWidth width(5, 2); + width.set_relative_costs(kDs4StepCosts); + for (int i = 0; i <= AdaptiveSpecWidth::kCostWarmupSamples; ++i) { + width.observe(2, 2, 36.3f); + width.observe(4, 4, 44.0f); + width.observe(5, 5, 52.5f); + } + for (int i = 0; i < AdaptiveSpecWidth::kCostWarmupSamples; ++i) { + width.observe(3, 3, 43.7f); + } + width.observe(3, 3, 58.9f); + // Held to 1.25x its unmeasured price (no more than the measured q2). + CHECK(width.measured_costs()[3] <= 1.25f * 36.3f + 1e-3f); + // Its live cost then tracks down from there. + for (int i = 0; i < 12; ++i) width.observe(3, 3, 43.7f); + CHECK(width.measured_costs()[3] < 44.0f); +} + TEST_CASE(AdaptiveSpecWidthFixture, cost_aware_never_exceeds_the_proposal) { AdaptiveSpecWidth width(5, 2, true); width.set_relative_costs({0.0f, 0.0f, 1.0f, 1.0f, 1.0f, 1.0f}); From d8f5c547b34d8236890eafcfb90d2caf97b90e43 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Thu, 8 Oct 2026 13:08:40 +0200 Subject: [PATCH 04/40] feat(cluster): two-box runtime with a raw-RDMA decode exchange A port of maikzz32's lucebox-halo-cluster runtime onto the in-graph collective of #700 and the cold-owner-None hybrid path of #701: - rank bootstrap over TCP (build, model and placement identity checked), a control channel on which rank 0 decides and the workers follow, an RCCL communicator for prefill-sized reductions, a worker main, and the expert placement JSON with its uniform and balanced builders (scripts/cluster/build_expert_placement.py); - fast_reduce, the decode all-reduce for discrete GPUs: raw RDMA verbs between two ranks with stream write/wait flags instead of a host round trip, a zero-copy bf16 wire, and exchange kernels that publish their own completion flags. Its switches treat 0 as off, so cluster mode can turn them on by default; - the RCCL environment of a rank, now with NCCL_PROTO=Simple, the only protocol that runs reliably on the R9700; - server flags --cluster-rank, --cluster-size, --cluster-head, --cluster-ifname, --cluster-ib-hca, --cluster-gid-index, --cluster-expert-placement, --cluster-shared-expert, --cluster-allreduce-dtype, --cluster-timeout-ms, --cluster-verify-hash. Built with -DLUCE_CLUSTER=ON (RCCL and libibverbs). Without it the RCCL communicator and fast_reduce compile as unavailable stubs, so the default build links and a rank falls back to RCCL. test_cluster_unit covers the protocol, placement and control channel in both builds. Co-Authored-By: Claude Opus 5.5 --- server/CMakeLists.txt | 96 +- .../scripts/cluster/build_expert_placement.py | 245 ++++ server/src/cluster/cluster_comm.cpp | 837 +++++++++++ server/src/cluster/cluster_comm.h | 140 ++ server/src/cluster/cluster_comm_kernels.cu | 62 + server/src/cluster/cluster_config.cpp | 286 ++++ server/src/cluster/cluster_config.h | 121 ++ server/src/cluster/cluster_control.cpp | 1069 +++++++++++++++ server/src/cluster/cluster_control.h | 133 ++ server/src/cluster/cluster_decision_hooks.cpp | 318 +++++ server/src/cluster/cluster_decision_hooks.h | 254 ++++ .../src/cluster/cluster_expert_placement.cpp | 617 +++++++++ server/src/cluster/cluster_expert_placement.h | 110 ++ server/src/cluster/cluster_head_backend.cpp | 641 +++++++++ server/src/cluster/cluster_head_backend.h | 204 +++ server/src/cluster/cluster_identity.h | 57 + server/src/cluster/cluster_protocol.cpp | 469 +++++++ server/src/cluster/cluster_protocol.h | 366 +++++ server/src/cluster/cluster_telemetry.h | 52 + server/src/cluster/cluster_worker_main.cpp | 483 +++++++ server/src/cluster/cluster_worker_main.h | 22 + server/src/cluster/fast_reduce.cpp | 1073 +++++++++++++++ server/src/cluster/fast_reduce.h | 116 ++ server/src/cluster/fast_reduce_kernels.cu | 308 +++++ server/src/common/backend_args.h | 3 + server/src/common/backend_factory.cpp | 1 + server/src/common/backend_factory.h | 1 + server/src/common/backend_plan.cpp | 1 + server/src/common/cluster_participant.h | 61 + server/src/common/cluster_view.h | 70 + server/src/common/host_stream_flag.cpp | 53 + server/src/common/host_stream_flag.h | 38 + server/src/common/pinned_stage.h | 151 ++ server/src/internal.h | 8 + server/src/server/server_main.cpp | 133 ++ server/test/unit/test_cluster_unit.cpp | 1220 +++++++++++++++++ 36 files changed, 9818 insertions(+), 1 deletion(-) create mode 100644 server/scripts/cluster/build_expert_placement.py create mode 100644 server/src/cluster/cluster_comm.cpp create mode 100644 server/src/cluster/cluster_comm.h create mode 100644 server/src/cluster/cluster_comm_kernels.cu create mode 100644 server/src/cluster/cluster_config.cpp create mode 100644 server/src/cluster/cluster_config.h create mode 100644 server/src/cluster/cluster_control.cpp create mode 100644 server/src/cluster/cluster_control.h create mode 100644 server/src/cluster/cluster_decision_hooks.cpp create mode 100644 server/src/cluster/cluster_decision_hooks.h create mode 100644 server/src/cluster/cluster_expert_placement.cpp create mode 100644 server/src/cluster/cluster_expert_placement.h create mode 100644 server/src/cluster/cluster_head_backend.cpp create mode 100644 server/src/cluster/cluster_head_backend.h create mode 100644 server/src/cluster/cluster_identity.h create mode 100644 server/src/cluster/cluster_protocol.cpp create mode 100644 server/src/cluster/cluster_protocol.h create mode 100644 server/src/cluster/cluster_telemetry.h create mode 100644 server/src/cluster/cluster_worker_main.cpp create mode 100644 server/src/cluster/cluster_worker_main.h create mode 100644 server/src/cluster/fast_reduce.cpp create mode 100644 server/src/cluster/fast_reduce.h create mode 100644 server/src/cluster/fast_reduce_kernels.cu create mode 100644 server/src/common/cluster_participant.h create mode 100644 server/src/common/cluster_view.h create mode 100644 server/src/common/host_stream_flag.cpp create mode 100644 server/src/common/host_stream_flag.h create mode 100644 server/src/common/pinned_stage.h create mode 100644 server/test/unit/test_cluster_unit.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 5f17c00779..47bc06a243 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 @@ -1373,6 +1439,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}) @@ -1681,6 +1753,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}) @@ -2583,6 +2676,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/scripts/cluster/build_expert_placement.py b/server/scripts/cluster/build_expert_placement.py new file mode 100644 index 0000000000..76b87da8fc --- /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 0000000000..3458e59e90 --- /dev/null +++ b/server/src/cluster/cluster_comm.cpp @@ -0,0 +1,837 @@ +// 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); + +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 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 0000000000..37095bdedd --- /dev/null +++ b/server/src/cluster/cluster_comm.h @@ -0,0 +1,140 @@ +// 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 *; + +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; + + // 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 0000000000..abd8722408 --- /dev/null +++ b/server/src/cluster/cluster_comm_kernels.cu @@ -0,0 +1,62 @@ +// 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. + +#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); +} + +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; +} + +} // namespace luce::cluster diff --git a/server/src/cluster/cluster_config.cpp b/server/src/cluster/cluster_config.cpp new file mode 100644 index 0000000000..e8c9aa0237 --- /dev/null +++ b/server/src/cluster/cluster_config.cpp @@ -0,0 +1,286 @@ +// 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"; + } + if (verify_hash_every < 0) { + return "--cluster-verify-hash must be non-negative, got " + std::to_string(verify_hash_every); + } + 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; + } + if (text == "rank0") { + out = SharedExpertMode::Rank0; + return true; + } + set_err(err, "--cluster-shared-expert expects replicate, shard, or rank0, 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"; + case SharedExpertMode::Rank0: return "rank0"; + } + 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 0000000000..846009e009 --- /dev/null +++ b/server/src/cluster/cluster_config.h @@ -0,0 +1,121 @@ +// 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): features ship as CLI flags. Environment variables are +// reserved for burn-in kill switches (LUCE_CLUSTER_NO_INGRAPH_ALLREDUCE, +// LUCE_CLUSTER_NO_GRAPH_CAPTURE) and debug tracing (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 + Rank0, // only rank 0 computes it and adds it after the 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::Replicate; + + // Reduction and robustness knobs. + AllreduceDType allreduce_dtype = AllreduceDType::Auto; + uint32_t timeout_ms = kClusterDefaultTimeoutMs; // --cluster-timeout-ms + + // Debug: cross-rank hidden-state hash every n steps (0 = off). + int verify_hash_every = 0; + + // --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 0000000000..424b0f3646 --- /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 0000000000..3673b3fbe8 --- /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. all + // HashProbe or RequestReport messages); 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 0000000000..6bc4c55168 --- /dev/null +++ b/server/src/cluster/cluster_decision_hooks.cpp @@ -0,0 +1,318 @@ +// cluster_decision_hooks.cpp - head / worker / local decision hooks. + +#include "cluster/cluster_decision_hooks.h" +#include "cluster/cluster_control.h" + +#include +#include +#include + +namespace luce::cluster { + +namespace { + +int32_t greedy_argmax(const float * logits, int n_vocab) { + if (!logits || n_vocab <= 0) return -1; + int32_t best = 0; + float best_v = logits[0]; + for (int i = 1; i < n_vocab; i++) { + if (logits[i] > best_v) { + best_v = logits[i]; + best = i; + } + } + return best; +} + +void set_err(std::string * err, const std::string & msg) { + if (err) *err = msg; +} + +} // namespace + +uint64_t hash_floats(const float * data, size_t n) { + if (!data || n == 0) return fnv1a64(nullptr, 0); + return fnv1a64(data, n * sizeof(float)); +} + +uint64_t hash_tokens(const std::vector & tokens) { + if (tokens.empty()) return fnv1a64(nullptr, 0); + return fnv1a64(tokens.data(), tokens.size() * sizeof(int32_t)); +} + +// ─── HeadHooks ───────────────────────────────────────────────────────── + +HeadHooks::HeadHooks(ClusterHeadControl & control, const HeadHooksConfig & cfg) + : control_(control), cfg_(cfg) {} + +void HeadHooks::reset_counters() { + Ds4ClusterHooks::reset_counters(); + first_mismatch_rank_ = -1; + first_mismatch_step_ = 0; +} + +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) { + if (token < 0) { + // Caller did not sample host-side. Only rank 0 owns an RNG. + if (sampler_) { + token = sampler_(logits, n_vocab, cfg); + } else if (cfg.temp > 0.0f) { + static thread_local std::mt19937_64 rng{cfg.seed != 0 ? cfg.seed : 0x5eedULL}; + std::vector no_history; + token = common::sample_logits(logits, n_vocab, cfg, no_history, rng); + } else { + token = greedy_argmax(logits, n_vocab); + } + if (token < 0) { + set_err(err, "head: sampling produced no token"); + 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::hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, + std::string * err) { + std::vector frames; + const uint64_t t0 = monotonic_us(); + const bool ok = control_.gather(MsgType::HashProbe, cfg_.timeout_ms, frames, err); + counters_.ctrl_wait_us += monotonic_us() - t0; + if (!ok) return false; + counters_.hash_probes++; + + bool mismatch = false; + for (size_t i = 0; i < frames.size(); i++) { + HashProbeMsg probe; + if (!parse_frame(frames[i], probe)) { + set_err(err, "head: malformed HashProbe from rank " + std::to_string(i + 1)); + return false; + } + if (probe.request_id != req || probe.step != step) { + set_err(err, "head: HashProbe from rank " + std::to_string(probe.rank) + + " is for req=" + std::to_string(probe.request_id) + + " step=" + std::to_string(probe.step) + + ", expected req=" + std::to_string(req) + + " step=" + std::to_string(step)); + return false; + } + const bool hc_diff = probe.hc_state_hash != hc_hash; + const bool tok_diff = probe.token_hash != tok_hash; + if (hc_diff || tok_diff) { + mismatch = true; + counters_.hash_mismatches++; + if (first_mismatch_rank_ < 0) { + first_mismatch_rank_ = probe.rank; + first_mismatch_step_ = step; + } + std::fprintf(stderr, + "[cluster] HASH MISMATCH req=%llu step=%u rank=%d: hc %016llx vs head %016llx%s, " + "tokens %016llx vs head %016llx%s%s\n", + (unsigned long long) req, step, probe.rank, + (unsigned long long) probe.hc_state_hash, (unsigned long long) hc_hash, + hc_diff ? " (DIFF)" : "", + (unsigned long long) probe.token_hash, (unsigned long long) tok_hash, + tok_diff ? " (DIFF)" : "", + probe.first_divergent_layer >= 0 + ? (" first_divergent_layer=" + std::to_string(probe.first_divergent_layer)).c_str() + : ""); + } else if (cfg_.trace) { + std::fprintf(stderr, "[cluster] hash ok req=%llu step=%u rank=%d hc=%016llx tok=%016llx\n", + (unsigned long long) req, step, probe.rank, + (unsigned long long) hc_hash, (unsigned long long) tok_hash); + } + } + if (mismatch && cfg_.strict_hash) { + set_err(err, "hash mismatch at step " + std::to_string(step) + + " (first divergent rank " + std::to_string(first_mismatch_rank_) + ")"); + return false; + } + 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; +} + +bool WorkerHooks::hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, + std::string * err) { + HashProbeMsg msg; + msg.request_id = req; + msg.step = step; + msg.rank = cfg_.rank; + msg.hc_state_hash = hc_hash; + msg.token_hash = tok_hash; + msg.first_divergent_layer = -1; // per-layer attribution is a LUCE_CLUSTER_TRACE follow-up + const uint64_t t0 = monotonic_us(); + const bool ok = control_.send(make_frame(msg), err); + counters_.ctrl_wait_us += monotonic_us() - t0; + if (!ok) return false; + counters_.hash_probes++; + 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 0000000000..5639e71e92 --- /dev/null +++ b/server/src/cluster/cluster_decision_hooks.h @@ -0,0 +1,254 @@ +// 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 and, every +// verify_hash_every steps, gathers HashProbe frames from the +// workers and compares them with its own hash. +// 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 +#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 hash_probes = 0; + uint64_t hash_mismatches = 0; // head only + 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` (including any budget-hook override) + // and computed `flags` with make_decision_flags. If the caller passes + // token < 0 the head hook samples itself (greedy argmax, or `cfg` via + // common::sample_logits with an internal RNG when cfg.temp > 0); this + // path exists for callers that do not sample host-side. + // 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; + + // Determinism probe (--cluster-verify-hash n). Worker sends its hashes; + // head gathers one HashProbe per worker and compares against its own. + // A mismatch is logged (first divergent rank and step) and counted; it + // is fatal only when strict mode is enabled on the head hooks. + virtual bool hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, + 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; + + // Steps between hash probes; 0 = off. + virtual int verify_hash_every() const { return 0; } + + virtual const HookCounters & counters() const { return counters_; } + virtual 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 hash_probe(uint64_t, uint32_t, uint64_t, uint64_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 ───────────────────────────────────────────────────────── +struct HeadHooksConfig { + uint32_t timeout_ms = 30000; // deadline for gather() + int verify_hash_every = 0; + bool strict_hash = false; // mismatch -> return false + bool trace = false; // LUCE_CLUSTER_TRACE +}; + +class HeadHooks final : public Ds4ClusterHooks { +public: + // `control` must outlive the hooks. Optional sampler used only when the + // caller passes token < 0. + HeadHooks(ClusterHeadControl & control, const HeadHooksConfig & cfg); + + using Sampler = std::function; + void set_sampler(Sampler s) { sampler_ = std::move(s); } + + 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 hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, + 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; } + int verify_hash_every() const override { return cfg_.verify_hash_every; } + + // First mismatch seen since reset_counters(); rank -1 = none. + int first_mismatch_rank() const { return first_mismatch_rank_; } + uint32_t first_mismatch_step() const { return first_mismatch_step_; } + void reset_counters() override; + +private: + ClusterHeadControl & control_; + HeadHooksConfig cfg_; + Sampler sampler_; + int first_mismatch_rank_ = -1; + uint32_t first_mismatch_step_ = 0; +}; + +// ─── WorkerHooks ─────────────────────────────────────────────────────── +struct WorkerHooksConfig { + int rank = -1; + uint32_t timeout_ms = 30000; // deadline for every blocking recv + int verify_hash_every = 0; +}; + +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 hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, + 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; } + int verify_hash_every() const override { return cfg_.verify_hash_every; } + + // 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_; +}; + +// FNV-1a over a float buffer / token vector; shared by backend and tests. +uint64_t hash_floats(const float * data, size_t n); +uint64_t hash_tokens(const std::vector & tokens); + +} // 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 0000000000..ad64c4a987 --- /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 0000000000..0a4b557274 --- /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 0000000000..2a3b67ad53 --- /dev/null +++ b/server/src/cluster/cluster_head_backend.cpp @@ -0,0 +1,641 @@ +// 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 + +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()); + // Launcher contract (scripts/cluster/launch_cluster.sh waits for this + // exact stdout prefix before it starts 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.verify_hash_every = cfg_.verify_hash_every; + 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); + } + + HeadHooksConfig hcfg; + hcfg.timeout_ms = cfg_.timeout_ms; + hcfg.verify_hash_every = cfg_.verify_hash_every; + hcfg.strict_hash = false; + hcfg.trace = cluster_env_trace(); + hooks_ = std::make_unique(control_, hcfg); + + // 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 verify_hash_every=%d\n", + cfg_.size, control_.n_workers(), rccl_version_string().c_str(), + cfg_.verify_hash_every); + std::fflush(stdout); +} + +// ─── Failure ──────────────────────────────────────────────────────────── + +void ClusterHeadBackend::fail_cluster(const std::string & reason, int code, uint64_t request_id) { + if (fatal_) return; + fatal_ = true; + { + std::lock_guard lk(fault_mu_); + if (fault_reason_.empty()) fault_reason_ = reason; + } + 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(); + if (exit_on_abort_) { + std::fprintf(stderr, "[cluster] head: exiting so the supervisor restarts all ranks\n"); + std::fflush(nullptr); + std::_Exit(3); + } +} + +// ─── M4: 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; + { + std::lock_guard lk(fault_mu_); + if (fault_reason_.empty()) fault_reason_ = reason; + } + // 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); + if (!exit_on_abort_) return; + // 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; + + std::string err; + if (!control_.broadcast(make_frame(msg), &err)) { + fail_cluster("request broadcast: " + err, kAbortCodeControl, request_id); + result.fail(GenerateErrorCode::DecodeFailed, "cluster: request broadcast failed: " + err); + return result; + } + + 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, req, io) + : inner_->generate_impl(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)))) { + fail_cluster("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)) { + fail_cluster("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)) { + fail_cluster("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; + last_report_.hash_probes = hooks_->counters().hash_probes; + last_report_.hash_mismatches = hooks_->counters().hash_mismatches; + last_report_.first_mismatch_rank = hooks_->first_mismatch_rank(); + last_report_.first_mismatch_step = hooks_->first_mismatch_step(); + + 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)) { + fail_cluster("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 (M4: ship it). + (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 0000000000..6b5672a0c1 --- /dev/null +++ b/server/src/cluster/cluster_head_backend.h @@ -0,0 +1,204 @@ +// 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 into `last_report()`. +// +// Failure policy (M1): any control-channel or collective failure broadcasts +// Abort, aborts the communicator and, unless exit_on_abort is false, ends +// the process with a non-zero status so a supervisor restarts all ranks +// together. Communicator re-initialization is M4. + +#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_view.h" +#include "common/model_backend.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; + uint64_t hash_probes = 0; + uint64_t hash_mismatches = 0; + int first_mismatch_rank = -1; + uint32_t first_mismatch_step = 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; + + // ── Cluster introspection ──────────────────────────────────────── + // Plain snapshots for the HTTP layer (common/cluster_view.h). + // telemetry describes the request that just finished and is read on the + // request thread; props is constant after init(). + + const ClusterRequestReport & last_report() const { return last_report_; } + const ClusterConfig & config() const { return cfg_; } + const std::vector & worker_identities() const { return hellos_; } + IClusterComm * comm() { return comm_.get(); } + bool fatal() const { return fatal_; } + void set_exit_on_abort(bool v) { exit_on_abort_ = v; } + + // 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. + 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}; + // Why the cluster died, set by whichever thread noticed. Read by + // cluster_request_telemetry so the failed request's response carries it + // in usage.timings.cluster.error instead of looking like an empty answer. + mutable std::mutex fault_mu_; + std::string fault_reason_; + 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 exit_on_abort_ = true; + 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 0000000000..017e70a31b --- /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 0000000000..e6b059337f --- /dev/null +++ b/server/src/cluster/cluster_protocol.cpp @@ -0,0 +1,469 @@ +// 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::HashProbe: return "HashProbe"; + 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.i32(verify_hash_every); + 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) || !r.i32(out.verify_hash_every)) { + 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 HashProbeMsg::encode(ByteWriter & w) const { + w.u64(request_id); + w.u32(step); + w.i32(rank); + w.u64(hc_state_hash); + w.u64(token_hash); + w.i32(first_divergent_layer); +} + +bool HashProbeMsg::decode(ByteReader & r, HashProbeMsg & out) { + return r.u64(out.request_id) && r.u32(out.step) && r.i32(out.rank) && + r.u64(out.hc_state_hash) && r.u64(out.token_hash) && + r.i32(out.first_divergent_layer); +} + +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 0000000000..351fc289ce --- /dev/null +++ b/server/src/cluster/cluster_protocol.h @@ -0,0 +1,366 @@ +// 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" +// 2: RequestMsg carries the inline snapshot position and separates the +// restore slot from the inline save slot, so the prefix cache can be +// replicated (a request may restore from one slot and check-point into +// another; before this the workers never took the inline snapshot at all). +inline constexpr uint16_t kProtocolVersion = 2; +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 + HashProbe = 9, // worker -> head (and head local), determinism check + 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 + int32_t verify_hash_every = 0; + 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); +}; + +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); +}; + +// Determinism probe (--cluster-verify-hash n). +struct HashProbeMsg { + uint64_t request_id = 0; + uint32_t step = 0; + int32_t rank = -1; + uint64_t hc_state_hash = 0; // FNV-1a over the final-layer HC state bytes + uint64_t token_hash = 0; // FNV-1a over accepted token ids so far + int32_t first_divergent_layer = -1; // filled only under LUCE_CLUSTER_TRACE + + static MsgType type() { return MsgType::HashProbe; } + void encode(ByteWriter & w) const; + static bool decode(ByteReader & r, HashProbeMsg & 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 hashes and determinism probes. +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 0000000000..d8534d4c37 --- /dev/null +++ b/server/src/cluster/cluster_telemetry.h @@ -0,0 +1,52 @@ +// 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 (Agent C) and the +// decision hooks / worker main (Agent D) can accumulate into the same struct; +// the head folds the per-rank totals into RequestReportMsg and +// usage.timings.cluster. + +#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 0000000000..8fcb07d586 --- /dev/null +++ b/server/src/cluster/cluster_worker_main.cpp @@ -0,0 +1,483 @@ +// 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 = m.force_ar || m.decode_mode == DecodeMode::Autoregressive; + // 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.verify_hash_every = welcome.verify_hash_every; + 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.rank = st.cfg.rank; + hcfg.timeout_ms = st.cfg.timeout_ms; + hcfg.verify_hash_every = st.cfg.verify_hash_every; + 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 0000000000..28f4c46be2 --- /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 +// (maikzz32/lucebox-halo-cluster). 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 0000000000..02f3647714 --- /dev/null +++ b/server/src/cluster/fast_reduce.cpp @@ -0,0 +1,1073 @@ +#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 +#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). +void fast_reduce_launch_pack(const float * src, void * dst, int n, int bf16, hipStream_t stream); +void fast_reduce_launch_pack_pub(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_hyb_combine_pub(const float * cold, const float * hot, const float * shared, void * own_row, + int n, int bf16, uint32_t * flag, uint32_t value, uint32_t * counter, + hipStream_t stream); +void fast_reduce_launch_hyb_final_pub(const void * own_row, const void * peer_row, float * dst, int n, int bf16, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream); +void fast_reduce_launch_add_zc(float * dst, const void * peer, int n, int bf16, hipStream_t stream); +void fast_reduce_launch_flag_sys(uint32_t * flag, uint32_t value, hipStream_t stream); +void fast_reduce_launch_hyb_combine(const float * cold, const float * hot, const float * shared, void * own_row, + int n, int bf16, hipStream_t stream); +void fast_reduce_launch_hyb_final(const void * own_row, const void * peer_row, float * dst, int n, int bf16, + hipStream_t stream); +void fast_reduce_launch_set_flag(uint32_t * pub_flag, uint32_t * slot_done, + uint32_t * progress, uint32_t seq, + uint32_t slot_span, uint64_t spin_limit, + hipStream_t stream); +void fast_reduce_launch_wait_flags(uint32_t * const * peer_flags, + uint32_t * timed_out, uint32_t * progress, + uint32_t seq, int n_peers, + uint64_t spin_limit, hipStream_t stream); +void fast_reduce_launch_add(float * dst, const float * scratch, int n_peers, + int n, int stride, uint32_t * slot_done, + uint32_t * progress, uint32_t seq, + 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; +} + +bool write_all(int fd, const void * p, size_t n) { + const char * c = (const char *) p; + while (n) { + const ssize_t w = ::write(fd, c, n); + if (w <= 0) return false; + c += w; + n -= (size_t) w; + } + return true; +} + +bool read_all(int fd, void * p, size_t n) { + char * c = (char *) p; + while (n) { + const ssize_t r = ::read(fd, c, n); + if (r <= 0) return false; + c += r; + n -= (size_t) r; + } + return true; +} + +// 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; + }; + + 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"); } + + std::vector peers((size_t) cfg.size, -1); + for (int i = 1; i < cfg.size; ++i) { + const int fd = ::accept(lst, nullptr, nullptr); + if (fd < 0) { ::close(lst); return fail("accept"); } + set_nodelay(fd); + int32_t r = -1; + Endpoint e{}; + if (!read_all(fd, &r, sizeof(r)) || !read_all(fd, &e, sizeof(e)) || + r <= 0 || r >= cfg.size) { + ::close(fd); ::close(lst); + return fail("read endpoint"); + } + all[(size_t) r] = e; + peers[(size_t) r] = fd; + } + ::close(lst); + for (int i = 1; i < cfg.size; ++i) { + if (!write_all(peers[(size_t) i], all.data(), sizeof(Endpoint) * all.size())) { + for (int fd : peers) if (fd >= 0) ::close(fd); + return fail("write table"); + } + } + for (int fd : peers) if (fd >= 0) ::close(fd); + 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; + if (!write_all(fd, &r, sizeof(r)) || !write_all(fd, &mine, sizeof(mine)) || + !read_all(fd, all.data(), sizeof(Endpoint) * all.size())) { + ::close(fd); + return fail("exchange"); + } + ::close(fd); + return true; +} + +} // namespace + +// A switch is on when set to anything but empty or "0". Cluster mode sets the +// measured configuration's switches by default (launch_profiles.h), so "0" is +// how an operator turns one off. +static bool fr_env_on(const char * name) { + const char * v = std::getenv(name); + return v && *v && std::strcmp(v, "0") != 0; +} + +// LUCE_CLUSTER_FAST_REDUCE_LEAN=1 drops the per-reduction slot-reuse waits and +// done writes on the zero-copy and hybrid paths. Neither carries data: 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 256 reductions earlier, and the GPU's own stream orders its +// writes behind its reads. On a discrete GPU each of these stream operations +// is a blit kernel with its own dispatch gap, four per layer. +static bool fast_reduce_lean() { + static const bool lean = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_LEAN"); + return lean; +} + +// The kernel that writes an exchange row also publishes its flag (one +// system-scope release by its grid's last block) instead of a +// hipStreamWriteValue32 after it: about 0.6 ms less per two-box verify step, +// same bits. LUCE_CLUSTER_FAST_REDUCE_KPUB=0 goes back to the stream writes. +static bool fast_reduce_kpub() { + static const bool on = [] { + const char * v = std::getenv("LUCE_CLUSTER_FAST_REDUCE_KPUB"); + return !(v && v[0] == '0' && v[1] == '\0'); + }(); + return on; +} + +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] done + uint32_t * flag_dev = nullptr; + + // Device-visible pointer table for the flags, one row per slot. The + // payloads need no such table: they are moved by the copy engine, which + // takes its addresses from the host. + uint32_t ** peer_flags = nullptr; // [slots][size-1] + uint32_t ** peer_flags_dev = nullptr; + + // Device memory the peers' partials are copied into before they are added. + // Adding straight out of the pinned staging buffer is what made the first + // version 250 times slower than the collective it replaces: that memory is + // uncached for the GPU. + float * scratch = nullptr; // [slots][size-1][max_elems] + + uint32_t * timed_out_host = nullptr; + uint32_t * timed_out_dev = nullptr; + // seq*4 + {1 published, 2 waiting, 3 done}: which reduction the GPU is in + // and what it is doing there. + uint32_t * progress_host = nullptr; + uint32_t * progress_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 = hot | part 1 = shared | 2 = final][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; + } + size_t done_off(int slot) const { + return pub_off(0) + (size_t) cfg.slots + (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); } + uint32_t * done_flag(int slot) { return flag_host + done_off(slot); } + size_t hybh_off(int slot) const { return done_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 + 4u * (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; } +uint64_t FastReduce::submitted() const { return p_ ? p_->seq.load() : 0; } +uint64_t FastReduce::timed_out() const { + return (p_ && p_->timed_out_host) ? *p_->timed_out_host : 0; +} + +bool FastReduce::init(const Config & cfg, std::string * err) { + auto & s = *p_; + s.cfg = cfg; + if (cfg.size <= 1) { s.up = false; return true; } + + 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 (hipHostMalloc((void **) &s.timed_out_host, sizeof(uint32_t) * 16, + hipHostMallocDefault) != hipSuccess) { + return bail("hipHostMalloc for the timeout word failed"); + } + *s.timed_out_host = 0; + if (hipHostMalloc((void **) &s.progress_host, sizeof(uint32_t) * 16, + hipHostMallocDefault) != hipSuccess) { + return bail("hipHostMalloc for the progress word failed"); + } + *s.progress_host = 0; + if (fr_env_on("LUCE_CLUSTER_HYBRID_EXCHANGE") && cfg.size == 2) { + // 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) 3 * (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 || + hipHostGetDevicePointer((void **) &s.timed_out_dev, s.timed_out_host, 0) != hipSuccess || + hipHostGetDevicePointer((void **) &s.progress_dev, s.progress_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; + + for (int r = 0; r < cfg.size; ++r) { + if (r == cfg.rank) continue; + ibv_qp_attr a{}; + a.qp_state = IBV_QPS_RTR; + // LUCE_CLUSTER_FAST_REDUCE_MTU=1024|2048|4096 (default 4096: the direct ConnectX link runs 4096). + { + static const int mtu_env = std::getenv("LUCE_CLUSTER_FAST_REDUCE_MTU") ? std::atoi(std::getenv("LUCE_CLUSTER_FAST_REDUCE_MTU")) : 4096; + a.path_mtu = mtu_env >= 4096 ? IBV_MTU_4096 : (mtu_env >= 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"); + } + } + + // ── pointer tables, one row per slot ────────────────────────────────── + const int n_peers = cfg.size - 1; + const size_t rows = (size_t) cfg.slots * (size_t) n_peers; + if (hipHostMalloc((void **) &s.peer_flags, rows * sizeof(uint32_t *), + hipHostMallocDefault) != hipSuccess) { + return bail("hipHostMalloc for the flag table failed"); + } + for (int slot = 0; slot < cfg.slots; ++slot) { + int k = 0; + for (int r = 0; r < cfg.size; ++r) { + if (r == cfg.rank) continue; + s.peer_flags[(size_t) slot * (size_t) n_peers + (size_t) k] = + s.flag_dev + s.visible_off(slot, r); + ++k; + } + } + if (hipHostGetDevicePointer((void **) &s.peer_flags_dev, s.peer_flags, 0) != hipSuccess) { + return bail("hipHostGetDevicePointer for the flag table failed"); + } + if (hipMalloc((void **) &s.scratch, + rows * (size_t) cfg.max_elems * sizeof(float)) != hipSuccess) { + return bail("hipMalloc for the peer scratch failed"); + } + + s.pending_n.assign((size_t) FastReduceImpl::kPendingRing, 0); + + // ── progress ────────────────────────────────────────────────────────── + s.up = true; + s.progress = std::thread([&s]() { + const int n_peers_l = s.cfg.size - 1; + 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; + } + }; + 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(); + static const double stall_s = std::getenv("LUCE_CLUSTER_FAST_REDUCE_STALL_S") + ? std::atof(std::getenv("LUCE_CLUSTER_FAST_REDUCE_STALL_S")) : 20.0; + while (*pub != (uint32_t) next) { + if (s.stop.load(std::memory_order_relaxed)) return; + if (++spins < 200000ull) continue; + spins = 0; + std::this_thread::yield(); + { + const double w = std::chrono::duration( + std::chrono::steady_clock::now() - wait_t0).count(); + if (w > stall_s && s.seq.load() >= next) { // allocated, never published + s.failed.store(true); + const uint64_t last = s.seq.load(); + 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) (last - next)); + const auto t0 = std::chrono::steady_clock::now(); + while (std::chrono::duration(std::chrono::steady_clock::now() - t0).count() < 10.0) { + for (uint64_t q = next; q <= s.seq.load(); ++q) release_one(q); + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + std::fflush(nullptr); + std::_Exit(4); + } + } + if (next <= 1) continue; + // The GPU published reduction p = next-1 and is waiting on the + // peers' flags for it. Done long ago is the normal idle case; + // published but not done means a peer never delivered. + const uint64_t p = next - 1; + const int ps = (int) (p % (uint64_t) s.cfg.slots); + if (*(volatile uint32_t *) s.done_flag(ps) == (uint32_t) p) continue; + if (fast_reduce_lean()) { + bool arrived = true; + for (int r = 0; r < s.cfg.size; ++r) { + if (r == s.cfg.rank) continue; + const uint32_t v = *(volatile uint32_t *) s.visible_flag(ps, r); + if (v != (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 done; 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 (waited < stall_s) continue; + // Stall: release every reduction the host has handed the GPU so + // its queue drains, make later submits no-ops, and exit once the + // GPU is idle. Exiting with a parked queue costs a GPU reset. + s.failed.store(true); + 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)); + for (uint64_t q = p; q <= std::max(p, last); ++q) release_one(q); + 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 l2 = std::max(p, s.seq.load()); + for (uint64_t q = p; q <= l2; ++q) release_one(q); // late submits too + if (*(volatile uint32_t *) s.done_flag((int) (l2 % (uint64_t) s.cfg.slots)) == (uint32_t) l2) break; + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + std::fprintf(stderr, "[fast-reduce] rank %d: GPU drained, exiting\n", s.cfg.rank); + std::fflush(nullptr); + std::_Exit(4); + } + 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); + static const bool bf16_wire = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16") && + fr_env_on("LUCE_CLUSTER_FAST_REDUCE_ZC"); + sge.length = (uint32_t) ((size_t) n * (bf16_wire && !exact_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) { + s.stop.store(true); + return; + } + } + // 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 it stops the thread and the + // waiting kernels time out into a counted, visible failure. + 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)\n", + ibv_wc_status_str(wc[(size_t) i].status), + (unsigned long long) wc[(size_t) i].wr_id); + s.stop.store(true); + return; + } + } + } 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), " + "gpu timeouts up to seq %u\n", + (unsigned long long) next, n, + *(volatile uint32_t *) s.timed_out_host); + } + ++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 + // The first handful of decisions, taken or declined, so a stall says which + // reduction it stopped at rather than only that it stopped. + static std::atomic traced{0}; + const bool trace = traced.fetch_add(1, std::memory_order_relaxed) < 24; + if (!s.up || n == 0 || n > (size_t) s.cfg.max_elems) { + if (trace) { + std::fprintf(stderr, "[fast-reduce] declined %zu floats (cap %d)\n", + n, s.cfg.max_elems); + } + return false; + } + if (trace) std::fprintf(stderr, "[fast-reduce] took %zu floats\n", n); + + 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; + + const int n_peers = s.cfg.size - 1; + auto stm = (hipStream_t) stream; + + // Loads, not cycles: roughly 100 ms of system-scope reads. What it waits + // on takes about fourteen microseconds, so this only fires when something + // is wrong -- and when it does it has to fire, not spin forever. + constexpr uint64_t kSpin = 200000ull; + + // Publish and wait are stream operations, not kernels. + // + // A kernel cannot do this cheaply on gfx1151. Raising a flag the host will + // see needs a system-scope release, which writes the L2 back; spinning on + // one the host raised needs a system-scope acquire, which invalidates it. + // Either way the weights the next layer is about to read are gone, and + // measured that cost 34 ms per reduction -- while the same path with the + // flags removed ran the step in 32.5 ms. Weakening the ordering does not + // help: relaxed is cheap and the host then does not see the flag at all. + // + // hipStreamWriteValue32 and hipStreamWaitValue32 are handled by the + // command processor instead, so the shader cores and their caches are not + // involved at all -- which is what this needs and what a kernel cannot be. + // LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS=1 restores the kernel form for + // comparison. + static const bool kernel_flags = + std::getenv("LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS") != nullptr; + + // Do not write into a slot the progress thread has not finished with. The + // kernel form waited on this; the stream form has to as well, or the GPU -- + // which runs a whole step's reductions ahead of the host -- laps the ring + // and overwrites a payload that has not been sent. That is what made the + // third request of a run stop rather than the first. + // The previous user of this slot was seq - slots, and it leaves exactly that + // value behind. Asking for anything higher is a wait nothing can satisfy. + static const bool zc_lean = fast_reduce_lean() && fr_env_on("LUCE_CLUSTER_FAST_REDUCE_ZC") && + s.cfg.size == 2; + if (!kernel_flags && !zc_lean && seq > (uint64_t) s.cfg.slots) { + const uint32_t need = (uint32_t) (seq - (uint64_t) s.cfg.slots); + if (hipStreamWaitValue32(stm, s.flag_dev + s.done_off(slot), need, + hipStreamWaitValueGte, 0xffffffffu) != hipSuccess) { + return false; + } + } + + // Zero-copy (discrete GPU): pack straight into the pinned row, publish, wait + // for the peer, add its row straight out of pinned memory, mark the slot done. + static const bool zc = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_ZC"); + static const int zc_bf16 = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16") ? 1 : 0; + const int wire_bf16 = exact ? 0 : zc_bf16; + if (zc && !kernel_flags && s.cfg.size == 2) { + float * own_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) s.cfg.rank * s.cfg.max_elems; + const int peer = 1 - s.cfg.rank; + const float * peer_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) peer * s.cfg.max_elems; + if (fast_reduce_kpub()) { + uint32_t * ctr = s.pub_counter(stm, slot, 0); + if (!ctr) return false; + fast_reduce_launch_pack_pub(data, own_row, (int) n, wire_bf16, + s.flag_dev + s.pub_off(slot), (uint32_t) seq, ctr, stm); + } else { + fast_reduce_launch_pack(data, own_row, (int) n, wire_bf16, stm); + if (hipStreamWriteValue32(stm, s.flag_dev + s.pub_off(slot), (uint32_t) seq, 0) != hipSuccess) return false; + } + if (hipStreamWaitValue32(stm, s.flag_dev + s.visible_off(slot, peer), (uint32_t) seq, + hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; + fast_reduce_launch_add_zc(data, peer_row, (int) n, wire_bf16, stm); + if (zc_lean) return true; + return hipStreamWriteValue32(stm, s.flag_dev + s.done_off(slot), (uint32_t) seq, 0) == hipSuccess; + } + + // Publish: the copy engine moves the partial into the pinned staging + // buffer, then the flag is written. Both are on the stream, so the flag + // cannot precede the data it announces. + if (hipMemcpyAsync(s.data_host + (size_t) slot * s.slot_stride() + + (size_t) s.cfg.rank * s.cfg.max_elems, + data, n * sizeof(float), hipMemcpyDefault, stm) != hipSuccess) { + return false; + } + if (kernel_flags) { + fast_reduce_launch_set_flag(s.flag_dev + s.pub_off(slot), + s.flag_dev + s.done_off(slot), + s.progress_dev, (uint32_t) seq, + (uint32_t) (s.cfg.slots - 2), kSpin, stm); + } else if (hipStreamWriteValue32(stm, s.flag_dev + s.pub_off(slot), + (uint32_t) seq, 0) != hipSuccess) { + return false; + } + + // Wait for the peers, then bring their partials into device memory before + // touching them: the staging buffers are uncached for the GPU, and adding + // straight out of them is what made the first version 250 times slower + // than the collective it replaces. + if (kernel_flags) { + fast_reduce_launch_wait_flags(s.peer_flags_dev + (size_t) slot * (size_t) n_peers, + s.timed_out_dev, s.progress_dev, + (uint32_t) seq, n_peers, kSpin, stm); + } else { + for (int r = 0; r < s.cfg.size; ++r) { + if (r == s.cfg.rank) continue; + if (hipStreamWaitValue32(stm, s.flag_dev + s.visible_off(slot, r), + (uint32_t) seq, hipStreamWaitValueEq, + 0xffffffffu) != hipSuccess) { + return false; + } + } + } + float * scratch_row = s.scratch + + ((size_t) slot * (size_t) n_peers) * (size_t) s.cfg.max_elems; + int k = 0; + for (int r = 0; r < s.cfg.size; ++r) { + if (r == s.cfg.rank) continue; + if (hipMemcpyAsync(scratch_row + (size_t) k * (size_t) s.cfg.max_elems, + s.data_host + (size_t) slot * s.slot_stride() + + (size_t) r * s.cfg.max_elems, + n * sizeof(float), hipMemcpyDefault, stm) != hipSuccess) { + return false; + } + ++k; + } + fast_reduce_launch_add(data, scratch_row, n_peers, (int) n, s.cfg.max_elems, + !kernel_flags ? nullptr : s.flag_dev + s.done_off(slot), + s.progress_dev, (uint32_t) seq, stm); + if (!kernel_flags && + hipStreamWriteValue32(stm, s.flag_dev + s.done_off(slot), + (uint32_t) seq, 0) != hipSuccess) { + return false; + } + 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, int part, 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 || part < 0 || part > 1) 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 (part == 0 && !fast_reduce_lean() && seq > (uint64_t) s.cfg.slots) { + const uint32_t need = (uint32_t) (seq - (uint64_t) s.cfg.slots); + if (hipStreamWaitValue32(stm, s.flag_dev + s.done_off(slot), need, + hipStreamWaitValueGte, 0xffffffffu) != hipSuccess) return false; + } + static const bool one_export = fr_env_on("LUCE_CLUSTER_HYBRID_ONE_EXPORT"); + const bool last_part = part == 1 || one_export; + if (last_part && fast_reduce_kpub()) { + uint32_t * ctr = s.pub_counter(stm, slot, 1); + if (!ctr) return false; + fast_reduce_launch_pack_pub(data, s.hyb_part(part, slot), (int) n, 0, + s.flag_dev + s.hybh_off(slot), (uint32_t) seq, ctr, stm); + return true; + } + fast_reduce_launch_pack(data, s.hyb_part(part, slot), (int) n, 0, stm); + if (last_part) { + return hipStreamWriteValue32(stm, s.flag_dev + s.hybh_off(slot), (uint32_t) seq, 0) == hipSuccess; + } + 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; + static const int zc_bf16 = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16") ? 1 : 0; + 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; + static const bool one_export = fr_env_on("LUCE_CLUSTER_HYBRID_ONE_EXPORT"); + if (fast_reduce_kpub()) { + 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_pub(data, s.hyb_part(0, slot), one_export ? nullptr : s.hyb_part(1, slot), + own_row, (int) n, zc_bf16, + 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_pub(own_row, peer_row, s.hyb_part(2, slot), (int) n, zc_bf16, + s.flag_dev + s.hybf_off(slot), (uint32_t) seq, ctr_fin, stm); + return true; + } + fast_reduce_launch_hyb_combine(data, s.hyb_part(0, slot), one_export ? nullptr : s.hyb_part(1, slot), own_row, (int) n, zc_bf16, stm); + if (hipStreamWriteValue32(stm, s.flag_dev + s.pub_off(slot), (uint32_t) seq, 0) != hipSuccess) return false; + 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(2, slot), (int) n, zc_bf16, stm); + return hipStreamWriteValue32(stm, s.flag_dev + s.hybf_off(slot), (uint32_t) seq, 0) == hipSuccess; +} + +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; + static const bool import_dma = std::getenv("LUCE_CLUSTER_HYBRID_IMPORT_KERNEL") == nullptr; + if (import_dma) { + if (hipMemcpyAsync(data, s.hyb_part(2, slot), n * sizeof(float), hipMemcpyHostToDevice, stm) != hipSuccess) return false; + } else { + fast_reduce_launch_pack(s.hyb_part(2, slot), data, (int) n, 0, stm); + } + if (fast_reduce_lean()) return true; + return hipStreamWriteValue32(stm, s.flag_dev + s.done_off(slot), (uint32_t) seq, 0) == 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; } + + if (s.peer_flags) { hipHostFree(s.peer_flags); s.peer_flags = nullptr; } + if (s.scratch) { hipFree(s.scratch); s.scratch = nullptr; } + if (s.data_host) { hipHostFree(s.data_host); s.data_host = nullptr; } + if (s.flag_host) { hipHostFree(s.flag_host); s.flag_host = nullptr; } + if (s.timed_out_host) { hipHostFree(s.timed_out_host); s.timed_out_host = nullptr; } + if (s.progress_host) { hipHostFree(s.progress_host); s.progress_host = nullptr; } + if (s.hyb_host) { 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 *, int, 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; } +uint64_t FastReduce::submitted() const { return 0; } +uint64_t FastReduce::timed_out() const { return 0; } + +} // 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 0000000000..637f29f166 --- /dev/null +++ b/server/src/cluster/fast_reduce.h @@ -0,0 +1,116 @@ +// fast_reduce.h - a lean all-reduce for the decode path's tiny payloads. +// +// WHY THIS EXISTS. A qwen4exp decode step on two nodes computes in 29.6 ms of +// GPU work and takes 41.0. The difference is ninety-seven all-reduces of +// 10 KiB each -- 117 us apiece. That is not the wire: ib_write_lat puts the +// same payload at 13.65 us, and the round trip through a pinned flag between +// the GPU and a host thread measures 0.79 us. It is what a general-purpose +// collective costs per call, ninety-seven times, and it is the single reason +// adding ranks to this model buys nothing: every axis worth sharding pays for +// its bytes with reductions at about 11.9 MB each, and every axis this model +// has sits on that line. +// +// 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. Unified memory. On Strix Halo there is no separate VRAM, so +// a pinned host buffer *is* GPU memory: the NIC writes into it and a kernel +// reads it without a copy and without GPUDirect. 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: sum of f32, small sizes, no +// in-place aliasing games, 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; + // Decode sizes only -- 10240 floats is the largest a step reduces -- + // and a ring deep enough that the guard against lapping it almost never + // fires. It has to be deep: the guard is a stream wait, and a stream + // wait that fires stops the GPU running ahead, which is where this + // path's speed comes from. Prefill's much larger reductions fall back + // to RCCL, where they cost nothing that matters. + int max_elems = 16384; // 64 KiB + int slots = 4096; // a step has ~97 + uint32_t timeout_ms = 30000; + }; + + // Brings the fabric up: 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. `exact` keeps this payload in f32 on + // the wire and in the add when LUCE_CLUSTER_FAST_REDUCE_BF16 rounds the + // others, 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, zero-copy two-rank mode): + // a rank whose partial is split between a dGPU (hot experts, shared + // expert) and the iGPU that owns its cold experts. The dGPU exports its + // parts to pinned memory while it waits anyway; the iGPU adds its part to + // them 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, int part, 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); + + // Reductions issued and how long the progress thread waited for the GPU, + // for the telemetry line. + uint64_t submitted() const; + uint64_t timed_out() const; + +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 0000000000..93e81376db --- /dev/null +++ b/server/src/cluster/fast_reduce_kernels.cu @@ -0,0 +1,308 @@ +// fast_reduce_kernels.cu - the device half of the lean decode-path reduction. +// +// Three tiny kernels per reduction, all on the backend's own stream so they sit +// in the graph exactly where the collective used to: +// +// set_flag raise this rank's publish flag once the payload has been +// copied out. The copy itself is a DMA, enqueued ahead of this +// on the same stream, so stream order is what publishes it. +// wait_flags spin until every peer's flag has arrived. +// add_peers add the peers' partials, by then copied into device memory. +// +// WHY THE PAYLOAD IS NOT MOVED BY A KERNEL. It was, and it cost 33 ms per +// reduction against RCCL's 117 us -- 250 times worse than the thing it was +// meant to replace. Pinned host memory is mapped uncached for the GPU, so a +// kernel reading or writing bulk data there crawls; the copy engine does not +// care, because it is not reading through the GPU's caches. So the payload goes +// by hipMemcpyAsync and only the flags -- four bytes, three times per +// reduction -- are touched by a kernel at zero-copy speed. +// +// WHY THE FLAGS ARE PINNED AND NOT MANAGED. 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, so they need +// no fault to be reached -- and the round trip through one measures 0.79 us +// (server/test/cluster_flag_latency.cu), which is what makes this worth doing. +// +// Every spin is bounded and counted. A rank that dies must surface as a visible +// failure, not as a GPU that never returns. + +#include + +#include + +namespace luce::cluster { + +namespace { +constexpr int kBlock = 256; + +// Ordering costs cache, and on this part it costs all of it. A system-scope +// release writes the L2 back and a system-scope acquire invalidates it, so a +// flag raised or spun on with the strong form evicts the weights the next layer +// is about to read. Four of them per reduction, ninety-seven reductions per +// step: measured, that alone was 34 ms per reduction against RCCL's 117 us, +// while the same path with the announcements removed ran the step in 32.5 ms. +// +// So each flag gets the weakest ordering that is still correct: +// +// relaxed system for anything the host reads or writes. The payload is +// ordered by the stream -- the copy is a separate +// operation that completes before the kernel starts -- so +// no release is needed to publish it. +// one acquire after a spin succeeds, not once per iteration. The +// invalidate has to happen; it does not have to happen a +// million times. +// agent scope for slot_done, which only this GPU's own later kernels +// read. Telling the whole system about it buys nothing. +__device__ __forceinline__ uint32_t sys_load_relaxed(const uint32_t * p) { + return __hip_atomic_load(p, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM); +} + +__device__ __forceinline__ void sys_store_relaxed(uint32_t * p, uint32_t v) { + __hip_atomic_store(p, v, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM); +} + +__device__ __forceinline__ void sys_acquire(const uint32_t * p) { + (void) __hip_atomic_load(p, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_SYSTEM); +} + +__device__ __forceinline__ uint32_t dev_load(const uint32_t * p) { + return __hip_atomic_load(p, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); +} + +__device__ __forceinline__ void dev_store(uint32_t * p, uint32_t v) { + __hip_atomic_store(p, v, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); +} +} // namespace + +__global__ void fast_reduce_set_flag_kernel(uint32_t * pub_flag, + uint32_t * slot_done, + uint32_t * progress, + uint32_t seq, + uint32_t slot_span, + uint64_t spin_limit) { + if (threadIdx.x || blockIdx.x) return; + // Do not overwrite a slot the progress thread has not finished with. The + // GPU can be a step's worth of reductions ahead, so the ring is defended + // rather than assumed wide enough. + if (seq > slot_span) { + uint64_t spins = 0; + while (dev_load(slot_done) + slot_span < seq) { + if (++spins > spin_limit) break; + } + } + sys_store_relaxed(progress, seq * 4u + 1u); + sys_store_relaxed(pub_flag, seq); +} + +__global__ void fast_reduce_wait_flags_kernel(uint32_t * const * __restrict__ peer_flags, + uint32_t * timed_out, + uint32_t * progress, + uint32_t seq, + int n_peers, + uint64_t spin_limit) { + if (threadIdx.x || blockIdx.x) return; + sys_store_relaxed(progress, seq * 4u + 2u); + for (int p = 0; p < n_peers; ++p) { + uint64_t spins = 0; + while (sys_load_relaxed(peer_flags[p]) != seq) { + if (++spins > spin_limit) { + sys_store_relaxed(timed_out, seq); + return; + } + } + } + // One invalidate, now that every peer has arrived, so the copies queued + // behind this read what the NIC wrote and not a cached line. + sys_acquire(peer_flags[0]); +} + +// The peers' partials are in device memory by now -- the copy engine moved them +// there after the wait -- so this is an ordinary cached add. +__global__ void fast_reduce_add_kernel(float * __restrict__ dst, + const float * __restrict__ scratch, + int n_peers, + int n, + int stride, + uint32_t * slot_done, + uint32_t * progress, + uint32_t seq) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + for (int k = i; k < n; k += blockDim.x * gridDim.x) { + float acc = dst[k]; + for (int p = 0; p < n_peers; ++p) { + acc += scratch[(size_t) p * (size_t) stride + (size_t) k]; + } + dst[k] = acc; + } + // The slot is free once its data has been consumed, not once it arrived. + // Saying so from one thread is safe only because the whole grid's writes + // are complete at the kernel boundary -- which is why the flag is raised + // here and not by a fence inside the loop. + // + // A null slot_done means the caller wants the arithmetic without the + // announcement: a system-scope release writes the L2 back, and the probe + // needs to price that separately from the copies. + if (i == 0 && slot_done) { + dev_store(slot_done, seq); + sys_store_relaxed(progress, seq * 4u + 3u); + } +} + +// ── Zero-copy path for discrete GPUs (LUCE_CLUSTER_FAST_REDUCE_ZC=1) ────── +// The partial goes straight into the pinned row the NIC sends from, and the +// peer's row is added straight out of pinned memory: no copy-engine hops. +// bf16 payloads 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); +} + +// LUCE_CLUSTER_FAST_REDUCE_KPUB: 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. Measured against a hipStreamWriteValue32 after the kernel: +// publish + host ack 9-11 us instead of 16-18 on gfx1201, 3.4 instead of 6.3 on +// gfx1151. A null flag publishes nothing. 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) { + if (!flag) return; + __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_zc_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) + shared in f32, the order +// the dGPU's graph adds them in, then the wire row; the final sum is +// bf16(own) + peer exactly as fast_reduce_add_zc_kernel forms it. +__global__ void fast_reduce_hyb_combine_kernel(const float * __restrict__ cold, const float * __restrict__ hot, + const float * __restrict__ shared, void * own_row, 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) { + const float own = shared ? (cold[k] + hot[k]) + shared[k] : cold[k] + hot[k]; + if (bf16) ((uint16_t *) own_row)[k] = fr_f32_to_bf16(own); + else ((float *) own_row)[k] = own; + } + 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, 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) dst[k] = fr_bf16_to_f32(((const uint16_t *) own_row)[k]) + fr_bf16_to_f32(((const uint16_t *) peer_row)[k]); + else dst[k] = ((const float *) own_row)[k] + ((const float *) peer_row)[k]; + } + fr_publish_last_block(flag, value, counter); +} + +void fast_reduce_launch_hyb_combine_pub(const float * cold, const float * hot, const float * shared, void * own_row, + int n, int bf16, uint32_t * flag, uint32_t value, uint32_t * counter, + hipStream_t stream) { + int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; + hipLaunchKernelGGL(fast_reduce_hyb_combine_kernel, dim3(grid), dim3(kBlock), 0, stream, cold, hot, shared, own_row, n, bf16, + flag, value, counter); +} + +void fast_reduce_launch_hyb_combine(const float * cold, const float * hot, const float * shared, void * own_row, + int n, int bf16, hipStream_t stream) { + fast_reduce_launch_hyb_combine_pub(cold, hot, shared, own_row, n, bf16, nullptr, 0, nullptr, stream); +} + +void fast_reduce_launch_hyb_final_pub(const void * own_row, const void * peer_row, float * dst, int n, int bf16, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream) { + int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; + hipLaunchKernelGGL(fast_reduce_hyb_final_kernel, dim3(grid), dim3(kBlock), 0, stream, own_row, peer_row, dst, n, bf16, + flag, value, counter); +} + +void fast_reduce_launch_hyb_final(const void * own_row, const void * peer_row, float * dst, int n, int bf16, + hipStream_t stream) { + fast_reduce_launch_hyb_final_pub(own_row, peer_row, dst, n, bf16, nullptr, 0, nullptr, stream); +} + +// A word another device polls, written after every earlier write of this +// stream's kernels left this device (peer writes, then the flag on the same +// path, so the reader sees the data once it sees the flag). +__global__ void fast_reduce_flag_sys_kernel(uint32_t * flag, uint32_t value) { + __threadfence_system(); + __hip_atomic_store(flag, value, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM); +} + +void fast_reduce_launch_flag_sys(uint32_t * flag, uint32_t value, hipStream_t stream) { + hipLaunchKernelGGL(fast_reduce_flag_sys_kernel, dim3(1), dim3(1), 0, stream, flag, value); +} + +void fast_reduce_launch_pack_pub(const float * src, void * dst, int n, int bf16, + uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream) { + int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; + hipLaunchKernelGGL(fast_reduce_pack_kernel, dim3(grid), dim3(kBlock), 0, stream, src, dst, n, bf16, + flag, value, counter); +} + +void fast_reduce_launch_pack(const float * src, void * dst, int n, int bf16, hipStream_t stream) { + fast_reduce_launch_pack_pub(src, dst, n, bf16, nullptr, 0, nullptr, stream); +} + +void fast_reduce_launch_add_zc(float * dst, const void * peer, int n, int bf16, hipStream_t stream) { + int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; + hipLaunchKernelGGL(fast_reduce_add_zc_kernel, dim3(grid), dim3(kBlock), 0, stream, dst, peer, n, bf16); +} + +void fast_reduce_launch_set_flag(uint32_t * pub_flag, uint32_t * slot_done, + uint32_t * progress, uint32_t seq, + uint32_t slot_span, uint64_t spin_limit, + hipStream_t stream) { + hipLaunchKernelGGL(fast_reduce_set_flag_kernel, dim3(1), dim3(1), 0, stream, + pub_flag, slot_done, progress, seq, slot_span, spin_limit); +} + +void fast_reduce_launch_wait_flags(uint32_t * const * peer_flags, + uint32_t * timed_out, uint32_t * progress, + uint32_t seq, int n_peers, + uint64_t spin_limit, hipStream_t stream) { + hipLaunchKernelGGL(fast_reduce_wait_flags_kernel, dim3(1), dim3(1), 0, stream, + peer_flags, timed_out, progress, seq, n_peers, spin_limit); +} + +void fast_reduce_launch_add(float * dst, const float * scratch, int n_peers, + int n, int stride, uint32_t * slot_done, + uint32_t * progress, uint32_t seq, + hipStream_t stream) { + int grid = (n + kBlock - 1) / kBlock; + if (grid < 1) grid = 1; + if (grid > 256) grid = 256; + hipLaunchKernelGGL(fast_reduce_add_kernel, dim3(grid), dim3(kBlock), 0, stream, + dst, scratch, n_peers, n, stride, slot_done, progress, seq); +} + +} // namespace luce::cluster diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index 21cd03589b..96e28a9741 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 { @@ -83,6 +84,8 @@ struct BackendArgs { // Expert ownership and routing files (see docs/DS41.md); empty = unused. std::string 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 33220a6165..b56cdf86e6 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -407,6 +407,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 8473515a1c..aa73a1c81c 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -90,6 +90,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 7e0df088d0..830543c984 100644 --- a/server/src/common/backend_plan.cpp +++ b/server/src/common/backend_plan.cpp @@ -176,6 +176,7 @@ BackendPreparation BackendPlanBuilder::resolve( plan.execution_.expert_placement = args.ds4_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 0000000000..2ac15389ee --- /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/cluster_view.h b/server/src/common/cluster_view.h new file mode 100644 index 0000000000..4cd54cac38 --- /dev/null +++ b/server/src/common/cluster_view.h @@ -0,0 +1,70 @@ +// cluster_view.h - dependency-free snapshots of cluster state for the HTTP +// layer. +// +// server/src/server/ has no business knowing about RCCL, the control channel +// or the DeepSeek4 runtime, and it must keep compiling with +// -DLUCE_CLUSTER=OFF. So the cluster backend copies what the API should +// expose into these plain structs, exactly the way get_routing_stats() hands +// out routing telemetry, and ModelBackend answers "no cluster" by default. +// +// Threading: cluster_request_telemetry() is filled by the rank-0 decorator at +// the end of a request and must be read on the same thread that ran it (the +// request worker), which is where GenTimings is assembled. cluster_props() is +// constant after bootstrap and is safe to read from any thread. + +#pragma once + +#include +#include +#include + +namespace luce::common { + +// One rank's contribution to the request that just finished. +struct ClusterRankTiming { + int rank = 0; + uint32_t steps = 0; + uint64_t compute_us = 0; // local prefill+decode wall time + uint64_t allreduce_calls = 0; + uint64_t allreduce_bytes = 0; + uint64_t allreduce_wait_us = 0; // host time blocked on collectives; 0 with + // the in-graph all-reduce + uint64_t ctrl_wait_us = 0; // host time blocked on the control channel + uint64_t peak_device_bytes = 0; // high-water mark, sampled at request ends +}; + +// Per-request cluster telemetry. `ranks[0]` is the head. +struct ClusterTelemetryView { + bool active = false; // false: not a cluster run + int size = 0; + uint64_t request_id = 0; + uint32_t steps = 0; + bool complete = false; // every worker reported + uint64_t head_ctrl_wait_us = 0; + uint64_t hash_probes = 0; + uint64_t hash_mismatches = 0; + int first_mismatch_rank = -1; + uint32_t first_mismatch_step = 0; + std::string error; // non-empty when the gather failed + std::vector ranks; +}; + +// Constant after bootstrap: what this cluster is, for /props. +struct ClusterPropsView { + bool active = false; + int size = 0; + int rank = 0; + std::string ifname; + std::string ib_hca; + int gid_index = 0; + std::string placement; // uniform | balanced | + uint64_t placement_hash = 0; + int replicate_hot = 0; + std::string shared_expert; // replicate | shard | rank0 + std::string allreduce_dtype; // f32 | bf16 | auto + bool ingraph_allreduce = false; // collectives run inside the graph + uint64_t resident_expert_bytes = 0; // this rank's expert shard + uint32_t timeout_ms = 0; +}; + +} // namespace luce::common diff --git a/server/src/common/host_stream_flag.cpp b/server/src/common/host_stream_flag.cpp new file mode 100644 index 0000000000..5aaf1dc7d5 --- /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 0000000000..398a606d87 --- /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/pinned_stage.h b/server/src/common/pinned_stage.h new file mode 100644 index 0000000000..b36df9b236 --- /dev/null +++ b/server/src/common/pinned_stage.h @@ -0,0 +1,151 @@ +// 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 +#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; } + + // Batched mode: set() collects its device uploads instead of queueing + // one copy each, and the caller sends batch() as a single launch (one + // kernel reading the pinned slices; a copy engine adds a start latency + // of 0.03-1 ms to every step on a discrete GPU). end_batch() before the + // next reset; the slices stay valid until then. + struct Pending { + const void * src; + void * dst; + size_t bytes; + }; + void begin_batch() { batching_ = true; pending_.clear(); } + const std::vector & batch() const { return pending_; } + void end_batch() { batching_ = false; pending_.clear(); } + + // 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); + if (batching_) { + pending_.push_back({slice, static_cast(t->data) + offset, bytes}); + return; + } + 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; + bool batching_ = false; + std::vector pending_; +}; + +} // namespace luce::common diff --git a/server/src/internal.h b/server/src/internal.h index dfd52d8759..519d6ade2c 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/server_main.cpp b/server/src/server/server_main.cpp index ed0847be3e..88709908c8 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" @@ -174,6 +177,18 @@ static void print_usage(const char * prog) { " Experts that keep a token's native routing and stay on\n" " the primary GPU (JSON {\"layer\": [ids]})\n" " --ds4-prefill DeepSeek4 prefill: exact, dense, or sparse\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: replicate)\n" + " --cluster-allreduce-dtype f32|bf16|auto (default: auto = f32)\n" + " --cluster-timeout-ms Collective/control watchdog (default: 30000)\n" + " --cluster-verify-hash Debug: cross-rank hidden-state hash every n steps\n" " (default: exact; dense/sparse are experimental\n" " and may change generated tokens)\n" " --fa-window Flash-attention sliding window (default: 0=full).\n" @@ -577,6 +592,72 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, : std::strcmp(flag, "--ds4-router-bias") == 0 ? bargs.ds4_router_bias : bargs.ds4_protected_experts; dst = path; + } else if (std::strncmp(argv[i], "--cluster-", 10) == 0) { + // Expert-parallel cluster flags (maikzz32/lucebox-halo-cluster). + 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-verify-hash") == 0) { + if (!int_value(bargs.cluster.verify_hash_every) || bargs.cluster.verify_hash_every < 0) { + std::fprintf(stderr, "[server] --cluster-verify-hash must be non-negative\n"); + return 2; + } + } 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; @@ -1195,6 +1276,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; @@ -1473,6 +1590,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(); // Cross-check the capability table against the backend that was actually // built. arch_supports_remote_draft() admitted this launch from the arch diff --git a/server/test/unit/test_cluster_unit.cpp b/server/test/unit/test_cluster_unit.cpp new file mode 100644 index 0000000000..68ea49eaac --- /dev/null +++ b/server/test/unit/test_cluster_unit.cpp @@ -0,0 +1,1220 @@ +// 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)); + + const std::vector toks = {1, 2, 3}; + CHECK_EQUAL(hash_tokens(toks), fnv1a64(toks.data(), toks.size() * sizeof(int32_t))); + CHECK_EQUAL(hash_tokens({}), fnv1a64(nullptr, 0)); + const float fl[2] = {1.0f, -2.0f}; + CHECK_EQUAL(hash_floats(fl, 2), fnv1a64(fl, sizeof(fl))); + CHECK_EQUAL(hash_floats(nullptr, 0), fnv1a64(nullptr, 0)); +} + +TEST_CASE(ClusterUnitFixture, MsgTypeNamesAreDistinct) { + std::set names; + for (uint16_t t = 1; t <= 13; t++) { + const char * n = msg_type_name((MsgType) t); + REQUIRE_NOT_NULL(n); + names.insert(n); + } + CHECK_EQUAL(names.size(), (size_t) 13); +} + +// ─── 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.verify_hash_every = 8; + 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); + CHECK_EQUAL(wo.verify_hash_every, 8); + 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); + // 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_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, HashProbeReportBackendOpRequestEndRoundTrip) { + HashProbeMsg hp; + hp.request_id = 1; + hp.step = 8; + hp.rank = 2; + hp.hc_state_hash = 0xAAAAAAAAAAAAAAAAull; + hp.token_hash = 0x5555555555555555ull; + hp.first_divergent_layer = 17; + HashProbeMsg hp_out; + check_roundtrip_and_truncation(hp, hp_out); + CHECK_EQUAL(hp_out.rank, 2); + CHECK_EQUAL(hp_out.hc_state_hash, hp.hc_state_hash); + CHECK_EQUAL(hp_out.token_hash, hp.token_hash); + CHECK_EQUAL(hp_out.first_divergent_layer, 17); + + 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); + REQUIRE(parse_shared_expert_mode("rank0", se, &err)); CHECK(se == SharedExpertMode::Rank0); + 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()); + CHECK_EQUAL(h.verify_hash_every(), 0); + 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.hash_probe(1, 0, 1, 2, &err)); + 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.verify_hash_every = 4; + 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.verify_hash_every, 4); + 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)); + + HeadHooksConfig hc; + hc.timeout_ms = 5000; + hc.verify_hash_every = 4; + HeadHooks head(lp.head, hc); + WorkerHooksConfig wc; + wc.rank = 1; + wc.timeout_ms = 5000; + wc.verify_hash_every = 4; + WorkerHooks worker(lp.worker, wc); + CHECK(head.is_head()); CHECK(head.is_cluster()); + CHECK_FALSE(worker.is_head()); CHECK(worker.is_cluster()); + CHECK_EQUAL(worker.verify_hash_every(), 4); + + 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, hash probe 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, w_hp = false, w_hp2 = 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); + w_hp = worker.hash_probe(req, 4, 0xAAAA, 0xBBBB, &w_err); // matches + w_hp2 = worker.hash_probe(req, 8, 0xAAAA, 0xCCCC, &w_err); // token hash differs + }); + + 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)); + REQUIRE(head.hash_probe(req, 4, 0xAAAA, 0xBBBB, &h_err)); + CHECK_EQUAL(head.counters().hash_mismatches, 0u); + CHECK_EQUAL(head.first_mismatch_rank(), -1); + REQUIRE(head.hash_probe(req, 8, 0xAAAA, 0xBBBB, &h_err)); // non-strict: logged, not fatal + CHECK_EQUAL(head.counters().hash_mismatches, 1u); + CHECK_EQUAL(head.first_mismatch_rank(), 1); + CHECK_EQUAL(head.first_mismatch_step(), 8u); + 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(w_hp); CHECK(w_hp2); + CHECK_EQUAL(head.counters().decisions, 1u); + CHECK_EQUAL(head.counters().drafts, 1u); + CHECK_EQUAL(head.counters().accepts, 1u); + CHECK_EQUAL(head.counters().hash_probes, 2u); + CHECK_EQUAL(worker.counters().decisions, 1u); + CHECK_EQUAL(worker.counters().hash_probes, 2u); + + // 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); + + // Head samples itself when the caller passes token < 0 (greedy). + int32_t sampled = -1; uint8_t sflags = 0; + std::thread wt2([&] { + int32_t t = -1; uint8_t fl = 0; std::string e; + worker.decide_next_token(req, 1, nullptr, 0, scfg, false, t, fl, &e); + }); + REQUIRE(head.decide_next_token(req, 1, logits, 4, scfg, false, sampled, sflags, &h_err)); + CHECK_EQUAL(sampled, 1); + wt2.join(); + } + + 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)); + + HeadHooksConfig hc; hc.timeout_ms = 5000; + HeadHooks head(lp.head, hc); + WorkerHooksConfig wc; wc.rank = 1; 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, HeadHooksFailWhenWorkerGone) { + LoopbackPair lp; + if (lp.port == 0) SKIP("no loopback port available"); + REQUIRE(lp.handshake(lp.worker_id)); + lp.worker.close(); + + HeadHooksConfig hc; + hc.timeout_ms = 500; + HeadHooks head(lp.head, hc); + std::string err; + // A gather with no live worker must return false within the deadline. + CHECK_FALSE(head.hash_probe(1, 0, 1, 2, &err)); + CHECK_FALSE(err.empty()); + lp.head.close(); +} + +#endif // !_WIN32 From 9f350f023c548edd6adb5b4aab129409b23eb842 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Thu, 8 Oct 2026 13:08:40 +0200 Subject: [PATCH 05/40] feat(ds41): expert-parallel DeepSeek V4.1 across two boxes Each rank keeps the single-box hybrid tiers (hot experts on the discrete GPU, the rest on the integrated GPU) for its share of the routed experts, and the two ranks reduce every MoE layer over the cluster exchange: - hybrid FFN exchange through the integrated GPU (it sums the cold and imported rows and sends one export), a lean exchange, and one export per layer; - attention heads split across the ranks, and a rank that loads only its attention slice; - greedy verify head and DSpark draft head with their vocabulary split across the ranks (exact argmax over both halves), and the DSpark drafter split across both ranks; - workers that skip the drafter they never use, a drafter swap for long prompts with chunks agreed across ranks, and a pipelined prefill only when every rank pipelines; - prefill that keeps the device hot and cold join and reduces in the HC post graph, with a bf16 wire for prefill-sized reductions; - the drafter's context window and the verify's window rows kept on the device, batched verify uploads, one wait per step for small transfers, and the Engram projection run inside the FFN waits. Cluster mode (--cluster-size > 1) installs the measured configuration as environment defaults (cluster_launch_env in launch_profiles.h) before a launch profile's, so a two-box run needs only the operator's inputs: expert budget, placement and router bias. Each default is a switch that 0 turns off, and an explicitly set variable always wins. A cluster rank never retries a failed prefill or a failed drafter swap-in on its own: the ranks run every chunk together. Co-Authored-By: Claude Opus 5.5 --- server/src/common/dspark_head.cpp | 254 ++- server/src/common/dspark_head.h | 37 + server/src/deepseek4/deepseek4_backend.cpp | 398 ++++- server/src/deepseek4/deepseek4_backend.h | 29 +- server/src/deepseek4/deepseek4_cluster.cpp | 904 ++++++++++ server/src/deepseek4/deepseek4_cluster.h | 287 +++ server/src/deepseek4/deepseek4_dspark.cpp | 54 +- server/src/deepseek4/deepseek4_dspark.h | 29 +- .../src/deepseek4/deepseek4_dspark_spec.cpp | 417 ++++- server/src/deepseek4/deepseek4_engram.h | 3 +- .../src/deepseek4/deepseek4_engram_apply.cpp | 15 +- .../src/deepseek4/deepseek4_fused_verify.inc | 746 +++++++- server/src/deepseek4/deepseek4_graph.cpp | 1534 ++++++++++++++--- server/src/deepseek4/deepseek4_internal.h | 29 + server/src/deepseek4/deepseek4_loader.cpp | 132 ++ server/src/server/launch_profiles.h | 56 + server/src/server/server_main.cpp | 35 +- server/test/unit/test_deepseek4_unit.cpp | 180 ++ server/test/unit/test_launch_policy.cpp | 14 + 19 files changed, 4756 insertions(+), 397 deletions(-) create mode 100644 server/src/deepseek4/deepseek4_cluster.cpp create mode 100644 server/src/deepseek4/deepseek4_cluster.h diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index 6bb68b164c..8a1571e317 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 c407800dea..a2aa173763 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/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 0e1a880b2c..00fa023c37 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3,6 +3,8 @@ #include "deepseek4_backend.h" #include "deepseek4_budget_hook.h" +#include "deepseek4_cluster.h" +#include "cluster/cluster_decision_hooks.h" #include "deepseek4_internal.h" #include "deepseek4_image_spans.h" #include "deepseek4_image_budget.h" @@ -1136,6 +1138,173 @@ void log_deepseek4_step_telemetry(const char * phase, DeepSeek4Backend::DeepSeek4Backend(DeepSeek4BackendConfig cfg) : cfg_(std::move(cfg)) {} +// ─── Expert-parallel cluster (after maikzz32/lucebox-halo-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) { + const char * fe = std::getenv("LUCE_CLUSTER_FAST_REDUCE"); + if (fe && std::atoi(fe) != 0) { + cluster::FastReduce::Config fc; + 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)); + auto fast = std::make_unique(); + std::string ferr; + if (fast->init(fc, &ferr)) { + 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 { + std::fprintf(stderr, "[deepseek4-cluster] rank %d fast reduce unavailable, keeping " + "RCCL: %s\n", cfg->rank, ferr.c_str()); + } + } + } + if (comm && !cluster_agree_prefill_chunk()) 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; +} + +// 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); + } + 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(); } @@ -1968,8 +2137,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(); @@ -1979,6 +2176,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) { @@ -2258,8 +2474,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_) { @@ -2278,6 +2510,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. @@ -2545,6 +2785,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"); @@ -2793,7 +3034,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; } @@ -3332,6 +3583,17 @@ bool DeepSeek4Backend::init_hybrid_model() { TargetLoadPlan plan; plan.skip_expert_tensors = true; plan.load_ds4_image_bias = !cfg_.mmproj_path.empty(); + // LUCE_CLUSTER_SLICE_WEIGHTS=1: a cluster rank loads only the dense + // slices its head split and shared-expert shard read (about 3.3 GB less + // on the primary GPU for DS4.1 on two ranks, for more hot experts). + if (cluster_ && cluster_->cfg && cluster_->cfg->size > 1 && + env_flag_enabled("LUCE_CLUSTER_SLICE_WEIGHTS")) { + 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()); @@ -3477,6 +3739,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(); @@ -4511,9 +4788,15 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, // finish_reason still reported "stop". This mirrors qwen35_backend's override-and-continue. bool budget_close_started = false; size_t close_inject_pos = 0; + // Cluster lockstep (maikzz32/lucebox-halo-cluster): 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; @@ -4569,8 +4852,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 { @@ -4588,13 +4888,16 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, // keep going. Runs before history.push_back so penalty history records what was // actually emitted. The rule lives in a header-only helper so it is testable without a // model; see deepseek4_budget_hook.h for why this overrides rather than appends. - { - bool hook_forced = false; + // Head/local only: the worker's token already carries the override. + bool hook_forced = false; + if (!worker) { next_token = luce::deepseek4::budget_hook_apply( budget_hook.close_token_ids, n_gen - generated, budget_hook.hard_limit_remaining, next_token, budget_close_started, close_inject_pos, hook_forced); if (hook_forced && forced_close_out) *forced_close_out = true; + } else if ((decision_flags & cluster::kDecisionBudget) && forced_close_out) { + *forced_close_out = true; } if (process_logits) { @@ -4612,7 +4915,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, hook_forced, 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; } } @@ -4631,6 +4960,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; @@ -4736,9 +5067,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; @@ -4757,11 +5089,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"); } } @@ -4844,6 +5181,34 @@ 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; } } + 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=*/false, + /*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 (env_flag_enabled("LUCE_DS4_TIMING")) { size_t nonfinite_logits = 0; for (float value : last_logits_) { @@ -4861,7 +5226,12 @@ GenerateResult DeepSeek4Backend::generate_from_state( float accept_rate = 0.0f; bool spec_ran = false; bool spec_forced_close = 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; @@ -4891,7 +5261,8 @@ GenerateResult DeepSeek4Backend::generate_from_state( expert_runtime_.compute ? &expert_runtime_ : nullptr, routing_stats_.get(), spec_sampler ? &*spec_sampler : nullptr, - spec_hook ? &*spec_hook : nullptr)) { + spec_hook ? &*spec_hook : nullptr, + cluster ? hooks_ : nullptr)) { result.fail(GenerateErrorCode::DecodeFailed, "DSpark speculative decode failed"); return result; @@ -5399,6 +5770,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 f135350f39..f03966d3e3 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,32 @@ class DeepSeek4Backend : public ModelBackend { return routing_stats_.get(); } + // ── Expert-parallel cluster (after maikzz32/lucebox-halo-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; + // Loads, on load_model()'s path, the size every rank must agree on. + bool cluster_agree_prefill_chunk(); + 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 0000000000..7894bd9e7c --- /dev/null +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -0,0 +1,904 @@ +// 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 { + +// LUCE_CLUSTER_TRACE_COLL=1: one stderr line per collective (rank, sequence, kind, elements, stream), so the +// two ranks' collective sequences can be diffed when a step hangs. +static void trace_coll(int rank, const char * kind, size_t n, const void * stream) { + static const bool on = std::getenv("LUCE_CLUSTER_TRACE_COLL") != nullptr; + static unsigned long long seq = 0; + if (!on) return; + std::fprintf(stderr, "[coll] r%d #%llu %s n=%zu stream=%p\n", rank, seq++, kind, n, stream); +} + +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 need decode tables registered per resident tensor +// (register_deepseek4_moe_hybrid_mix_tables), which today requires a +// materialized GPU cold owner. Refuse instead of silently loading garbage. +// Supporting them needs register_deepseek4_moe_hybrid_mix_tables to accept +// hot-only storage (cold owner None). +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); + trace_coll(rt.rank(), "host-allreduce", n, (const void *) stream); + + 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; + } + if (cfg.shared_expert == cluster::SharedExpertMode::Rank0) { + set_err(err, std::string("shared expert mode '") + + cluster::shared_expert_mode_name(cfg.shared_expert) + + "' is not implemented yet; use replicate or shard"); + 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; +} + +// ─── Diagnostics ──────────────────────────────────────────────────────── + +void ds4_cluster_checksum(const float * data, size_t n, double * sum, double * sum_abs) { + double s = 0.0; + double a = 0.0; + if (data) { + for (size_t i = 0; i < n; ++i) { + s += (double) data[i]; + a += (double) (data[i] < 0.0f ? -data[i] : data[i]); + } + } + if (sum) *sum = s; + if (sum_abs) *sum_abs = a; +} + +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() { + static const bool enabled = [] { + const char * v = std::getenv("LUCE_CLUSTER_PREFILL_PIPELINE"); + return ds4_cluster_prefill_device_join_enabled() && v && *v && std::strcmp(v, "0") != 0; + }(); + return enabled; +} + +bool ds4_cluster_prefill_attention_split(int context_end) { + static const bool enabled = [] { + const char * v = std::getenv("LUCE_CLUSTER_PREFILL_ATTN_SPLIT"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + static const int min_rows = [] { + const char * v = std::getenv("LUCE_CLUSTER_PREFILL_ATTN_SPLIT_MIN"); + return v && *v ? std::max(0, std::atoi(v)) : 0; + }(); + return enabled && context_end >= min_rows; +} + +bool ds4_cluster_prefill_device_join_enabled() { + static const bool enabled = [] { + const char * v = std::getenv("LUCE_CLUSTER_PREFILL_DEVICE_JOIN"); + return !v || !*v || std::strcmp(v, "0") != 0; + }(); + return enabled; +} + +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; + if (rt->cfg && rt->cfg->shared_expert == cluster::SharedExpertMode::Rank0) { + 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; +} + +bool ds4_cluster_env_prefill_single_token() { + static const bool enabled = [] { + const char * v = std::getenv("LUCE_CLUSTER_PREFILL_SINGLE_TOKEN"); + return v && v[0] && std::strcmp(v, "0") != 0; + }(); + return enabled; +} + +// ─── 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 ───────────────────────────────────────────────────────── + +// Path 3b: 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; + }(); + if (prefill_bf16 && n >= (size_t) 262144 && rt->scratch_bf16 && rt->scratch_elems >= n) { + trace_coll(rt->rank(), "graph-allreduce-bf16", n, stream); + 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)) { + trace_coll(rt->rank(), "fast-allreduce", n, stream); + rt->telemetry.allreduce_calls += 1; + rt->telemetry.allreduce_bytes += (uint64_t) n * sizeof(float); + return; + } + trace_coll(rt->rank(), "graph-allreduce", n, stream); + 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_export0_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, 0, rt->hyb_export0++)) ds4_hyb_fail(rt, "export (hot)"); +} +static void ds4_hyb_export1_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, 1, rt->hyb_export1++)) ds4_hyb_fail(rt, "export (shared)"); +} +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_hot_node, + ggml_tensor ** export_shared_node) { + if (!ctx || !hot || !shared || !cold || !combine_node || !export_hot_node || !export_shared_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. + // LUCE_CLUSTER_HYBRID_ONE_EXPORT=1: export hot + shared as one partial + // (half the dGPU link bytes; the iGPU adds cold + (hot + shared)). + static const bool one_export = [] { + const char * v = std::getenv("LUCE_CLUSTER_HYBRID_ONE_EXPORT"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + ggml_tensor * ex0 = nullptr; + ggml_tensor * ex1 = nullptr; + if (one_export) { + ggml_tensor * hs = ggml_add(ctx, hot, shared); + ex1 = ggml_cluster_allreduce_inplace(ctx, hs, &ds4_hyb_export0_cb, &rt); + ex0 = ex1; + } else { + ex0 = ggml_cluster_allreduce_inplace(ctx, hot, &ds4_hyb_export0_cb, &rt); + ex1 = ggml_cluster_allreduce_inplace(ctx, shared, &ds4_hyb_export1_cb, &rt); + } + ggml_tensor * comb = ggml_cluster_allreduce_inplace(ctx, cold, &ds4_hyb_combine_cb, &rt); + // The import overwrites the exported shared partial with the reduced sum. + ggml_tensor * imp = ggml_cluster_allreduce_inplace(ctx, ex1, &ds4_hyb_import_cb, &rt); + *combine_node = comb; + *export_hot_node = ex0; + *export_shared_node = ex1; + return imp; +} + +ggml_cluster_allreduce_fn ds4_cluster_allreduce_fn() { + return &ds4_cluster_allreduce_graph_callback; +} + +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(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + ggml_tensor * partial, + int layer, + DeepSeek4StepTelemetry * telemetry, + std::string * err) { + if (!partial) { + set_err(err, "cluster all-reduce: null partial tensor"); + return false; + } + if (partial->type != GGML_TYPE_F32 || !ggml_is_contiguous(partial)) { + set_err(err, "cluster all-reduce: partial must be contiguous F32"); + return false; + } + const auto t0 = ClusterClock::now(); + const size_t n = (size_t) ggml_nelements(partial); + const bool ok = allreduce_device(rt, backend, partial->data, n, layer, + /*force_wait=*/false, telemetry, err); + if (telemetry) telemetry->cluster_allreduce_us += us_since(t0); + return ok; +} + +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 { + +void broadcast_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; + if (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)) { + trace_coll(rt->rank(), "fast-broadcast", n, stream); + return; + } + trace_coll(rt->rank(), "graph-broadcast", n, stream); + 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; + } + } +} + +// Rank 0's values bit for bit: f32 on the wire even where the fast reduce +// rounds its payloads to bf16. No upload hook (that belongs to the Engram keys). +void broadcast_exact_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; + if (rt->fast && rt->fast->ok() && + rt->fast->submit_from_root((float *) data, n, stream, /*root=*/0, /*exact=*/true)) { + trace_coll(rt->rank(), "fast-broadcast-exact", n, stream); + return; + } + trace_coll(rt->rank(), "graph-broadcast", n, stream); + 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; + } + } +} + +} // namespace + +ggml_tensor * ds4_cluster_broadcast_exact_node(ggml_context * ctx, + ggml_tensor * value, + Ds4ClusterRuntime & rt) { + 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, &broadcast_exact_graph_callback, &rt); +} + +ggml_tensor * ds4_cluster_broadcast_node(ggml_context * ctx, + ggml_tensor * value, + Ds4ClusterRuntime & rt) { + 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); + // The in-graph collective node runs its callback in place on a copy of + // `value`; this callback is a broadcast from rank 0. + return ggml_cluster_allreduce(ctx, value, &broadcast_graph_callback, &rt); +} + +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); + trace_coll(rt.rank(), "host-broadcast", n, (const void *) stream); + 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 0000000000..bb9d6eb6bc --- /dev/null +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -0,0 +1,287 @@ +// deepseek4_cluster.h - DeepSeek4 expert-parallel cluster runtime. +// +// Every rank loads the full dense model 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 routed partial identical on every rank before the shared expert +// is added and 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: SharedExpertMode::Replicate only. Every rank computes it +// locally from the same normalized input and adds it after the all-reduce; +// the routed partial is 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. + uint64_t hyb_export0 = 0, hyb_export1 = 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; + // Path 3b: 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. + int attn_head_begin = 0; + int attn_head_count = 0; + // Shared-expert sharding (--cluster-shared-expert shard): 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); + +// Diagnostics: sum and sum of |x| over n floats (LUCE_CLUSTER_TRACE lines). +void ds4_cluster_checksum(const float * data, size_t n, double * sum, double * sum_abs); + +// LUCE_CLUSTER_PREFILL_SINGLE_TOKEN=1: run cluster prefill token by token +// (n_tokens == 1 chunks through the decode path) to bisect multi-token +// prefill numerics. Cached after the first call. +bool ds4_cluster_env_prefill_single_token(); + +// All-reduce a device-resident F32 partial [n_embd, n_tokens] in place on the +// backend stream. Stays stream-ordered (no host sync) unless rt.trace is set, +// in which case it waits with cfg->timeout_ms and logs the layer. Updates +// rt.telemetry and, when given, telemetry->cluster_allreduce_*. +bool ds4_cluster_allreduce_layer(Ds4ClusterRuntime & rt, + ggml_backend_t backend, + ggml_tensor * partial, + int layer, + DeepSeek4StepTelemetry * telemetry, + std::string * err); + +// Path 3b: 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); + +// The same sum-all-reduce as a raw callback, for graphs built outside the +// deepseek4 layer. server/src/common/dspark_head.cpp needs it to combine the +// vocabulary-sliced logits of the DSpark head, and must not gain a cluster +// dependency to do so, so it takes an ggml_cluster_allreduce_fn and a void* +// rather than the runtime itself. Pass `&rt` as the user pointer: errors land +// in rt.node_error and byte counts in the usual telemetry. +ggml_cluster_allreduce_fn ds4_cluster_allreduce_fn(); + +// 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_hot_node, + ggml_tensor ** export_shared_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(); +// LUCE_CLUSTER_PREFILL_DEVICE_JOIN (default on): cluster prefill keeps the +// device-resident hot+cold join and reduces in the HC-post graph. +bool ds4_cluster_prefill_device_join_enabled(); +// LUCE_CLUSTER_PREFILL_PIPELINE=1 (default off): band-pipelined prefill on a +// cluster rank, band count agreed across ranks. Off by default: at the 14 GB +// hot budget the agreed bands shrink to 512 rows and prefill gets slower. +bool ds4_cluster_prefill_pipeline_enabled(); + +// Attention head split in the uncached (prefill) graphs: true when a band +// whose context ends at `context_end` attends with this rank's heads and +// all-reduces the partial (opt-in, LUCE_CLUSTER_PREFILL_ATTN_SPLIT=1: measured +// slower at 2K and even at 30K, where the head-independent indexer dominates; LUCE_CLUSTER_PREFILL_ATTN_SPLIT_MIN sets the context +// it starts at, default 0). A pure function of the band, so every rank +// agrees. +bool ds4_cluster_prefill_attention_split(int context_end); + +// 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 ad990cf205..594aa85221 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 60a251560e..f9935956fd 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -40,6 +40,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 @@ -79,13 +83,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`. @@ -322,6 +344,9 @@ bool run_deepseek4_dspark_spec_decode( DSparkSpecSampling * sampling = nullptr, // Non-null for requests with a thinking budget: the hook is applied // inside speculative decode (deepseek4_budget_hook.h). - DSparkBudgetHook * budget_hook = nullptr); + DSparkBudgetHook * budget_hook = nullptr, + // Cluster lockstep hooks (from maikzz32's lucebox-halo-cluster): 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 a48ddc3a55..f11b8752b8 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 "deepseek4_budget_hook.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,8 +334,10 @@ 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 +// 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 verifies q5 for 1.3x the q2 step, not 1.6x) 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 @@ -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. About 100 MB, 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; @@ -938,8 +1061,34 @@ bool run_deepseek4_dspark_spec_decode( MoeExpertComputeRuntime * expert_runtime, MoeHybridRoutingStats * routing_stats, DSparkSpecSampling * sampling, - DSparkBudgetHook * budget_hook) { + DSparkBudgetHook * budget_hook, + luce::cluster::Ds4ClusterHooks * cluster_hooks) { const int n_embd = target_w.n_embd; + // Cluster lockstep (from maikzz32's lucebox-halo-cluster). 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; + // LUCE_CLUSTER_WORKER_KEEPALIVE=1: between its verify and the head's next + // draft a worker's GPU is idle for milliseconds, and its first launch after + // that started ~1 ms late (wake-up), which the head then waited for at the + // first exchange. One sleeping wave keeps it awake until the draft arrives. + const bool worker_keepalive = cluster_spec && !cluster_spec_head && !draft_split && + spec_env_flag("LUCE_CLUSTER_WORKER_KEEPALIVE"); + struct KeepaliveOff { + ggml_backend_t backend; bool on; + ~KeepaliveOff() { if (on) ggml_backend_cuda_keepalive(backend, false, 0); } + } keepalive_off{backend, worker_keepalive}; const bool hook_on = budget_hook && !budget_hook->close_ids.empty(); luce::deepseek4::SpecBudgetHookState hook_st; bool close_emitted = false; // the hook's close token reached the stream @@ -1057,6 +1206,14 @@ bool run_deepseek4_dspark_spec_decode( width_cost_ms[(size_t) width] = kDs4VerifyWidthCostMs[width]; } width_controller.set_relative_costs(width_cost_ms); + // Step costs belong to the machine, not to the text: every request starts + // from the costs the earlier ones measured (their warmup hold re-arms). + static std::mutex measured_width_costs_mu; + static std::vector measured_width_costs; + { + std::lock_guard lock(measured_width_costs_mu); + width_controller.set_measured_costs(measured_width_costs); + } if (timing && width_controller.enabled()) { std::fprintf(stderr, "[ds4-spec] adaptive width policy=%s\n", @@ -1144,6 +1301,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 703 MB Q8_0 lm_head every step (~1.1 ms 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) { const SpecClock::time_point step_t0 = SpecClock::now(); const int ctx_len = feat_count < n_swa ? feat_count : n_swa; @@ -1154,7 +1386,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 +1434,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 +1472,7 @@ bool run_deepseek4_dspark_spec_decode( draft_tok.push_back(budget_hook->close_ids[hook_st.inject_pos + (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 +1480,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: 1.67 -> 1.13 ms for four columns on gfx1201). + 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 +1564,58 @@ 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) { + // A worker's keepalive wave may share a hardware queue with the + // compute stream and would hold the head's graph until the wave's + // cap (~40 ms a step): stop it, queue the graph, then queue the + // wave behind it. + if (worker_keepalive) ggml_backend_cuda_keepalive(backend, false, 0); + 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; + } + if (worker_keepalive) ggml_backend_cuda_keepalive(backend, true, 40); + } + // 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 (worker_keepalive) ggml_backend_cuda_keepalive(backend, false, 0); + 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); @@ -1372,6 +1683,7 @@ bool run_deepseek4_dspark_spec_decode( const bool verify_ok = target.verify_batch(draft_tok, pos, verify_last, &tgt_am); tm_verify += spec_ms_since(t0); + if (worker_keepalive) ggml_backend_cuda_keepalive(backend, true, 40); if (probe_inflight) { const SpecClock::time_point probe_t0 = SpecClock::now(); deepseek4_dspark_draft_wait(drafter_backend); @@ -1478,6 +1790,32 @@ bool run_deepseek4_dspark_spec_decode( // The hook replaced this step's bonus with the close token (and may // have cut the step short there). const bool close_bonus = hook_on && !hook_was_started && hook_st.started; + // 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; + } + } const int matched = accept - 1; // accepted candidates const int commit_pos = pos + accept; // seed + accepted candidates in KV @@ -1518,7 +1856,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]. @@ -1646,7 +1984,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; + } } if (hook_on && close_emitted) budget_hook->fired = true; @@ -1672,6 +2052,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(measured_width_costs_mu); + measured_width_costs = width_controller.measured_costs(); + } + 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 9c7d02b2d5..bee231e30b 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 22ffc04980..aca977700b 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 353f184fbe..c4e056834e 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]; @@ -254,15 +260,25 @@ static std::vector ds4_fused_stream_jobs( // the launch. static enum ggml_status ds4_fused_graph_compute( const DeepSeek4Weights & w, DeepSeek4FusedDecodeGraph & fg, ggml_backend_t backend, - MoeHybridStorage * hybrid, const char * who) { + MoeHybridStorage * hybrid, const char * who, + const std::function & while_running = {}) { MoeStreamedMailbox * mailbox = hybrid && hybrid->expert_cache ? hybrid->expert_cache->mailbox() : nullptr; const std::vector jobs = mailbox ? ds4_fused_stream_jobs(fg) : std::vector{}; if (!jobs.empty()) mailbox->begin(jobs, ds4_fused_invalid_route()); - enum ggml_status status = fg.sched - ? ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf) - : ggml_backend_graph_compute(backend, fg.sg.gf); + enum ggml_status status; + if (fg.sched && jobs.empty() && while_running) { + // Host work that needs no device result (building the next shape) + // runs while the GPUs execute this graph. + status = ggml_backend_sched_graph_compute_async(fg.sched, fg.sg.gf); + if (status == GGML_STATUS_SUCCESS) while_running(); + ggml_backend_sched_synchronize(fg.sched); + } else { + status = fg.sched + ? ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf) + : ggml_backend_graph_compute(backend, fg.sg.gf); + } if (!jobs.empty()) { std::string error; if (!mailbox->end(&error)) { @@ -514,7 +530,7 @@ static bool ds4_build_fused_verify_graph( } step_graph_free(fg.sg); fg.reset_nodes(); - ex.reset(); + { const bool keep_vocab_split = ex.vocab_split; ex.reset(); ex.vocab_split = keep_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 +542,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 (from maikzz32's lucebox-halo-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 +586,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 (ds4_lane_batch_comp_enabled() && !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 +727,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 +779,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 +1077,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 +1086,89 @@ 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 (LUCE_DS4_LANE_BATCH_SELECT=0 + // keeps one chain per lane and concatenates their selections). + static const bool lane_batch_select = [] { + const char * value = std::getenv("LUCE_DS4_LANE_BATCH_SELECT"); + return !value || std::strcmp(value, "0") != 0; + }(); + const bool external_selection = lane_batch && lane_batch_select && + 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 && ds4_lane_batch_comp_enabled() && 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 +1243,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 +1282,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 +1386,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 +1542,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 +1583,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 +1622,86 @@ 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 * export_hot = nullptr; + ggml_tensor * export_shared = nullptr; + ggml_tensor * sum = ds4_cluster_hybrid_exchange_nodes( + ctx, hyb_ex_inputs->main_output, shared, hyb_ex_inputs->peer_raw, + *cluster_rt, &combine, &export_hot, &export_shared); + if (!sum) return false; + ggml_build_forward_expand(gf, export_hot); + ggml_build_forward_expand(gf, export_shared); + cluster_nodes.push_back(export_shared); + ggml_set_output(combine); + ggml_build_forward_expand(gf, combine); + hyb_ex_inputs->cold_nodes.push_back(combine); + cluster_nodes.push_back(export_hot); + // 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 +1742,31 @@ 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); + 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; + 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 (paged_mode || ds4_env_flag("LUCE_DS4_GPU_ARGMAX_VERIFY")) { + } + 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 +1789,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 +1886,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 +1909,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 +2031,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 +2107,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 +2192,18 @@ 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); + static bool vocab_split_reported = false; + if (vocab_split_env && !vocab_split_reported) { + vocab_split_reported = true; + std::fprintf(stderr, "[ds4-fused-verify] vocab split %d (hybrid %d argmax_only %d argmax_out %d cluster %d)\n", + vocab_split ? 1 : 0, hybrid ? 1 : 0, hooks->prefer_argmax_only ? 1 : 0, + hooks->argmax_out ? 1 : 0, cache.cluster_rt ? cache.cluster_rt->size() : 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) { @@ -1824,6 +2223,36 @@ static int ds4_try_fused_verify_step( ((int64_t) padded << 4) | (int64_t) (b_idx + 1)); } } + // The key a step at (kv, qq) would look up: the rules above. + auto key_for = [&](int kv, int qq) { + std::vector k; + k.reserve(key.size()); + k.push_back(w.n_swa); + k.push_back(qq); + k.push_back(token_ids ? 1 : 0); + k.push_back(hybrid ? 1 : 0); + k.push_back((int64_t) hooks->capture_layer_ids->size()); + for (int layer : *hooks->capture_layer_ids) k.push_back(layer); + k.push_back(vocab_split ? 1 : 0); + const int nl = ds4_fused_attention_lanes(w, qq); + const int alq = qq / nl; + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + DeepSeek4LayerCache & lc = ds4_comp_cache(cache, w, il); + for (int lane = 0; lane < nl; ++lane) { + const int lks = kv + lane * alq; + const int padded = ds4_fused_lane_padded(w, lc, il, lks + alq - 1); + int b_idx = -1; + if (ratio > 0) { + for (int t = 0; t < alq; ++t) { + if (((lks + t + 1) % ratio) == 0) { b_idx = t; break; } + } + } + k.push_back(((int64_t) padded << 4) | (int64_t) (b_idx + 1)); + } + } + return k; + }; vc.counter++; const size_t slot_limit = hybrid @@ -1863,6 +2292,7 @@ static int ds4_try_fused_verify_step( // that state is live can replay kernels against stale allocations. fg->destroy(vc.backend, vc.peer_backend); ex->reset(); + ex->vocab_split = vocab_split; const auto build_t0 = Ds4TimingClock::now(); if (!ds4_build_fused_verify_graph(mc, *fg, *ex, backend, w, &cache, @@ -1880,14 +2310,60 @@ 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 = ds4_pinned_step_io_enabled() && + 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). + Ds4ClusterRuntime * engram_rt = 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 && ds4_engram_async_enabled() && + 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(); + // LUCE_DS4_STAGE_BATCH=1: the step's pinned uploads leave as one copy + // kernel before the single wait below, instead of a dozen copies of + // which the larger ones wait for a copy engine to start. + static const bool stage_batch_env = [] { + const char * v = std::getenv("LUCE_DS4_STAGE_BATCH"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + const bool stage_batch = stage_batch_env && pinned_io && ggml_backend_is_cuda(backend); + if (stage_batch) vc.stage.begin_batch(); ds4_fv_set(fg->inp_embed, embed, sizeof(float) * (size_t) w.n_embd * q); std::vector iv(q); std::vector lv(q); @@ -1899,6 +2375,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); @@ -1945,6 +2425,12 @@ static int ds4_try_fused_verify_step( } ds4_fv_set(fg->i64_bundle, i64v.data(), sizeof(int64_t) * i64v.size()); + static const bool set_timing = std::getenv("LUCE_DS4_SET_TIMING") != nullptr; + static double st_inputs = 0, st_mask = 0, st_hash = 0, st_sync = 0; static int st_n = 0; + auto st_now = [] { return std::chrono::steady_clock::now(); }; + auto st_us = [](auto a, auto b) { return std::chrono::duration(b - a).count(); }; + const auto st_t1 = st_now(); + st_inputs += st_us(set_t0, st_t1); // Causal mask values. At long context this tensor is multi-megabyte, // while only the ring slots overwritten by this q-wide batch and the // padded compressed tail change between verifier steps. Build those @@ -2048,14 +2534,22 @@ 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)); + if (pinned_io) { + vc.stage.set(backend, fg->mask_bundle, maskv.data() + range.first, + range.first * sizeof(float), + (range.second - range.first) * sizeof(float)); + } else { + ggml_backend_tensor_set( + fg->mask_bundle, maskv.data() + range.first, + range.first * sizeof(float), + (range.second - range.first) * sizeof(float)); + } } } } + const auto st_t2 = st_now(); + st_mask += st_us(st_t1, st_t2); if (token_ids) { for (int il = 0; il < w.n_layer; ++il) { ggml_tensor * hids = fg->hash_ids[(size_t) il]; @@ -2076,10 +2570,49 @@ 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. + const auto st_t3 = st_now(); + st_hash += st_us(st_t2, st_t3); + if (stage_batch) { + const auto & pending = vc.stage.batch(); + if (!pending.empty()) { + std::vector descs; + descs.reserve(pending.size()); + for (const auto & p : pending) descs.push_back({p.src, p.dst, p.bytes}); + ggml_backend_cuda_copy_batch_async(backend, descs.data(), (int) descs.size()); + } + vc.stage.end_batch(); + } + if (pinned_io) ggml_backend_synchronize(backend); + st_sync += st_us(st_t3, st_now()); + if (set_timing && ++st_n % 64 == 0) { + std::fprintf(stderr, "[ds4-set-t] per verify over %d: inputs %.0f mask %.0f hash %.0f sync %.0f us (pinned %d)\n", + st_n, st_inputs / st_n, st_mask / st_n, st_hash / st_n, st_sync / st_n, pinned_io ? 1 : 0); + } + // 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()); } @@ -2114,8 +2647,65 @@ static int ds4_try_fused_verify_step( /*skip_property_check=*/force_verify_graph_replay, /*ds4_mix_mmv_max_tokens=*/0, /*mmvq_batch_invariant=*/n_lanes > 1); + // Experimental, off by default: on the two-box code run it filled the slot + // cache and the graphs thrashed through HIP-graph warmup (verify 42 -> 82 ms). + // LUCE_DS4_VERIFY_PREBUILD= (hybrid): while this graph runs, build up + // to n slots for the shapes the next step will most likely look up (this + // width, the next start after a full or one-short acceptance) when a free + // slot is left. A miss costs a ~15 ms rebuild on the critical path; here + // the host would only wait for the GPUs. + static const int prebuild_max = [] { + const char * v = std::getenv("LUCE_DS4_VERIFY_PREBUILD"); + return v && *v ? std::max(0, std::atoi(v)) : 0; + }(); + std::function prebuild; + if (prebuild_max > 0 && hybrid && q > 1) { + prebuild = [&]() { + int built = 0; + for (int a = q; a >= 1 && built < prebuild_max; --a) { + std::vector k = key_for(kv_start + a, q); + bool have = false; + for (size_t i = 0; i < slot_limit && !have; ++i) { + have = vc.slots[i].built() && vc.slots[i].shape_key == k; + } + if (have) continue; + size_t free_slot = slot_limit; + for (size_t i = 0; i < slot_limit; ++i) { + if (!vc.slots[i].built() && &vc.slots[i] != fg) { free_slot = i; break; } + } + if (free_slot == slot_limit) return; // never evict under a running graph + DeepSeek4FusedDecodeGraph & pf = vc.slots[free_slot]; + Ds4FusedVerifyCache::Extra & pe = vc.extra[free_slot]; + pf.destroy(vc.backend, vc.peer_backend); + pe.reset(); + pe.vocab_split = vocab_split; + const auto pb_t0 = Ds4TimingClock::now(); + if (!ds4_build_fused_verify_graph(mc, pf, pe, backend, w, &cache, + hc_weights, hc_out_weights, hash_tables, + kv_start + a, q, token_ids != nullptr, + *hooks->capture_layer_ids, hybrid, + std::move(k))) { + pf.destroy(vc.backend, vc.peer_backend); + pe.reset(); + return; + } + pf.last_use = 0; // unused prebuilt slots are the first to go + ++built; + static int reported = 0; + if (reported++ < 32) { + std::fprintf(stderr, "[ds4-fused-verify] prebuilt q=%d kv=%d in %.1f ms\n", q, + kv_start + a, ds4_elapsed_us(pb_t0, Ds4TimingClock::now()) / 1000.0); + } + } + }; + } const enum ggml_status compute_status = - ds4_fused_graph_compute(w, *fg, backend, hybrid, "ds4-fused-verify"); + ds4_fused_graph_compute(w, *fg, backend, hybrid, "ds4-fused-verify", prebuild); + 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 +2719,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 +2760,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 84674dd35f..1a82ec5441 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 (maikzz32/lucebox-halo-cluster): + // 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 rank holds only its attention heads " + "(LUCE_CLUSTER_SLICE_WEIGHTS) 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,47 @@ 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. + ggml_tensor * out_b_w = head_split + ? 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 +2358,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 +2375,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 +2393,75 @@ 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; +} + +// A lane batch runs a ratio-1 kv source's compressor for every lane in one +// width-q pass (LUCE_DS4_LANE_BATCH_COMP=0 keeps one pass per lane). Ratio 1 +// keeps no state: each lane's latent and index key are per-column products, +// and the MMVF block size does not depend on the column count on RDNA4, so +// the batched pass writes the rows the lanes would have written. +static bool ds4_lane_batch_comp_enabled() { + static const bool enabled = [] { + const char * v = std::getenv("LUCE_DS4_LANE_BATCH_COMP"); + return !v || std::strcmp(v, "0") != 0; + }(); + 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 +2509,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 +2528,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 +2670,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 +2826,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 +2840,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 +2903,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 +2912,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 +3048,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 +3163,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 +3327,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 +3433,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 +3556,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 +3627,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 +3657,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 +3711,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 +3912,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 +4107,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 +4128,7 @@ struct DeepSeek4PrefillHcPostGraph { block_out_cold = nullptr; split = nullptr; owner_join = false; + cluster_reduce = false; } }; @@ -4351,6 +4708,70 @@ static ggml_tensor * build_shared_ffn( return ggml_mul_mat(ctx, L.ffn_down_shexp, mid_sh); } +// Cluster shared-expert sharding (maikzz32/lucebox-halo-cluster): 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 must add it to the routed partial BEFORE the all-reduce, and the +// reduction then sums both contributions in one message. +// +// gate and up are sliced by output row, which is a contiguous view. `down` +// contracts over the intermediate axis, so slicing it would need a strided +// view of every row; instead the local intermediate is zero-padded back to +// full width and the full `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); + // LUCE_DS4_SHEXP_SLICE_VIEW=1: contract over this rank's slice only, + // through a column view of down (row stride unchanged): half the down + // weights read and no zero padding. Bit-identical to the padded product + // on the RDNA4 Q8_0 matvec, which gives the slice's blocks the lanes they + // take in the whole rows (ggml_backend_cuda_mul_mat_whole_row_lanes); + // other kernels sum the slice in their own order. + static const bool slice_view = [] { + const char * v = std::getenv("LUCE_DS4_SHEXP_SLICE_VIEW"); + return v && *v && std::strcmp(v, "0") != 0; + }(); + if (slice_view && (int64_t) ff_count != n_ff && + ff_begin % ggml_blck_size(L.ffn_down_shexp->type) == 0 && + ff_count % ggml_blck_size(L.ffn_down_shexp->type) == 0) { + 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; + } + if ((int64_t) ff_count != n_ff) { + 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; + } + mid_sh = padded; + } + return ggml_mul_mat(ctx, L.ffn_down_shexp, mid_sh); +} + bool init_deepseek4_streamed_expert_cache( const DeepSeek4Weights & w, const MoeHybridStorage & hybrid, @@ -6164,6 +6585,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 +6789,150 @@ 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: about 0.7 ms a step on the rank that holds the tables, +// and 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; +}; + +// LUCE_DS4_ENGRAM_ASYNC=0 keeps the in-line verify Engram read. +static bool ds4_engram_async_enabled() { + static const bool enabled = [] { + const char * v = std::getenv("LUCE_DS4_ENGRAM_ASYNC"); + return !(v && v[0] == '0' && v[1] == '\0'); + }(); + return enabled; +} + +// LUCE_DS4_PINNED_STEP_IO=0 keeps the blocking per-tensor verify uploads and +// readbacks. +static bool ds4_pinned_step_io_enabled() { + static const bool enabled = [] { + const char * v = std::getenv("LUCE_DS4_PINNED_STEP_IO"); + return !(v && v[0] == '0' && v[1] == '\0'); + }(); + return enabled; +} + 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 +6978,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 +6995,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; // the head covers this rank's vocabulary half (argmax only) 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 +7022,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 +7579,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 +7618,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 +7651,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 +8075,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 +8106,14 @@ 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() && + ds4_cluster_prefill_device_join_enabled(); 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 +8415,40 @@ 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) { + if (cluster_rt->cfg && + cluster_rt->cfg->shared_expert == cluster::SharedExpertMode::Rank0) { + std::fprintf(stderr, + "[deepseek4-cluster] layer %d: shared expert mode %s not implemented\n", + layer, cluster::shared_expert_mode_name(cluster_rt->cfg->shared_expert)); + return false; + } + 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 +8459,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 +8492,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 +8843,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); @@ -10027,11 +10737,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 +11167,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 +11255,19 @@ 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; + if (cache.cluster_rt && range_has_engram) { + // Engram tables live on rank 0: its rows replace every other rank's. + 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; + } + } + 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 +11351,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 +11797,15 @@ 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 attends with + // its own output groups and one all-reduce sums the partials. + // Every rank takes the same decision (from the band alone). + 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 || + ds4_cluster_prefill_attention_split(kv_start + n_tokens)); 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 +11814,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); @@ -11542,6 +12299,7 @@ bool deepseek4_step_layer_range( !device_input_env || !*device_input_env || std::strcmp(device_input_env, "0") != 0; const bool ffn_device_join_possible = + (!cache.cluster_rt || ds4_cluster_prefill_device_join_enabled()) && device_input_enabled && moe_expert_major_prefill_enabled(n_tokens) && layer_storage.cold_backend_kind == @@ -11596,7 +12354,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,11 +12756,27 @@ 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(); + 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 && ds4_cluster_prefill_pipeline_enabled()) { + 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) { @@ -12388,9 +13163,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 +13183,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 +13228,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 +13260,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 +13280,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 ggml_add(ctx, routed, shared); } // Read a small F32 GPU tensor (HC scale, [k]) into host floats. @@ -12517,119 +13359,164 @@ namespace { // 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 { +// 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; + PinnedStage stage; // inputs and the readback of one projection + + // The window stays on the draft device (LUCE_DS4_DRAFT_DEVICE_CONTEXT), + // [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{}; + } }; +// Default on; LUCE_DS4_DRAFT_DEVICE_CONTEXT=0 keeps the host-side window. +static bool dspark_device_context_enabled() { + static const bool on = [] { + const char * v = std::getenv("LUCE_DS4_DRAFT_DEVICE_CONTEXT"); + return !(v && v[0] == '0' && v[1] == '\0'); + }(); + return on; +} + thread_local DsparkContextKvProjector g_dspark_ctx_kv; +// Projects n_cols new feature columns to the drafter's per-layer context K/V. +// out != null: read the [head_dim, n_cols, n_layer] result back (one wait). +// out == null: leave it on the device 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) { + std::vector * out, + ggml_tensor ** dev_out = nullptr) { 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); + // Inputs and the readback go through pinned slices: one wait in all. + const size_t in_bytes = sizeof(float) * (size_t) fc_in * n_cols; + const size_t out_bytes = sizeof(float) * (size_t) head_dim * n_cols * w.n_layer; + const bool pinned = ds4_pinned_step_io_enabled() && + P.stage.reserve(backend, in_bytes + out_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) { + if (pinned) { + 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()); + } else { + ggml_backend_tensor_set(G.inp_features, features, 0, in_bytes); + ggml_backend_tensor_set(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()); + if (!out) { + if (dev_out) *dev_out = G.out; + return true; + } + out->resize((size_t) head_dim * n_cols * w.n_layer); + const void * src = pinned ? P.stage.get(backend, G.out, 0, out_bytes) : nullptr; + ggml_backend_synchronize(backend); + if (src) std::memcpy(out->data(), src, out_bytes); + else ggml_backend_tensor_get(G.out, out->data(), 0, out_bytes); return true; } @@ -12639,12 +13526,82 @@ 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 = nullptr) { 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 (out_dev_kv && dspark_device_context_enabled() && 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, + nullptr, &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.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_len <= 0) { P.end_pos = committed; P.valid = 0; @@ -12671,7 +13628,7 @@ static bool dspark_update_context_kv_cache( 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)) { + backend, d, new_features, n_new, first_new_pos, &P.projected)) { return false; } @@ -12722,8 +13679,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 +13708,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 +13739,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 +13836,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 +13864,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 +13920,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 +14007,34 @@ 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 = ds4_pinned_step_io_enabled() && + 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(); + const auto set_in = [&](ggml_tensor * t, const void * data, size_t bytes) { + if (pinned_io) C.stage.set(backend, t, data, 0, bytes); + else ggml_backend_tensor_set(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 +14053,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 +14069,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 +14083,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 +14121,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 = pinned_io ? C.stage.get(backend, C.out, 0, hidden_bytes) : nullptr; + const void * conf_src = pinned_io && 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 f2bd8fccfe..4dc41438ae 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 @@ -507,7 +519,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 +610,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 +855,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 009ed16265..f3c0fe7559 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,74 @@ 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; + } + 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)) { + s.kind = DS4TensorSlice::None; // not block aligned: keep the full tensor + } + return s; + }; + size_t sliced_bytes_saved = 0; + std::vector allocs; allocs.reserve(n_tensors); size_t total_buf_size = 0; @@ -2275,6 +2367,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( @@ -2385,6 +2485,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 +2605,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/server/launch_profiles.h b/server/src/server/launch_profiles.h index 408c15ba8a..d5db315bb4 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -232,6 +232,62 @@ 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 one to 0 turns it off. +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_ZC", "1"}, + {"LUCE_CLUSTER_FAST_REDUCE_BF16", "1"}, + {"LUCE_CLUSTER_FAST_REDUCE_LEAN", "1"}, + {"LUCE_CLUSTER_FAST_REDUCE_STALL_S", "15"}, + {"LUCE_CLUSTER_HYBRID_EXCHANGE", "1"}, + {"LUCE_CLUSTER_HYBRID_ONE_EXPORT", "1"}, + // Prefill-sized reductions stay on RCCL, in bf16. + {"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"}, + {"LUCE_CLUSTER_SLICE_WEIGHTS", "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_SHEXP_SLICE_VIEW", "1"}, + {"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, the ROCmFP2 prefetch loop and one peer-copy kernel per + // layer. + {"LUCE_MOE_Q8_HANDOFF", "1"}, + {"LUCE_MMID_GROUPED", "1"}, + {"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 88709908c8..3127181872 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -1178,27 +1178,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); diff --git a/server/test/unit/test_deepseek4_unit.cpp b/server/test/unit/test_deepseek4_unit.cpp index 52c5223028..a152912909 100644 --- a/server/test/unit/test_deepseek4_unit.cpp +++ b/server/test/unit/test_deepseek4_unit.cpp @@ -6252,6 +6252,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 ..."); @@ -8984,6 +9159,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; @@ -9131,6 +9310,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_launch_policy.cpp b/server/test/unit/test_launch_policy.cpp index c8d2b27f76..2c9ddd4f87 100644 --- a/server/test/unit/test_launch_policy.cpp +++ b/server/test/unit/test_launch_policy.cpp @@ -149,6 +149,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); @@ -267,6 +280,7 @@ TEST_CASE(LaunchPolicyFixture, launch_policy_suite) { test_auto_device_counts_the_kv_cache(); 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(); From 228b5bf2eca1958012c6054cf232f42d85c53cae Mon Sep 17 00:00:00 2001 From: mrciffa Date: Thu, 8 Oct 2026 16:22:43 +0000 Subject: [PATCH 06/40] perf(ggml-cuda): one launch for an RDNA4 Q8_0 gate/up pair with SwiGLU-DS4 A shared expert's gate and up products (Q8_0, several columns in a verify) run as two batch-invariant RDNA4 launches and SwiGLU-DS4 as a third. mul_mat_vec_q8_0_rdna4 gets a GLU variant that reads both matrices with one load of each activation fragment, keeps each product's partials and reduction exactly as the single kernel does, and writes ggml_cuda_op_swiglu_ds4_single(gate, up) in its epilogue: the value the two products and the GLU kernel give, in one launch. The graph matcher also finds the pair with weight views between the products (the cluster shared expert's row slices; views launch nothing). ggml_cuda_mmvq_rdna4_glu_pair gates it: RDNA4, batch-invariant products, 2..8 columns, a Q8_0 pair on the same activations, a non-swapped SwiGLU-DS4 that reads exactly their outputs, no biases. LUCE_CUDA_MMVQ_Q8_RDNA4_GLU=0 keeps three launches. The pair has its own launcher on one block shape (two rows per wave, four waves) whose partials of both products fit the LDS at eight columns: seven new kernel instances, the dispatch table unchanged. test_ds4_fused_ops_cuda: the pair against the products kept apart, and a whole-row-lane column slice against the zero-padded product, bit for bit at 1..8 columns. 0 failures on gfx1201 and gfx1151; a kernel trace on gfx1201 shows the fused launch at 2, 4, 5 and 8 columns. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 44 +++++ .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 168 ++++++++++++++++-- .../llama.cpp/ggml/src/ggml-cuda/mmvq.cuh | 2 + server/test/unit/test_ds4_fused_ops_cuda.cpp | 131 ++++++++++++++ 4 files changed, 334 insertions(+), 11 deletions(-) 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 a65c8af374..029d5b3d92 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 @@ -3140,6 +3140,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]; @@ -6284,6 +6297,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; 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 fa94bf0081..811492e6cc 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -2602,7 +2602,10 @@ 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, @@ -2610,7 +2613,7 @@ static __global__ void mul_mat_vec_q8_0_rdna4( 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 kb_off) { + 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"); @@ -2632,22 +2635,30 @@ 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; + } // 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 @@ -2682,12 +2693,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) @@ -2699,15 +2724,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; @@ -2722,6 +2759,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; + } } } @@ -2731,7 +2776,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; } } @@ -2746,6 +2796,44 @@ static bool mmvq_q8_0_rdna4_enabled() { return enabled; } +// 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). LUCE_CUDA_MMVQ_Q8_RDNA4_GLU=0 +// keeps three launches. +bool ggml_cuda_mmvq_rdna4_glu_pair(const ggml_tensor * gate, const ggml_tensor * up, const ggml_tensor * glu) { + static const bool enabled = []() { + const char * e = std::getenv("LUCE_CUDA_MMVQ_Q8_RDNA4_GLU"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + if (!enabled || !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; +} + template static void mul_mat_vec_q8_0_rdna4_launch_cfg( const void * vx, const void * vy, float * dst, @@ -2760,7 +2848,31 @@ static void mul_mat_vec_q8_0_rdna4_launch_cfg( 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, kb_off); + 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 @@ -2791,13 +2903,34 @@ 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), 2..8 columns, no column range. 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, const int kb_off, 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(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, \ @@ -2993,7 +3126,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()) { @@ -3001,11 +3139,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, kb_off, 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) \ @@ -3564,7 +3706,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]); 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 af82a0b89d..ce116473af 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); diff --git a/server/test/unit/test_ds4_fused_ops_cuda.cpp b/server/test/unit/test_ds4_fused_ops_cuda.cpp index 4606e61599..d4b4697766 100644 --- a/server/test/unit/test_ds4_fused_ops_cuda.cpp +++ b/server/test/unit/test_ds4_fused_ops_cuda.cpp @@ -475,6 +475,124 @@ bool run_glu_q8(ggml_backend_t backend, int n_ff, int n_used, int tokens, uint32 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). +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). +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) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + 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. @@ -793,6 +911,19 @@ int main() { 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 pair_cases[][5] = {{1024, 64, 32, 0, 5}, {1024, 64, 32, 32, 4}, {5120, 96, 48, 48, 2}, {1024, 64, 32, 32, 8}}; + 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; From 40c4808ef06b5850d5acd940106af95d2e3ea290 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Thu, 8 Oct 2026 16:22:43 +0000 Subject: [PATCH 07/40] perf(ggml-cuda): grouped ROCmFP2 expert products use the prefetching loop The grouped MoE MMVQ (LUCE_MMID_GROUPED) ran ROCmFP2 experts on the generic loop, where every iteration waits for a row's scale byte before the next loads leave: on the Strix Halo iGPU the cold experts read at ~100 GB/s. The per-expert kernel already has a prefetching loop for this (mul_mat_vec_q_moe_fp2_prefetch: the loads of several K iterations first, then the dot products in the same kbx order into the same accumulators, so the result is bit-identical). The grouped kernel now uses it too, and both kernels read one switch: on by default on gfx1151, LUCE_CUDA_MMVQ_MOE_FP2_PREFETCH=0/1 overrides (cluster mode sets 1). Two boxes, DS4.1, with the gate/up launch (previous commit), 2 x 3 runs against both off, same text on both prompts: step 42.1-42.4 -> 41.5-41.8 ms (code), 42.1-43.9 -> 41.4-41.7 ms (2K); code 100.7-101.2 -> 102.1-102.7 tok/s, 2K 55.8-58.8 -> 59.3-60.0. test_deepseek4_mmid_grouped_cuda passes on gfx1201 and gfx1151 with the switch at its default and forced on. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 46 +++++++++++++------ 1 file changed, 32 insertions(+), 14 deletions(-) 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 811492e6cc..6c96ee0619 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -1874,7 +1874,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, @@ -1950,6 +1951,12 @@ 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}; + // 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)); @@ -2049,7 +2056,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, @@ -2067,7 +2075,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, @@ -2076,7 +2084,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, @@ -2085,6 +2093,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, @@ -2136,6 +2156,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>( @@ -2453,16 +2480,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) \ From 4f2c69333d7d951973339fb7a1d4a72e868c80a3 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:46:09 +0200 Subject: [PATCH 08/40] fix(ggml-cuda): CUDA builds of the gate/up pair and the grouped prefetch The two commits above built on HIP only. ggml_cuda_mmvq_rdna4_glu_pair sat in the HIP-only RDNA4 block while ggml-cuda.cu calls it on every build: it is now defined everywhere and answers false outside HIP. nvcc rejects the grouped kernel's unbraced else before the K loop ("expected a }"): the loop now sits in braces, the way the per-expert kernel writes it. No change on HIP. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 146 +++++++++--------- 1 file changed, 76 insertions(+), 70 deletions(-) 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 6c96ee0619..9634dd505f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -1956,40 +1956,41 @@ static __global__ void mul_mat_vec_q_moe_grouped( 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)); + } 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); + } } } } @@ -2814,44 +2815,6 @@ static bool mmvq_q8_0_rdna4_enabled() { return enabled; } -// 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). LUCE_CUDA_MMVQ_Q8_RDNA4_GLU=0 -// keeps three launches. -bool ggml_cuda_mmvq_rdna4_glu_pair(const ggml_tensor * gate, const ggml_tensor * up, const ggml_tensor * glu) { - static const bool enabled = []() { - const char * e = std::getenv("LUCE_CUDA_MMVQ_Q8_RDNA4_GLU"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); - if (!enabled || !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; -} - template static void mul_mat_vec_q8_0_rdna4_launch_cfg( const void * vx, const void * vy, float * dst, @@ -2968,6 +2931,49 @@ 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). LUCE_CUDA_MMVQ_Q8_RDNA4_GLU=0 +// keeps three launches. +bool ggml_cuda_mmvq_rdna4_glu_pair(const ggml_tensor * gate, const ggml_tensor * up, const ggml_tensor * glu) { +#if defined(GGML_USE_HIP) + static const bool enabled = []() { + const char * e = std::getenv("LUCE_CUDA_MMVQ_Q8_RDNA4_GLU"); + return !(e && e[0] == '0' && e[1] == '\0'); + }(); + if (!enabled || !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 From a49bc3ef00259831dfdecb68f0ff28731180a280 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 8 Oct 2026 22:46:24 +0200 Subject: [PATCH 09/40] test(ds4): the whole-row-lane slice is bit-exact on HIP, close on CUDA ggml_backend_cuda_mul_mat_whole_row_lanes promises the padded product's bits only where the lanes are implemented (HIP); other kernels compute the range's product as usual. The column-slice check now compares bits on HIP and the value to 1e-5 relative elsewhere: on the DGX GB10 (sm_121) the slice and the padded product differ in the last bit (-1.77656889 vs -1.77656901), which failed the CUDA GPU job. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/test/unit/test_ds4_fused_ops_cuda.cpp | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/server/test/unit/test_ds4_fused_ops_cuda.cpp b/server/test/unit/test_ds4_fused_ops_cuda.cpp index d4b4697766..f5f99272cf 100644 --- a/server/test/unit/test_ds4_fused_ops_cuda.cpp +++ b/server/test/unit/test_ds4_fused_ops_cuda.cpp @@ -536,7 +536,15 @@ bool run_glu_pair(ggml_backend_t backend, int n_in, int n_full, int n_ff, int ff // 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). +// 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(); @@ -580,7 +588,10 @@ bool run_col_slice(ggml_backend_t backend, int n_out, int n_full, int k_begin, i 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) { + 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; From 53d673dd32667fe385e02b88357dc72290439947 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 9 Oct 2026 00:44:42 +0200 Subject: [PATCH 10/40] fix(cluster): worker exit, two unparsed flags, help text, watchdog at 0 - A cluster worker runs to completion inside load_model and has no HTTP server; main() went on to call run() on the null server, so every clean shutdown of rank 1 crashed. It now returns the worker's exit code. - --cluster-selftest and --cluster-replicate-hot were documented, wired to the head, the worker and the placement, but rejected by the parser. - The cluster lines of --help sat inside --ds4-prefill's description. - LUCE_CLUSTER_FAST_REDUCE_STALL_S=0 made every wait a stall; 0 now turns the watchdog off, and the defaults comment says which entries are sizes. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/cluster_config.h | 7 ++++--- server/src/cluster/fast_reduce.cpp | 2 +- server/src/server/launch_profiles.h | 4 +++- server/src/server/server_main.cpp | 16 ++++++++++++++-- 4 files changed, 22 insertions(+), 7 deletions(-) diff --git a/server/src/cluster/cluster_config.h b/server/src/cluster/cluster_config.h index 846009e009..1293e37fd2 100644 --- a/server/src/cluster/cluster_config.h +++ b/server/src/cluster/cluster_config.h @@ -10,9 +10,10 @@ // (communicator handle, resolved placement, peer identities) live in the // cluster runtime objects, not here. // -// Policy (variables.md): features ship as CLI flags. Environment variables are -// reserved for burn-in kill switches (LUCE_CLUSTER_NO_INGRAPH_ALLREDUCE, -// LUCE_CLUSTER_NO_GRAPH_CAPTURE) and debug tracing (LUCE_CLUSTER_TRACE). +// 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 diff --git a/server/src/cluster/fast_reduce.cpp b/server/src/cluster/fast_reduce.cpp index 02f3647714..69b0075eee 100644 --- a/server/src/cluster/fast_reduce.cpp +++ b/server/src/cluster/fast_reduce.cpp @@ -594,7 +594,7 @@ bool FastReduce::init(const Config & cfg, std::string * err) { { const double w = std::chrono::duration( std::chrono::steady_clock::now() - wait_t0).count(); - if (w > stall_s && s.seq.load() >= next) { // allocated, never published + if (stall_s > 0.0 && w > stall_s && s.seq.load() >= next) { // allocated, never published s.failed.store(true); const uint64_t last = s.seq.load(); std::fprintf(stderr, diff --git a/server/src/server/launch_profiles.h b/server/src/server/launch_profiles.h index d5db315bb4..c25a4b7889 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -235,7 +235,9 @@ inline const std::vector & launch_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 one to 0 turns it off. +// 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 diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 3127181872..80f8d1c096 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -177,6 +177,8 @@ static void print_usage(const char * prog) { " Experts that keep a token's native routing and stay on\n" " the primary GPU (JSON {\"layer\": [ids]})\n" " --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" @@ -188,9 +190,9 @@ static void print_usage(const char * prog) { " --cluster-shared-expert replicate|shard (default: replicate)\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-verify-hash Debug: cross-rank hidden-state hash every n steps\n" - " (default: exact; dense/sparse are experimental\n" - " and may change generated tokens)\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" @@ -649,6 +651,13 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, 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 if (std::strcmp(flag, "--cluster-verify-hash") == 0) { if (!int_value(bargs.cluster.verify_hash_every) || bargs.cluster.verify_hash_every < 0) { std::fprintf(stderr, "[server] --cluster-verify-hash must be non-negative\n"); @@ -2354,6 +2363,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)); } From bcbfc8ca558fb1653aaba242f164ec054fe7e93d Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:40:52 +0200 Subject: [PATCH 11/40] fix(cluster): bounded fast-reduce handshake, stall clock, RDMA-error release - The bootstrap exchange waits at most timeout_ms for each accept, read and write (cluster_control's deadline helpers), closes every socket it opened on failure and refuses a duplicate rank: a worker that died during its load fails rank 0's set-up instead of parking it. - The unpublished-reduction watchdog starts its clock when the host hands the reduction out, not when the progress thread starts waiting, so an idle server no longer trips it between requests. - A failed ibv_post_send or a failed completion now releases every reduction still in flight and exits, like a stall, instead of stopping the progress thread with the GPU parked on a flag that never comes. Both use one release_and_exit, which says whether the GPU drained or was only released. - The zero-copy path allocates no device scratch (80 MiB at the cluster sizes); the per-device publish counters are freed at shutdown; the submit trace and the unused system-scope flag kernel are gone; the kernel-flag and import switches read "0" as off like the others; comments match the code. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/fast_reduce.cpp | 220 ++++++++++++---------- server/src/cluster/fast_reduce.h | 20 +- server/src/cluster/fast_reduce_kernels.cu | 44 ++--- 3 files changed, 146 insertions(+), 138 deletions(-) diff --git a/server/src/cluster/fast_reduce.cpp b/server/src/cluster/fast_reduce.cpp index 69b0075eee..8cb0e97349 100644 --- a/server/src/cluster/fast_reduce.cpp +++ b/server/src/cluster/fast_reduce.cpp @@ -7,6 +7,9 @@ #include #include +#include + +#include "cluster/cluster_control.h" #include #include @@ -16,6 +19,7 @@ #include #include +#include #include #include #include @@ -39,7 +43,6 @@ void fast_reduce_launch_hyb_combine_pub(const float * cold, const float * hot, c void fast_reduce_launch_hyb_final_pub(const void * own_row, const void * peer_row, float * dst, int n, int bf16, uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream); void fast_reduce_launch_add_zc(float * dst, const void * peer, int n, int bf16, hipStream_t stream); -void fast_reduce_launch_flag_sys(uint32_t * flag, uint32_t value, hipStream_t stream); void fast_reduce_launch_hyb_combine(const float * cold, const float * hot, const float * shared, void * own_row, int n, int bf16, hipStream_t stream); void fast_reduce_launch_hyb_final(const void * own_row, const void * peer_row, float * dst, int n, int bf16, @@ -86,28 +89,6 @@ bool set_nodelay(int fd) { return setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one)) == 0; } -bool write_all(int fd, const void * p, size_t n) { - const char * c = (const char *) p; - while (n) { - const ssize_t w = ::write(fd, c, n); - if (w <= 0) return false; - c += w; - n -= (size_t) w; - } - return true; -} - -bool read_all(int fd, void * p, size_t n) { - char * c = (char *) p; - while (n) { - const ssize_t r = ::read(fd, c, n); - if (r <= 0) return false; - c += r; - n -= (size_t) r; - } - return true; -} - // 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, @@ -120,6 +101,10 @@ bool exchange(const FastReduce::Config & cfg, const Endpoint & mine, 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); @@ -133,29 +118,47 @@ bool exchange(const FastReduce::Config & cfg, const Endpoint & mine, 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); return fail("accept"); } + if (fd < 0) { ::close(lst); close_peers(); return fail("accept"); } set_nodelay(fd); int32_t r = -1; Endpoint e{}; - if (!read_all(fd, &r, sizeof(r)) || !read_all(fd, &e, sizeof(e)) || - r <= 0 || r >= cfg.size) { - ::close(fd); ::close(lst); - return fail("read endpoint"); + 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_all(peers[(size_t) i], all.data(), sizeof(Endpoint) * all.size())) { - for (int fd : peers) if (fd >= 0) ::close(fd); - return fail("write table"); + 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); } } - for (int fd : peers) if (fd >= 0) ::close(fd); + close_peers(); return true; } @@ -180,10 +183,12 @@ bool exchange(const FastReduce::Config & cfg, const Endpoint & mine, if (fd < 0) return fail("connect"); set_nodelay(fd); const int32_t r = cfg.rank; - if (!write_all(fd, &r, sizeof(r)) || !write_all(fd, &mine, sizeof(mine)) || - !read_all(fd, all.data(), sizeof(Endpoint) * all.size())) { + 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("exchange"); + return fail_io("exchange", ioerr); } ::close(fd); return true; @@ -199,6 +204,15 @@ static bool fr_env_on(const char * name) { return v && *v && std::strcmp(v, "0") != 0; } +// Two ranks with LUCE_CLUSTER_FAST_REDUCE_ZC=1 (and the stream flags) take the +// zero-copy path: pack into the pinned row, add the peer's row straight out of +// pinned memory. It needs no device scratch. +static bool fr_zero_copy(int size) { + static const bool zc = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_ZC") && + !fr_env_on("LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS"); + return zc && size == 2; +} + // LUCE_CLUSTER_FAST_REDUCE_LEAN=1 drops the per-reduction slot-reuse waits and // done writes on the zero-copy and hybrid paths. Neither carries data: the // ranks advance in lockstep (a rank publishes reduction q only after it added @@ -552,7 +566,8 @@ bool FastReduce::init(const Config & cfg, std::string * err) { if (hipHostGetDevicePointer((void **) &s.peer_flags_dev, s.peer_flags, 0) != hipSuccess) { return bail("hipHostGetDevicePointer for the flag table failed"); } - if (hipMalloc((void **) &s.scratch, + if (!fr_zero_copy(cfg.size) && + hipMalloc((void **) &s.scratch, rows * (size_t) cfg.max_elems * sizeof(float)) != hipSuccess) { return bail("hipMalloc for the peer scratch failed"); } @@ -562,7 +577,6 @@ bool FastReduce::init(const Config & cfg, std::string * err) { // ── progress ────────────────────────────────────────────────────────── s.up = true; s.progress = std::thread([&s]() { - const int n_peers_l = s.cfg.size - 1; uint64_t next = 1; std::vector wc(64); // Stall release: every flag a stream can be parked on for reduction q, @@ -578,12 +592,42 @@ bool FastReduce::init(const Config & cfg, std::string * err) { *(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); + const auto drain_t0 = std::chrono::steady_clock::now(); + bool drained = false; // the lean paths write no done flags: wait the full time + while (!drained && 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); // late submits too + drained = *(volatile uint32_t *) s.done_flag((int) (last % (uint64_t) s.cfg.slots)) == (uint32_t) last; + if (!drained) std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + std::fprintf(stderr, "[fast-reduce] rank %d: %s, exiting\n", s.cfg.rank, + drained ? "GPU drained" : "GPU released"); + 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; static const double stall_s = std::getenv("LUCE_CLUSTER_FAST_REDUCE_STALL_S") ? std::atof(std::getenv("LUCE_CLUSTER_FAST_REDUCE_STALL_S")) : 20.0; while (*pub != (uint32_t) next) { @@ -591,24 +635,20 @@ bool FastReduce::init(const Config & cfg, std::string * err) { 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() - wait_t0).count(); - if (stall_s > 0.0 && w > stall_s && s.seq.load() >= next) { // allocated, never published - s.failed.store(true); - const uint64_t last = s.seq.load(); + 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) (last - next)); - const auto t0 = std::chrono::steady_clock::now(); - while (std::chrono::duration(std::chrono::steady_clock::now() - t0).count() < 10.0) { - for (uint64_t q = next; q <= s.seq.load(); ++q) release_one(q); - std::this_thread::sleep_for(std::chrono::milliseconds(20)); - } - std::fflush(nullptr); - std::_Exit(4); + (unsigned long long) (s.seq.load() - next)); + release_and_exit(next); } } if (next <= 1) continue; @@ -644,28 +684,14 @@ bool FastReduce::init(const Config & cfg, std::string * err) { (unsigned long long) p, (unsigned long long) s.seq.load()); } - if (waited < stall_s) continue; - // Stall: release every reduction the host has handed the GPU so - // its queue drains, make later submits no-ops, and exit once the - // GPU is idle. Exiting with a parked queue costs a GPU reset. - s.failed.store(true); + 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)); - for (uint64_t q = p; q <= std::max(p, last); ++q) release_one(q); - 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 l2 = std::max(p, s.seq.load()); - for (uint64_t q = p; q <= l2; ++q) release_one(q); // late submits too - if (*(volatile uint32_t *) s.done_flag((int) (l2 % (uint64_t) s.cfg.slots)) == (uint32_t) l2) break; - std::this_thread::sleep_for(std::chrono::milliseconds(20)); - } - std::fprintf(stderr, "[fast-reduce] rank %d: GPU drained, exiting\n", s.cfg.rank); - std::fflush(nullptr); - std::_Exit(4); + 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 @@ -720,25 +746,25 @@ bool FastReduce::init(const Config & cfg, std::string * err) { w.next = &fw; ibv_send_wr * bad = nullptr; if (ibv_post_send(s.qps[(size_t) r], &w, &bad) != 0) { - s.stop.store(true); - return; + 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 it stops the thread and the - // waiting kernels time out into a counted, visible failure. + // 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)\n", + "[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); - s.stop.store(true); - return; + release_and_exit(oldest_in_flight(next)); } } } while (done > 0); @@ -762,18 +788,7 @@ bool FastReduce::init(const Config & cfg, std::string * err) { 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 - // The first handful of decisions, taken or declined, so a stall says which - // reduction it stopped at rather than only that it stopped. - static std::atomic traced{0}; - const bool trace = traced.fetch_add(1, std::memory_order_relaxed) < 24; - if (!s.up || n == 0 || n > (size_t) s.cfg.max_elems) { - if (trace) { - std::fprintf(stderr, "[fast-reduce] declined %zu floats (cap %d)\n", - n, s.cfg.max_elems); - } - return false; - } - if (trace) std::fprintf(stderr, "[fast-reduce] took %zu floats\n", n); + 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); @@ -805,8 +820,7 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { // involved at all -- which is what this needs and what a kernel cannot be. // LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS=1 restores the kernel form for // comparison. - static const bool kernel_flags = - std::getenv("LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS") != nullptr; + static const bool kernel_flags = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS"); // Do not write into a slot the progress thread has not finished with. The // kernel form waited on this; the stream form has to as well, or the GPU -- @@ -815,8 +829,7 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { // third request of a run stop rather than the first. // The previous user of this slot was seq - slots, and it leaves exactly that // value behind. Asking for anything higher is a wait nothing can satisfy. - static const bool zc_lean = fast_reduce_lean() && fr_env_on("LUCE_CLUSTER_FAST_REDUCE_ZC") && - s.cfg.size == 2; + static const bool zc_lean = fast_reduce_lean() && fr_zero_copy(s.cfg.size); if (!kernel_flags && !zc_lean && seq > (uint64_t) s.cfg.slots) { const uint32_t need = (uint32_t) (seq - (uint64_t) s.cfg.slots); if (hipStreamWaitValue32(stm, s.flag_dev + s.done_off(slot), need, @@ -827,10 +840,9 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { // Zero-copy (discrete GPU): pack straight into the pinned row, publish, wait // for the peer, add its row straight out of pinned memory, mark the slot done. - static const bool zc = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_ZC"); static const int zc_bf16 = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16") ? 1 : 0; const int wire_bf16 = exact ? 0 : zc_bf16; - if (zc && !kernel_flags && s.cfg.size == 2) { + if (fr_zero_copy(s.cfg.size)) { float * own_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) s.cfg.rank * s.cfg.max_elems; const int peer = 1 - s.cfg.rank; const float * peer_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) peer * s.cfg.max_elems; @@ -997,7 +1009,7 @@ bool FastReduce::hyb_import(float * data, size_t n, void * stream, uint64_t k) { auto stm = (hipStream_t) stream; if (hipStreamWaitValue32(stm, s.flag_dev + s.hybf_off(slot), (uint32_t) seq, hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; - static const bool import_dma = std::getenv("LUCE_CLUSTER_HYBRID_IMPORT_KERNEL") == nullptr; + static const bool import_dma = !fr_env_on("LUCE_CLUSTER_HYBRID_IMPORT_KERNEL"); if (import_dma) { if (hipMemcpyAsync(data, s.hyb_part(2, slot), n * sizeof(float), hipMemcpyHostToDevice, stm) != hipSuccess) return false; } else { @@ -1032,13 +1044,21 @@ void FastReduce::shutdown() { if (s.pd) { ibv_dealloc_pd(s.pd); s.pd = nullptr; } if (s.ctx) { ibv_close_device(s.ctx); s.ctx = nullptr; } - if (s.peer_flags) { hipHostFree(s.peer_flags); s.peer_flags = nullptr; } - if (s.scratch) { hipFree(s.scratch); s.scratch = nullptr; } - if (s.data_host) { hipHostFree(s.data_host); s.data_host = nullptr; } - if (s.flag_host) { hipHostFree(s.flag_host); s.flag_host = nullptr; } - if (s.timed_out_host) { hipHostFree(s.timed_out_host); s.timed_out_host = nullptr; } - if (s.progress_host) { hipHostFree(s.progress_host); s.progress_host = nullptr; } - if (s.hyb_host) { hipHostFree(s.hyb_host); s.hyb_host = 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.peer_flags) { (void) hipHostFree(s.peer_flags); s.peer_flags = nullptr; } + if (s.scratch) { (void) hipFree(s.scratch); s.scratch = nullptr; } + 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.timed_out_host) { (void) hipHostFree(s.timed_out_host); s.timed_out_host = nullptr; } + if (s.progress_host) { (void) hipHostFree(s.progress_host); s.progress_host = nullptr; } + if (s.hyb_host) { (void) hipHostFree(s.hyb_host); s.hyb_host = nullptr; } } } // namespace luce::cluster diff --git a/server/src/cluster/fast_reduce.h b/server/src/cluster/fast_reduce.h index 637f29f166..48f389acc8 100644 --- a/server/src/cluster/fast_reduce.h +++ b/server/src/cluster/fast_reduce.h @@ -16,9 +16,9 @@ // 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. Unified memory. On Strix Halo there is no separate VRAM, so -// a pinned host buffer *is* GPU memory: the NIC writes into it and a kernel -// reads it without a copy and without GPUDirect. It also assumes the ranks run +// WHAT IT ASSUMES. Pinned host memory that both the NIC and the GPU reach: the +// NIC writes each peer's row into it, and the GPU reads it there or after a DMA +// copy, 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. @@ -56,12 +56,12 @@ class FastReduce { int gid_index = 1; // RoCE v2 std::string bootstrap_host; // rank 0's address int bootstrap_port = 9500; - // Decode sizes only -- 10240 floats is the largest a step reduces -- - // and a ring deep enough that the guard against lapping it almost never - // fires. It has to be deep: the guard is a stream wait, and a stream - // wait that fires stops the GPU running ahead, which is where this - // path's speed comes from. Prefill's much larger reductions fall back - // to RCCL, where they cost nothing that matters. + // Sized for a plain decode step (10240 floats is the largest it + // reduces); DeepSeek V4 sizes both for its verify batch. The ring has + // to be deep: the guard against lapping it is a stream wait, and a + // stream wait that fires stops the GPU running ahead, which is where + // this path's speed comes from. Prefill's much larger reductions fall + // back to RCCL, where they cost nothing that matters. int max_elems = 16384; // 64 KiB int slots = 4096; // a step has ~97 uint32_t timeout_ms = 30000; @@ -104,7 +104,7 @@ class FastReduce { 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); - // Reductions issued and how long the progress thread waited for the GPU, + // Reductions issued, and flag waits of the kernel form that timed out, // for the telemetry line. uint64_t submitted() const; uint64_t timed_out() const; diff --git a/server/src/cluster/fast_reduce_kernels.cu b/server/src/cluster/fast_reduce_kernels.cu index 93e81376db..e4bab2fff7 100644 --- a/server/src/cluster/fast_reduce_kernels.cu +++ b/server/src/cluster/fast_reduce_kernels.cu @@ -1,21 +1,20 @@ // fast_reduce_kernels.cu - the device half of the lean decode-path reduction. // -// Three tiny kernels per reduction, all on the backend's own stream so they sit -// in the graph exactly where the collective used to: +// The flags normally go by stream memory operations (hipStreamWriteValue32 and +// hipStreamWaitValue32, in fast_reduce.cpp). The kernels here are the rest: // -// set_flag raise this rank's publish flag once the payload has been -// copied out. The copy itself is a DMA, enqueued ahead of this -// on the same stream, so stream order is what publishes it. -// wait_flags spin until every peer's flag has arrived. -// add_peers add the peers' partials, by then copied into device memory. +// pack, add_zc the zero-copy path: pack this rank's row into pinned memory +// (f32 or bf16) and add the peer's row straight from it. +// add the copy path: add the peers' partials, by then copied into +// device memory. +// hyb_* the two-rank hybrid exchange (FastReduce::hyb_combine). +// set_flag, the kernel form of the flags, kept for comparison +// wait_flags (LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS=1). // -// WHY THE PAYLOAD IS NOT MOVED BY A KERNEL. It was, and it cost 33 ms per -// reduction against RCCL's 117 us -- 250 times worse than the thing it was -// meant to replace. Pinned host memory is mapped uncached for the GPU, so a -// kernel reading or writing bulk data there crawls; the copy engine does not -// care, because it is not reading through the GPU's caches. So the payload goes -// by hipMemcpyAsync and only the flags -- four bytes, three times per -// reduction -- are touched by a kernel at zero-copy speed. +// WHY THE COPY PATH MOVES THE PAYLOAD BY DMA. A kernel moving it cost 33 ms per +// reduction against RCCL's 117 us. Pinned host memory is mapped uncached for +// the GPU, so a kernel reading or writing bulk data there crawls; the copy +// engine does not care, because it is not reading through the GPU's caches. // // WHY THE FLAGS ARE PINNED AND NOT MANAGED. gfx1151 reports XNACK disabled, so // the GPU cannot take a page fault, and a managed page the CPU has touched @@ -23,8 +22,9 @@ // no fault to be reached -- and the round trip through one measures 0.79 us // (server/test/cluster_flag_latency.cu), which is what makes this worth doing. // -// Every spin is bounded and counted. A rank that dies must surface as a visible -// failure, not as a GPU that never returns. +// 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. +// The kernel form's spin is bounded and counted (FastReduce::timed_out). #include @@ -250,18 +250,6 @@ void fast_reduce_launch_hyb_final(const void * own_row, const void * peer_row, f fast_reduce_launch_hyb_final_pub(own_row, peer_row, dst, n, bf16, nullptr, 0, nullptr, stream); } -// A word another device polls, written after every earlier write of this -// stream's kernels left this device (peer writes, then the flag on the same -// path, so the reader sees the data once it sees the flag). -__global__ void fast_reduce_flag_sys_kernel(uint32_t * flag, uint32_t value) { - __threadfence_system(); - __hip_atomic_store(flag, value, __ATOMIC_RELEASE, __HIP_MEMORY_SCOPE_SYSTEM); -} - -void fast_reduce_launch_flag_sys(uint32_t * flag, uint32_t value, hipStream_t stream) { - hipLaunchKernelGGL(fast_reduce_flag_sys_kernel, dim3(1), dim3(1), 0, stream, flag, value); -} - void fast_reduce_launch_pack_pub(const float * src, void * dst, int n, int bf16, uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream) { int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; From 4e6cdcb4ae450600f38f18157aaaa264ae3a405f Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:40:52 +0200 Subject: [PATCH 12/40] fix(ds4): ranks agree on fast reduce and the drafter; unaligned slices fail the load - Every rank keeps the fast reduce only if all ranks brought it up: one rank on RCCL and the other on the raw-RDMA path desynced the collectives. - After the comm is up the ranks compare drafters (none, whole, split). Rank 0 decides whether a request drafts, so without a drafter there every rank decodes AR; otherwise every rank must draft the way rank 0 does, and a mismatch fails the load with a message naming both sides. - A cluster slice the quant blocks cannot follow fails the load. Keeping that one tensor whole left it mismatched with its sliced partner. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/deepseek4/deepseek4_backend.cpp | 65 +++++++++++++++++----- server/src/deepseek4/deepseek4_backend.h | 4 +- server/src/deepseek4/deepseek4_loader.cpp | 14 ++++- 3 files changed, 68 insertions(+), 15 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 00fa023c37..2629758b06 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1199,9 +1199,15 @@ bool DeepSeek4Backend::set_cluster(const cluster::ClusterConfig * cfg, 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) { - cluster::FastReduce::Config fc; fc.rank = cfg->rank; fc.size = cfg->size; fc.hca = cfg->ib_hca.substr(0, cfg->ib_hca.find(':')); @@ -1212,20 +1218,28 @@ bool DeepSeek4Backend::set_cluster(const cluster::ClusterConfig * cfg, 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)); - auto fast = std::make_unique(); - std::string ferr; - if (fast->init(fc, &ferr)) { - 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 { - std::fprintf(stderr, "[deepseek4-cluster] rank %d fast reduce unavailable, keeping " - "RCCL: %s\n", cfg->rank, ferr.c_str()); - } + 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()) return false; + if (comm && (!cluster_agree_prefill_chunk() || !cluster_agree_drafter())) return false; return true; } if (backend_ && !w_.layers.empty()) { @@ -1240,6 +1254,31 @@ bool DeepSeek4Backend::set_cluster(const cluster::ClusterConfig * cfg, 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() { diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index f03966d3e3..8b77b508c2 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -157,8 +157,10 @@ class DeepSeek4Backend : public ModelBackend, public IClusterParticipant { cluster::Ds4ClusterHooks * hooks_ = nullptr; bool cluster_active() const; bool cluster_worker() const; - // Loads, on load_model()'s path, the size every rank must agree on. + // 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_; diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index f3c0fe7559..cf9dd2b656 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -2293,6 +2293,9 @@ bool load_deepseek4_gguf_partial(const std::string & path, 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; @@ -2338,7 +2341,8 @@ bool load_deepseek4_gguf_partial(const std::string & path, } 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)) { - s.kind = DS4TensorSlice::None; // not block aligned: keep the full tensor + if (slice_error.empty()) slice_error = name; + s.kind = DS4TensorSlice::None; } return s; }; @@ -2394,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; From 8108b66187befa0c224aa4e736664a21f6d5bdd0 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:40:52 +0200 Subject: [PATCH 13/40] refactor(ds4): drop dead cluster helpers and a prefill opt-in that measured slower - ds4_cluster_checksum and ds4_cluster_env_prefill_single_token had no callers. - LUCE_CLUSTER_PREFILL_ATTN_SPLIT forced the prefill head split on unsliced weights and measured slower at 2K and 30K; the split now follows the sliced weights alone, as it did by default. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/deepseek4/deepseek4_cluster.cpp | 35 ---------------------- server/src/deepseek4/deepseek4_cluster.h | 16 ---------- server/src/deepseek4/deepseek4_graph.cpp | 10 +++---- 3 files changed, 4 insertions(+), 57 deletions(-) diff --git a/server/src/deepseek4/deepseek4_cluster.cpp b/server/src/deepseek4/deepseek4_cluster.cpp index 7894bd9e7c..01b70c1272 100644 --- a/server/src/deepseek4/deepseek4_cluster.cpp +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -462,21 +462,6 @@ bool ds4_cluster_init_experts(const std::string & model_path, return true; } -// ─── Diagnostics ──────────────────────────────────────────────────────── - -void ds4_cluster_checksum(const float * data, size_t n, double * sum, double * sum_abs) { - double s = 0.0; - double a = 0.0; - if (data) { - for (size_t i = 0; i < n; ++i) { - s += (double) data[i]; - a += (double) (data[i] < 0.0f ? -data[i] : data[i]); - } - } - if (sum) *sum = s; - if (sum_abs) *sum_abs = a; -} - bool ds4_cluster_ingraph_allreduce_enabled() { static const bool enabled = [] { const char * v = std::getenv("LUCE_CLUSTER_NO_INGRAPH_ALLREDUCE"); @@ -499,18 +484,6 @@ bool ds4_cluster_prefill_pipeline_enabled() { return enabled; } -bool ds4_cluster_prefill_attention_split(int context_end) { - static const bool enabled = [] { - const char * v = std::getenv("LUCE_CLUSTER_PREFILL_ATTN_SPLIT"); - return v && *v && std::strcmp(v, "0") != 0; - }(); - static const int min_rows = [] { - const char * v = std::getenv("LUCE_CLUSTER_PREFILL_ATTN_SPLIT_MIN"); - return v && *v ? std::max(0, std::atoi(v)) : 0; - }(); - return enabled && context_end >= min_rows; -} - bool ds4_cluster_prefill_device_join_enabled() { static const bool enabled = [] { const char * v = std::getenv("LUCE_CLUSTER_PREFILL_DEVICE_JOIN"); @@ -559,14 +532,6 @@ bool ds4_cluster_fused_graph_available(const Ds4ClusterRuntime * rt) { return true; } -bool ds4_cluster_env_prefill_single_token() { - static const bool enabled = [] { - const char * v = std::getenv("LUCE_CLUSTER_PREFILL_SINGLE_TOKEN"); - return v && v[0] && std::strcmp(v, "0") != 0; - }(); - return enabled; -} - // ─── Route masking ────────────────────────────────────────────────────── void ds4_cluster_mask_routes(const Ds4ClusterRuntime & rt, diff --git a/server/src/deepseek4/deepseek4_cluster.h b/server/src/deepseek4/deepseek4_cluster.h index bb9d6eb6bc..4b5aab2192 100644 --- a/server/src/deepseek4/deepseek4_cluster.h +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -157,14 +157,6 @@ void ds4_cluster_mask_routes(const Ds4ClusterRuntime & rt, int route_width, int n_tokens); -// Diagnostics: sum and sum of |x| over n floats (LUCE_CLUSTER_TRACE lines). -void ds4_cluster_checksum(const float * data, size_t n, double * sum, double * sum_abs); - -// LUCE_CLUSTER_PREFILL_SINGLE_TOKEN=1: run cluster prefill token by token -// (n_tokens == 1 chunks through the decode path) to bisect multi-token -// prefill numerics. Cached after the first call. -bool ds4_cluster_env_prefill_single_token(); - // All-reduce a device-resident F32 partial [n_embd, n_tokens] in place on the // backend stream. Stays stream-ordered (no host sync) unless rt.trace is set, // in which case it waits with cfg->timeout_ms and logs the layer. Updates @@ -224,14 +216,6 @@ bool ds4_cluster_prefill_device_join_enabled(); // hot budget the agreed bands shrink to 512 rows and prefill gets slower. bool ds4_cluster_prefill_pipeline_enabled(); -// Attention head split in the uncached (prefill) graphs: true when a band -// whose context ends at `context_end` attends with this rank's heads and -// all-reduces the partial (opt-in, LUCE_CLUSTER_PREFILL_ATTN_SPLIT=1: measured -// slower at 2K and even at 30K, where the head-independent indexer dominates; LUCE_CLUSTER_PREFILL_ATTN_SPLIT_MIN sets the context -// it starts at, default 0). A pure function of the band, so every rank -// agrees. -bool ds4_cluster_prefill_attention_split(int context_end); - // 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 diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 1a82ec5441..a98b2952e9 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -11798,14 +11798,12 @@ bool deepseek4_step_layer_range( 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 attends with - // its own output groups and one all-reduce sums the partials. - // Every rank takes the same decision (from the band alone). + // (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 || - ds4_cluster_prefill_attention_split(kv_start + n_tokens)); + 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, From 3a89de1c271c834b816de14f41ef5abefe482bc1 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:40:52 +0200 Subject: [PATCH 14/40] fix(ggml-cuda): the DS4 HC boundary barrier survives CUDA graph replay The fused boundary kernel synchronized its blocks on a running device count against a target the host computed at each launch. A CUDA graph replays the launch with the target it captured while the count keeps growing, so from the second replay on the barrier let every block through. The barrier now keeps its whole state on the device: the last block to arrive resets the count and advances a generation the others wait on. HIP builds, which do not capture graphs, run the same barrier. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../llama.cpp/ggml/src/ggml-cuda/common.cuh | 6 +- .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu | 71 ++++++++++--------- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 3 +- 3 files changed, 40 insertions(+), 40 deletions(-) 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 2693e050d9..8f7bd3e93f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh @@ -1495,10 +1495,8 @@ struct ggml_backend_cuda_context { std::vector luce_q8_memo; // DS4 HC boundary launches (ds4-hc.cu) synchronize their blocks on this - // device counter; launches on this context's stream run in order, so the - // running count of blocks launched is each launch's target. - unsigned long long * ds4_hc_barrier = nullptr; - unsigned long long ds4_hc_barrier_blocks = 0; + // 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; 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 2ee53c53b7..e5368c6f21 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 @@ -692,6 +692,34 @@ static __device__ __forceinline__ float ds4_hc_f32(float x) { 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, @@ -705,7 +733,7 @@ static __global__ void __launch_bounds__(1024) ds4_hc_boundary_kernel( 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 long long * barrier, const unsigned long long barrier_target, + 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; @@ -844,18 +872,7 @@ static __global__ void __launch_bounds__(1024) ds4_hc_boundary_kernel( const float mean_out = tmp_out / n_embd; const float scale_out = rsqrtf(mean_out + eps_out); - __syncthreads(); - if (tid == 0) { - __threadfence(); - atomicAdd(barrier, 1ull); - while (atomicAdd(barrier, 0ull) < barrier_target) { -#if defined(GGML_USE_HIP) - __builtin_amdgcn_s_sleep(1); -#endif - } - __threadfence(); - } - __syncthreads(); + 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; @@ -895,7 +912,7 @@ static __global__ void __launch_bounds__(1024) ds4_hc_boundary_split_kernel( 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 long long * barrier, const unsigned long long barrier_target, + 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; @@ -1058,18 +1075,7 @@ static __global__ void __launch_bounds__(1024) ds4_hc_boundary_split_kernel( const float mean_out = tmp_out / n_embd; const float scale_out = rsqrtf(mean_out + eps_out); - __syncthreads(); - if (tid == 0) { - __threadfence(); - atomicAdd(barrier, 1ull); - while (atomicAdd(barrier, 0ull) < barrier_target) { -#if defined(GGML_USE_HIP) - __builtin_amdgcn_s_sleep(1); -#endif - } - __threadfence(); - } - __syncthreads(); + 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; @@ -1116,15 +1122,12 @@ void ggml_cuda_ds4_hc_boundary( 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); - // Launches on this context's stream run in order, so a running count of - // the blocks they launched is each launch's barrier target. + // 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, sizeof(unsigned long long))); - CUDA_CHECK(cudaMemsetAsync(ctx.ds4_hc_barrier, 0, sizeof(unsigned long long), ctx.stream())); - ctx.ds4_hc_barrier_blocks = 0; + 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())); } - ctx.ds4_hc_barrier_blocks += (unsigned long long) n_tokens; 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 @@ -1148,7 +1151,7 @@ void ggml_cuda_ds4_hc_boundary( 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, ctx.ds4_hc_barrier_blocks, \ + 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) { @@ -1165,7 +1168,7 @@ void ggml_cuda_ds4_hc_boundary( 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, ctx.ds4_hc_barrier_blocks, \ + 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; 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 029d5b3d92..8bec1728b0 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 @@ -5696,8 +5696,7 @@ static bool ggml_cuda_try_rope_q8(ggml_backend_cuda_context & ctx, ggml_cgraph * static int ggml_cuda_try_ds4_hc_boundary(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { if (!ggml_cuda_ds4_hc_boundary_enabled()) return -1; const int n = cgraph->n_nodes; - // The next node at or after j that launches something; the skipped views - // are recorded so their single reader can be checked. + // 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; From ab496e10d12acf95323369d6ee8b42a2da8cba4b Mon Sep 17 00:00:00 2001 From: mrciffa Date: Fri, 9 Oct 2026 20:25:34 +0200 Subject: [PATCH 15/40] fix(ggml-cuda): a single-token decode no longer aborts the RDNA4 gate/up pair The generic gate/GLU fusion forms a one-column gate/up pair with SwiGLU-DS4 (a decode step without a drafter, or the last tokens of a reply), and the RDNA4 Q8_0 dispatch took it to the pair launcher, which had no one-column case: "unreachable q8_0 rdna4 glu width" on a two-box rank. One column is now a case of the same launcher, so a decode step runs the kernel a verify's columns run and the two stay batch-invariant. test_ds4_fused_ops_cuda checks the pair at one column too, bit for bit against the products kept apart. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 5 ++++- server/test/unit/test_ds4_fused_ops_cuda.cpp | 6 ++++-- 2 files changed, 8 insertions(+), 3 deletions(-) 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 9634dd505f..9dfc5e9d2b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -2884,7 +2884,9 @@ 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), 2..8 columns, no column range. +// 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, @@ -2900,6 +2902,7 @@ static void mul_mat_vec_q8_0_rdna4_launch( 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); diff --git a/server/test/unit/test_ds4_fused_ops_cuda.cpp b/server/test/unit/test_ds4_fused_ops_cuda.cpp index f5f99272cf..aa537e1ff9 100644 --- a/server/test/unit/test_ds4_fused_ops_cuda.cpp +++ b/server/test/unit/test_ds4_fused_ops_cuda.cpp @@ -478,7 +478,8 @@ bool run_glu_q8(ggml_backend_t backend, int n_ff, int n_used, int tokens, uint32 // 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). +// 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(); @@ -925,7 +926,8 @@ int main() { { // 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 pair_cases[][5] = {{1024, 64, 32, 0, 5}, {1024, 64, 32, 32, 4}, {5120, 96, 48, 48, 2}, {1024, 64, 32, 32, 8}}; + 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; } From 5306ca579036c2d189c9f45560ffba40fcd783f5 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Fri, 9 Oct 2026 20:29:19 +0200 Subject: [PATCH 16/40] perf(ds41): two-box prefill pipelines its bands when they are long enough A cluster rank kept prefill bands whole unless LUCE_CLUSTER_PREFILL_PIPELINE was set, because at 512-row bands the pipeline lost. With bands of 1536 rows or more, one part's cold FFN running beside the next part's attention wins: two boxes, a 2K prompt in one band, 649-669 -> 710-727 tok/s (2 runs each). The pipeline is now on whenever every band of the pass is that long; the switch still forces it either way. Every rank sees the same bands, so the decision needs no extra agreement. Each prefill chunk reads all of a rank's experts once, so the rank with the least free memory sets the cluster's chunk for every prompt. Rank 0 now names that rank when it holds the chunk below the pipeline's length (a 2K prompt then runs as two chunks: 558-598 tok/s): a smaller expert budget there leaves room for longer chunks. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/deepseek4/deepseek4_backend.cpp | 10 ++++++++++ server/src/deepseek4/deepseek4_cluster.cpp | 11 +++++++---- server/src/deepseek4/deepseek4_cluster.h | 12 ++++++++---- server/src/deepseek4/deepseek4_graph.cpp | 5 +++-- 4 files changed, 28 insertions(+), 10 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 2dcd39a2d9..225f430188 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1302,6 +1302,16 @@ bool DeepSeek4Backend::cluster_agree_prefill_chunk() { "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 diff --git a/server/src/deepseek4/deepseek4_cluster.cpp b/server/src/deepseek4/deepseek4_cluster.cpp index 01b70c1272..bbd375b23c 100644 --- a/server/src/deepseek4/deepseek4_cluster.cpp +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -476,12 +476,15 @@ bool ds4_cluster_ingraph_allreduce_enabled() { return enabled; } -bool ds4_cluster_prefill_pipeline_enabled() { - static const bool enabled = [] { +bool ds4_cluster_prefill_pipeline_enabled(const std::vector & bands) { + if (!ds4_cluster_prefill_device_join_enabled() || bands.empty()) return false; + static const int forced = [] { // -1 unset, 0 off, 1 on const char * v = std::getenv("LUCE_CLUSTER_PREFILL_PIPELINE"); - return ds4_cluster_prefill_device_join_enabled() && v && *v && std::strcmp(v, "0") != 0; + return v && *v ? (std::strcmp(v, "0") != 0 ? 1 : 0) : -1; }(); - return enabled; + if (forced >= 0) return forced == 1; + return std::all_of(bands.begin(), bands.end(), + [](int rows) { return rows >= kDs4ClusterPipelineMinBand; }); } bool ds4_cluster_prefill_device_join_enabled() { diff --git a/server/src/deepseek4/deepseek4_cluster.h b/server/src/deepseek4/deepseek4_cluster.h index 4b5aab2192..3243925d95 100644 --- a/server/src/deepseek4/deepseek4_cluster.h +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -211,10 +211,14 @@ bool ds4_cluster_attention_parallel_enabled(); // LUCE_CLUSTER_PREFILL_DEVICE_JOIN (default on): cluster prefill keeps the // device-resident hot+cold join and reduces in the HC-post graph. bool ds4_cluster_prefill_device_join_enabled(); -// LUCE_CLUSTER_PREFILL_PIPELINE=1 (default off): band-pipelined prefill on a -// cluster rank, band count agreed across ranks. Off by default: at the 14 GB -// hot budget the agreed bands shrink to 512 rows and prefill gets slower. -bool ds4_cluster_prefill_pipeline_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 (two boxes, 2K prompt in one band: 649-669 +// -> 710-727 tok/s); smaller bands split into parts too short to pay for it +// (512-row bands measured slower). 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 diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index a98b2952e9..348f638e44 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -12759,13 +12759,13 @@ bool deepseek4_prefill_layer_major( // 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(); + 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; bool pipelined = pipeline_requested && cluster_pipeline && device_residual && !cache.pipeline_off; - if (cluster_multi && pipeline_requested && ds4_cluster_prefill_pipeline_enabled()) { + 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)) { @@ -12781,6 +12781,7 @@ bool deepseek4_prefill_layer_major( 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()); From 11a92e027140da723e9c12731bf9c51cfceaafd4 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 12:35:15 +0200 Subject: [PATCH 17/40] perf(ds41): two-box prefill runs the R9700's experts on the grouped kernels too The ds41-gorgon profile sends only the cold owner's experts through the grouped MUL_MAT_ID prefill graph (LUCE_DS4_GROUPED_HYBRID_PREFILL=1); the R9700's hot experts took the per-expert path. On a cluster rank the R9700 does about 70% of the routed work in prefill, so cluster mode now sets 2 (both owners). Two boxes, 2K prompt in one band: the hot owner's prefill time 776 -> 582-599 ms, prefill 688 -> 711-731 tok/s; decode unchanged (code 85.5-88.1, 2K 50.8-51.5). The 2K text differs at near-ties, as any change of the prefill kernels does. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/server/launch_profiles.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/server/src/server/launch_profiles.h b/server/src/server/launch_profiles.h index 16558c885c..64426170b5 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -290,10 +290,11 @@ inline const std::vector & cluster_launch_env() { {"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, the ROCmFP2 prefetch loop and one peer-copy kernel per - // layer. + // 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"}, From ffc9649c346b07765657c2651cd3157213349d98 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 13:15:12 +0200 Subject: [PATCH 18/40] perf(ggml-cuda): ROCmFP2 expert prefill on the RDNA4 matrix cores The F16 WMMA routed-expert GEMM for large mul_mat_id batches (ROCmFP2/FP3 MIX and plain ROCmFP2) ran on RDNA3.5 only; an R9700 (gfx1201) took MMQ. gfx12 has the same 16x16x16 F16 WMMA with another fragment layout (a lane holds 8 of the K values and C rows 8*hi + e instead of all 16 and 2*e + hi), so the kernel now builds for both and picks the layout per arch. Two boxes, DS4.1, 2K prompt: the R9700's expert time in prefill 600 -> 463 ms, prefill 715-733 -> 744-756 tok/s; decode unchanged. The kernel decodes the same F16 weights, so results stay within F16 rounding of MMQ: test_rocmfp_mix_gateup_glu passes on gfx1201 through path=mix_wmma, and test_deepseek4_mmid_grouped_cuda on both GPUs. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../ggml/src/ggml-cuda/mix-wmma-moe.cu | 53 +++++++++++++------ .../test/unit/test_rocmfp_mix_gateup_glu.cpp | 2 +- 2 files changed, 38 insertions(+), 17 deletions(-) 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 fdd3947ab5..4b08b101c5 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/test/unit/test_rocmfp_mix_gateup_glu.cpp b/server/test/unit/test_rocmfp_mix_gateup_glu.cpp index 370e31f369..89758fac97 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 From b3871251772d60ffcb692ea5d60ef35d04205a6a Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 13:15:12 +0200 Subject: [PATCH 19/40] perf(cluster): an 8-bit prefill exchange (opt-in) LUCE_CLUSTER_PREFILL_Q8=1 sends prefill-sized partials as blocks of 32 int8 values with one f32 scale, all-gathered (1.125 bytes a value against the bf16 all-reduce's 2 per direction), and every rank sums the ranks' dequantized values in rank order, so the result is the same on all ranks. On a link held to about 1.8 GB/s by the NIC's x2 slot, one exchange of a 1K-row band took 5.95 ms in bf16 and 4.36 ms in 8 bits (waits for the other rank included). Two boxes, 2K prompt: 744-756 -> 789-807 tok/s. Off by default: it rounds every prefill partial to 8 bits, so outputs change at near-ties; it stays opt-in until a quality check against bf16. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/cluster_comm.cpp | 34 ++++++++++++++ server/src/cluster/cluster_comm.h | 20 +++++++++ server/src/cluster/cluster_comm_kernels.cu | 52 +++++++++++++++++++++- server/src/deepseek4/deepseek4_cluster.cpp | 21 +++++++++ 4 files changed, 125 insertions(+), 2 deletions(-) diff --git a/server/src/cluster/cluster_comm.cpp b/server/src/cluster/cluster_comm.cpp index 3458e59e90..ee7e846d3d 100644 --- a/server/src/cluster/cluster_comm.cpp +++ b/server/src/cluster/cluster_comm.cpp @@ -177,6 +177,9 @@ static_assert(sizeof(ncclUniqueId) == kRcclUniqueIdSize, // 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 { @@ -332,6 +335,37 @@ class RcclComm final : public IClusterComm { 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; diff --git a/server/src/cluster/cluster_comm.h b/server/src/cluster/cluster_comm.h index 37095bdedd..ff6f298ab7 100644 --- a/server/src/cluster/cluster_comm.h +++ b/server/src/cluster/cluster_comm.h @@ -51,6 +51,13 @@ struct ClusterCommStats { // 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; @@ -72,6 +79,19 @@ class IClusterComm { 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; diff --git a/server/src/cluster/cluster_comm_kernels.cu b/server/src/cluster/cluster_comm_kernels.cu index abd8722408..f7eae31fbe 100644 --- a/server/src/cluster/cluster_comm_kernels.cu +++ b/server/src/cluster/cluster_comm_kernels.cu @@ -2,8 +2,9 @@ // // 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. +// 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 @@ -39,6 +40,36 @@ __global__ void k_bf16_to_f32(const uint16_t * __restrict__ in, float * __restri 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; @@ -59,4 +90,21 @@ bool cluster_kernel_bf16_to_f32(const uint16_t * in, float * out, size_t n, void 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/deepseek4/deepseek4_cluster.cpp b/server/src/deepseek4/deepseek4_cluster.cpp index bbd375b23c..514c1912c8 100644 --- a/server/src/deepseek4/deepseek4_cluster.cpp +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -574,6 +574,27 @@ static void ds4_cluster_allreduce_graph_callback(void * user, void * data, 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 >= (size_t) 262144 && n % cluster::kQ8ExchangeBlock == 0 && rt->scratch_f32 && + rt->scratch_elems * sizeof(float) >= (size_t) rt->comm->size() * cluster::q8_exchange_slot_bytes(n)) { + trace_coll(rt->rank(), "graph-allreduce-q8", n, stream); + 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 >= (size_t) 262144 && rt->scratch_bf16 && rt->scratch_elems >= n) { trace_coll(rt->rank(), "graph-allreduce-bf16", n, stream); std::string err16; From 5e9448748aa663f39605b4f94797ae1f1db819c4 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 13:22:48 +0200 Subject: [PATCH 20/40] perf(ds41): cluster prefill exchanges 8-bit blocks by default The quality check the 8-bit exchange waited for, two boxes, harness suites, greedy: Math500 9/10 (bf16: 8/10), GSM8K 9/10 (bf16: 10/10), one answer flipping each way, 18/20 both. Small, but no sign that rounding the prefill partials to 8-bit blocks costs answers, and 2K prefill goes 744-756 -> 789-807 tok/s. LUCE_CLUSTER_PREFILL_Q8=0 returns to the bf16 all-reduce. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/server/launch_profiles.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/server/src/server/launch_profiles.h b/server/src/server/launch_profiles.h index 64426170b5..41b052910c 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -263,7 +263,9 @@ inline const std::vector & cluster_launch_env() { {"LUCE_CLUSTER_FAST_REDUCE_STALL_S", "15"}, {"LUCE_CLUSTER_HYBRID_EXCHANGE", "1"}, {"LUCE_CLUSTER_HYBRID_ONE_EXPORT", "1"}, - // Prefill-sized reductions stay on RCCL, in bf16. + // 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 From c8d12671fd71056dec629f18b321a0f3571e0d90 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 13:37:53 +0200 Subject: [PATCH 21/40] perf(ggml-cuda): a prefill band's HC mix in one pass over the residual The HC pre-mix of a prefill band is RMS_NORM of the [n_hc * n_embd, T] residual read only by a product against 24 F16 rows. It ran as a norm, an F16 copy of the normed rows and a GEMM 24 rows tall: three passes over a band's ~83 MB residual for 1 GFLOP. ggml_cuda_ds4_rms_norm_mix computes the 24 dot products and the sum of squares in one pass, four tokens a block, F32 throughout. Only for bands of 64 tokens or more: decode and verify keep their batch-invariant products. Two boxes, DS4.1, 2K prompt: the HC pre-mix 235 -> 106 ms per prefill. test_ds4_fused_ops_cuda checks it against a double-precision reference on gfx1201 and gfx1151 (the unfused path rounds its input to 16 bits on some GPUs, so it is no reference for an F32 kernel). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu | 76 +++++++++++++++++++ .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh | 8 ++ .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 30 ++++++++ server/test/unit/test_ds4_fused_ops_cuda.cpp | 70 +++++++++++++++++ 4 files changed, 184 insertions(+) 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 e5368c6f21..456f2c7d38 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)) @@ -1399,3 +1401,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 51d4e84e4c..f9b52f6516 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 @@ -24,6 +24,14 @@ void ggml_cuda_ds4_hc_boundary( 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/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 59aa4acb75..a782b7f3e4 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 @@ -5611,6 +5611,35 @@ static int ggml_cuda_try_rms_norm_rope_f16(ggml_backend_cuda_context & ctx, ggml 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 @@ -6097,6 +6126,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud 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; diff --git a/server/test/unit/test_ds4_fused_ops_cuda.cpp b/server/test/unit/test_ds4_fused_ops_cuda.cpp index aa537e1ff9..d2c5da07f5 100644 --- a/server/test/unit/test_ds4_fused_ops_cuda.cpp +++ b/server/test/unit/test_ds4_fused_ops_cuda.cpp @@ -15,6 +15,7 @@ #include "ggml-cuda.h" #include "ggml.h" +#include #include #include #include @@ -475,6 +476,71 @@ bool run_glu_q8(ggml_backend_t backend, int n_ff, int n_used, int tokens, uint32 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 @@ -926,6 +992,10 @@ int main() { { // 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) { From 60364d60502dc8dc49ba3d9aa6fa24c95d495989 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Sun, 11 Oct 2026 13:37:53 +0200 Subject: [PATCH 22/40] perf(ds41): a cluster band's Engram rows cross the link once A layer-major prefill band holds every Engram layer's rows, and rank 0 broadcast all of them at each Engram layer the band reached: twice per band on V4.1 (layers 1 and 14), about 100 MB a time for a 1K-row band. The band now remembers that its rows reached every rank. The prefill timing line counts the broadcast (engram_read) and the pipelined joins (hc_post). Two boxes, 2K prompt, with the one-pass HC mix: prefill 788-809 -> 877-902 tok/s; decode unchanged (code 86.8-88.9, 2K 53.8-53.9). Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/deepseek4/deepseek4_graph.cpp | 21 ++++++++++++++++----- server/src/deepseek4/deepseek4_internal.h | 3 +++ 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 348f638e44..d5cabad0c4 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -10706,10 +10706,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; @@ -11257,8 +11260,12 @@ bool deepseek4_step_layer_range( } std::vector & engram_keys_mut = band_engram_keys && range_has_engram ? *band_engram_keys : engram_keys_local; - if (cache.cluster_rt && range_has_engram) { + // 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)) { @@ -11266,6 +11273,8 @@ bool deepseek4_step_layer_range( 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( @@ -12805,6 +12814,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); } @@ -12830,6 +12840,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 diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 4dc41438ae..31f420c6a2 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -501,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. From 6a28b6a4ccb79278fae79f374efb5a44ba54b0e3 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:12:12 +0200 Subject: [PATCH 23/40] fix(cluster): rank 0 decodes in the mode it broadcast Rank 0 picks each request's decode mode for every rank and sends it in the request message: sampled and budgeted requests decode AR. The workers followed it, but rank 0 handed the original request to its own backend, which speculates on both now that main runs sampling and the thinking budget inside DSpark. The ranks then sat in two different decode loops. Rank 0 now decodes the same forced-AR request the workers rebuild, through one rule both sides share (request_decodes_ar). Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/cluster_head_backend.cpp | 14 ++++++++++++-- server/src/cluster/cluster_protocol.h | 8 ++++++++ server/src/cluster/cluster_worker_main.cpp | 2 +- server/test/unit/test_cluster_unit.cpp | 5 +++++ 4 files changed, 26 insertions(+), 3 deletions(-) diff --git a/server/src/cluster/cluster_head_backend.cpp b/server/src/cluster/cluster_head_backend.cpp index 2a3b67ad53..c39fbaa199 100644 --- a/server/src/cluster/cluster_head_backend.cpp +++ b/server/src/cluster/cluster_head_backend.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include namespace luce::cluster { @@ -413,10 +414,19 @@ GenerateResult ClusterHeadBackend::run_request(const GenerateRequest & req, cons 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, req, io) - : inner_->generate_impl(req, io); + 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 diff --git a/server/src/cluster/cluster_protocol.h b/server/src/cluster/cluster_protocol.h index 351fc289ce..48feee79a4 100644 --- a/server/src/cluster/cluster_protocol.h +++ b/server/src/cluster/cluster_protocol.h @@ -219,6 +219,14 @@ struct RequestMsg { 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 diff --git a/server/src/cluster/cluster_worker_main.cpp b/server/src/cluster/cluster_worker_main.cpp index 8fcb07d586..7a9240fcb0 100644 --- a/server/src/cluster/cluster_worker_main.cpp +++ b/server/src/cluster/cluster_worker_main.cpp @@ -58,7 +58,7 @@ GenerateRequest request_from_msg(const RequestMsg & m) { req.sampler.seed = m.seed; req.do_sample = m.temperature > 0.0f; req.stream = false; - req.force_ar_decode = m.force_ar || m.decode_mode == DecodeMode::Autoregressive; + 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; diff --git a/server/test/unit/test_cluster_unit.cpp b/server/test/unit/test_cluster_unit.cpp index 68ea49eaac..0229866ae9 100644 --- a/server/test/unit/test_cluster_unit.cpp +++ b/server/test/unit/test_cluster_unit.cpp @@ -382,6 +382,10 @@ TEST_CASE(ClusterUnitFixture, RequestRoundTripEmptyAndLarge) { 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. @@ -411,6 +415,7 @@ TEST_CASE(ClusterUnitFixture, RequestRoundTripEmptyAndLarge) { 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); } } From cb6305678a6487d60ba59afb8d26cf32f4b317db Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:30:47 +0200 Subject: [PATCH 24/40] refactor(ds4): drop the verify stage-batch and graph-prebuild experiments Both were off by default and measured no better: LUCE_DS4_STAGE_BATCH (the step's pinned uploads as one copy kernel) gained nothing, and LUCE_DS4_VERIFY_PREBUILD (building the next shapes while a verify runs) filled the slot cache and made the two-box verify 42 -> 82 ms. The pinned stage loses its batch mode and ds4_fused_graph_compute is main's again. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/common/pinned_stage.h | 21 --- .../src/deepseek4/deepseek4_fused_verify.inc | 121 +----------------- 2 files changed, 5 insertions(+), 137 deletions(-) diff --git a/server/src/common/pinned_stage.h b/server/src/common/pinned_stage.h index b36df9b236..5fce857b20 100644 --- a/server/src/common/pinned_stage.h +++ b/server/src/common/pinned_stage.h @@ -19,7 +19,6 @@ #include #include #include -#include namespace luce::common { @@ -59,20 +58,6 @@ class PinnedStage { void reset() { used_ = 0; } - // Batched mode: set() collects its device uploads instead of queueing - // one copy each, and the caller sends batch() as a single launch (one - // kernel reading the pinned slices; a copy engine adds a start latency - // of 0.03-1 ms to every step on a discrete GPU). end_batch() before the - // next reset; the slices stay valid until then. - struct Pending { - const void * src; - void * dst; - size_t bytes; - }; - void begin_batch() { batching_ = true; pending_.clear(); } - const std::vector & batch() const { return pending_; } - void end_batch() { batching_ = false; pending_.clear(); } - // 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) { @@ -94,10 +79,6 @@ class PinnedStage { return; } std::memcpy(slice, data, bytes); - if (batching_) { - pending_.push_back({slice, static_cast(t->data) + offset, bytes}); - return; - } ggml_backend_tensor_set_async(backend, t, slice, offset, bytes); } @@ -144,8 +125,6 @@ class PinnedStage { size_t size_ = 0; size_t used_ = 0; ggml_backend_t backend_ = nullptr; - bool batching_ = false; - std::vector pending_; }; } // namespace luce::common diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index c4e056834e..d8a4a4af8a 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -260,25 +260,15 @@ static std::vector ds4_fused_stream_jobs( // the launch. static enum ggml_status ds4_fused_graph_compute( const DeepSeek4Weights & w, DeepSeek4FusedDecodeGraph & fg, ggml_backend_t backend, - MoeHybridStorage * hybrid, const char * who, - const std::function & while_running = {}) { + MoeHybridStorage * hybrid, const char * who) { MoeStreamedMailbox * mailbox = hybrid && hybrid->expert_cache ? hybrid->expert_cache->mailbox() : nullptr; const std::vector jobs = mailbox ? ds4_fused_stream_jobs(fg) : std::vector{}; if (!jobs.empty()) mailbox->begin(jobs, ds4_fused_invalid_route()); - enum ggml_status status; - if (fg.sched && jobs.empty() && while_running) { - // Host work that needs no device result (building the next shape) - // runs while the GPUs execute this graph. - status = ggml_backend_sched_graph_compute_async(fg.sched, fg.sg.gf); - if (status == GGML_STATUS_SUCCESS) while_running(); - ggml_backend_sched_synchronize(fg.sched); - } else { - status = fg.sched - ? ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf) - : ggml_backend_graph_compute(backend, fg.sg.gf); - } + enum ggml_status status = fg.sched + ? ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf) + : ggml_backend_graph_compute(backend, fg.sg.gf); if (!jobs.empty()) { std::string error; if (!mailbox->end(&error)) { @@ -2223,36 +2213,6 @@ static int ds4_try_fused_verify_step( ((int64_t) padded << 4) | (int64_t) (b_idx + 1)); } } - // The key a step at (kv, qq) would look up: the rules above. - auto key_for = [&](int kv, int qq) { - std::vector k; - k.reserve(key.size()); - k.push_back(w.n_swa); - k.push_back(qq); - k.push_back(token_ids ? 1 : 0); - k.push_back(hybrid ? 1 : 0); - k.push_back((int64_t) hooks->capture_layer_ids->size()); - for (int layer : *hooks->capture_layer_ids) k.push_back(layer); - k.push_back(vocab_split ? 1 : 0); - const int nl = ds4_fused_attention_lanes(w, qq); - const int alq = qq / nl; - for (int il = 0; il < w.n_layer; ++il) { - const int ratio = (int) w.compress_ratios[il]; - DeepSeek4LayerCache & lc = ds4_comp_cache(cache, w, il); - for (int lane = 0; lane < nl; ++lane) { - const int lks = kv + lane * alq; - const int padded = ds4_fused_lane_padded(w, lc, il, lks + alq - 1); - int b_idx = -1; - if (ratio > 0) { - for (int t = 0; t < alq; ++t) { - if (((lks + t + 1) % ratio) == 0) { b_idx = t; break; } - } - } - k.push_back(((int64_t) padded << 4) | (int64_t) (b_idx + 1)); - } - } - return k; - }; vc.counter++; const size_t slot_limit = hybrid @@ -2355,15 +2315,6 @@ static int ds4_try_fused_verify_step( ~EngramJoin() { finish(); } } engram_join{engram_async ? vc.engram_async.get() : nullptr, telemetry}; const auto set_t0 = Ds4TimingClock::now(); - // LUCE_DS4_STAGE_BATCH=1: the step's pinned uploads leave as one copy - // kernel before the single wait below, instead of a dozen copies of - // which the larger ones wait for a copy engine to start. - static const bool stage_batch_env = [] { - const char * v = std::getenv("LUCE_DS4_STAGE_BATCH"); - return v && *v && std::strcmp(v, "0") != 0; - }(); - const bool stage_batch = stage_batch_env && pinned_io && ggml_backend_is_cuda(backend); - if (stage_batch) vc.stage.begin_batch(); ds4_fv_set(fg->inp_embed, embed, sizeof(float) * (size_t) w.n_embd * q); std::vector iv(q); std::vector lv(q); @@ -2579,16 +2530,6 @@ static int ds4_try_fused_verify_step( // GPU for the worker, not on the host. const auto st_t3 = st_now(); st_hash += st_us(st_t2, st_t3); - if (stage_batch) { - const auto & pending = vc.stage.batch(); - if (!pending.empty()) { - std::vector descs; - descs.reserve(pending.size()); - for (const auto & p : pending) descs.push_back({p.src, p.dst, p.bytes}); - ggml_backend_cuda_copy_batch_async(backend, descs.data(), (int) descs.size()); - } - vc.stage.end_batch(); - } if (pinned_io) ggml_backend_synchronize(backend); st_sync += st_us(st_t3, st_now()); if (set_timing && ++st_n % 64 == 0) { @@ -2647,60 +2588,8 @@ static int ds4_try_fused_verify_step( /*skip_property_check=*/force_verify_graph_replay, /*ds4_mix_mmv_max_tokens=*/0, /*mmvq_batch_invariant=*/n_lanes > 1); - // Experimental, off by default: on the two-box code run it filled the slot - // cache and the graphs thrashed through HIP-graph warmup (verify 42 -> 82 ms). - // LUCE_DS4_VERIFY_PREBUILD= (hybrid): while this graph runs, build up - // to n slots for the shapes the next step will most likely look up (this - // width, the next start after a full or one-short acceptance) when a free - // slot is left. A miss costs a ~15 ms rebuild on the critical path; here - // the host would only wait for the GPUs. - static const int prebuild_max = [] { - const char * v = std::getenv("LUCE_DS4_VERIFY_PREBUILD"); - return v && *v ? std::max(0, std::atoi(v)) : 0; - }(); - std::function prebuild; - if (prebuild_max > 0 && hybrid && q > 1) { - prebuild = [&]() { - int built = 0; - for (int a = q; a >= 1 && built < prebuild_max; --a) { - std::vector k = key_for(kv_start + a, q); - bool have = false; - for (size_t i = 0; i < slot_limit && !have; ++i) { - have = vc.slots[i].built() && vc.slots[i].shape_key == k; - } - if (have) continue; - size_t free_slot = slot_limit; - for (size_t i = 0; i < slot_limit; ++i) { - if (!vc.slots[i].built() && &vc.slots[i] != fg) { free_slot = i; break; } - } - if (free_slot == slot_limit) return; // never evict under a running graph - DeepSeek4FusedDecodeGraph & pf = vc.slots[free_slot]; - Ds4FusedVerifyCache::Extra & pe = vc.extra[free_slot]; - pf.destroy(vc.backend, vc.peer_backend); - pe.reset(); - pe.vocab_split = vocab_split; - const auto pb_t0 = Ds4TimingClock::now(); - if (!ds4_build_fused_verify_graph(mc, pf, pe, backend, w, &cache, - hc_weights, hc_out_weights, hash_tables, - kv_start + a, q, token_ids != nullptr, - *hooks->capture_layer_ids, hybrid, - std::move(k))) { - pf.destroy(vc.backend, vc.peer_backend); - pe.reset(); - return; - } - pf.last_use = 0; // unused prebuilt slots are the first to go - ++built; - static int reported = 0; - if (reported++ < 32) { - std::fprintf(stderr, "[ds4-fused-verify] prebuilt q=%d kv=%d in %.1f ms\n", q, - kv_start + a, ds4_elapsed_us(pb_t0, Ds4TimingClock::now()) / 1000.0); - } - } - }; - } const enum ggml_status compute_status = - ds4_fused_graph_compute(w, *fg, backend, hybrid, "ds4-fused-verify", prebuild); + 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; From 3645375f347c893fb36c722e9bf6d8a141d731e6 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:30:48 +0200 Subject: [PATCH 25/40] fix(cluster): shard the shared expert by default Each rank now computes half of the shared expert's intermediate axis and sends its partial in the routed all-reduce, the configuration the measured numbers and the documented command use. Replicate computed the whole expert on both ranks: two boxes, code 89.7-91.9 -> 93.2-97.4 tok/s, 2K prose 52.5 -> 56.3-56.5, at 44.2 -> 41.7 ms a code step. The rank0 mode, accepted by the parser but refused at load, is gone. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/cluster_config.cpp | 7 +------ server/src/cluster/cluster_config.h | 3 +-- server/src/deepseek4/deepseek4_cluster.cpp | 9 --------- server/src/deepseek4/deepseek4_cluster.h | 23 ++++++++++++---------- server/src/deepseek4/deepseek4_graph.cpp | 7 ------- server/src/server/server_main.cpp | 2 +- server/test/unit/test_cluster_unit.cpp | 2 +- 7 files changed, 17 insertions(+), 36 deletions(-) diff --git a/server/src/cluster/cluster_config.cpp b/server/src/cluster/cluster_config.cpp index e8c9aa0237..0bbae99e85 100644 --- a/server/src/cluster/cluster_config.cpp +++ b/server/src/cluster/cluster_config.cpp @@ -168,11 +168,7 @@ bool parse_shared_expert_mode(const std::string & text, SharedExpertMode & out, out = SharedExpertMode::Shard; return true; } - if (text == "rank0") { - out = SharedExpertMode::Rank0; - return true; - } - set_err(err, "--cluster-shared-expert expects replicate, shard, or rank0, got '" + text + "'"); + set_err(err, "--cluster-shared-expert expects replicate or shard, got '" + text + "'"); return false; } @@ -221,7 +217,6 @@ const char * shared_expert_mode_name(SharedExpertMode mode) { switch (mode) { case SharedExpertMode::Replicate: return "replicate"; case SharedExpertMode::Shard: return "shard"; - case SharedExpertMode::Rank0: return "rank0"; } return "unknown"; } diff --git a/server/src/cluster/cluster_config.h b/server/src/cluster/cluster_config.h index 1293e37fd2..33b90a0604 100644 --- a/server/src/cluster/cluster_config.h +++ b/server/src/cluster/cluster_config.h @@ -31,7 +31,6 @@ inline constexpr uint32_t kClusterDefaultTimeoutMs = 30000; 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 - Rank0, // only rank 0 computes it and adds it after the all-reduce }; // Wire precision of the per-layer partial-sum reduction. @@ -70,7 +69,7 @@ struct ClusterConfig { PlacementSource placement_source = PlacementSource::Uniform; std::string placement_file; // when placement_source == File int replicate_hot = 0; // --cluster-replicate-hot - SharedExpertMode shared_expert = SharedExpertMode::Replicate; + SharedExpertMode shared_expert = SharedExpertMode::Shard; // Reduction and robustness knobs. AllreduceDType allreduce_dtype = AllreduceDType::Auto; diff --git a/server/src/deepseek4/deepseek4_cluster.cpp b/server/src/deepseek4/deepseek4_cluster.cpp index 514c1912c8..d1351ebcf7 100644 --- a/server/src/deepseek4/deepseek4_cluster.cpp +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -251,12 +251,6 @@ bool ds4_cluster_build_placement(const cluster::ClusterConfig & cfg, set_err(err, "cluster config is not enabled or rank is out of range"); return false; } - if (cfg.shared_expert == cluster::SharedExpertMode::Rank0) { - set_err(err, std::string("shared expert mode '") + - cluster::shared_expert_mode_name(cfg.shared_expert) + - "' is not implemented yet; use replicate or shard"); - return false; - } const MoeHybridConfig mcfg = ds4_cluster_moe_config(w); MoeHybridRoutingStats stats; @@ -508,9 +502,6 @@ bool ds4_cluster_attention_parallel_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; - if (rt->cfg && rt->cfg->shared_expert == cluster::SharedExpertMode::Rank0) { - return false; - } for (int32_t owner : rt->placement.owner) { if (owner == cluster::kReplicated) { static bool warned = false; diff --git a/server/src/deepseek4/deepseek4_cluster.h b/server/src/deepseek4/deepseek4_cluster.h index 3243925d95..85135323ad 100644 --- a/server/src/deepseek4/deepseek4_cluster.h +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -1,12 +1,13 @@ // deepseek4_cluster.h - DeepSeek4 expert-parallel cluster runtime. // -// Every rank loads the full dense model 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 routed partial identical on every rank before the shared expert -// is added and HC-post runs. +// Every rank loads the dense model (with LUCE_CLUSTER_SLICE_WEIGHTS, only the +// attention heads and shared-expert slice it computes) 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 @@ -14,9 +15,11 @@ // call between per-layer graph computes inside // deepseek4_step_layer_range / eval_ds4_layer_range_hybrid_ffn. // -// Shared expert: SharedExpertMode::Replicate only. Every rank computes it -// locally from the same normalized input and adds it after the all-reduce; -// the routed partial is evaluated with the shared-expert tensors removed +// 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 diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index d5cabad0c4..2ed50a38eb 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -8424,13 +8424,6 @@ static bool eval_ds4_layer_range_hybrid_ffn( // Slot = token index inside this batch (identical on all ranks). const MoeLayerDesc shared_desc = desc; if (cluster_rt) { - if (cluster_rt->cfg && - cluster_rt->cfg->shared_expert == cluster::SharedExpertMode::Rank0) { - std::fprintf(stderr, - "[deepseek4-cluster] layer %d: shared expert mode %s not implemented\n", - layer, cluster::shared_expert_mode_name(cluster_rt->cfg->shared_expert)); - return false; - } 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 diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 6530a3ce90..b663b8d3d9 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -197,7 +197,7 @@ static void print_usage(const char * prog) { " --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: replicate)\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" diff --git a/server/test/unit/test_cluster_unit.cpp b/server/test/unit/test_cluster_unit.cpp index 0229866ae9..dde9254f2c 100644 --- a/server/test/unit/test_cluster_unit.cpp +++ b/server/test/unit/test_cluster_unit.cpp @@ -608,7 +608,7 @@ TEST_CASE(ClusterUnitFixture, ClusterConfigParsers) { 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); - REQUIRE(parse_shared_expert_mode("rank0", se, &err)); CHECK(se == SharedExpertMode::Rank0); + 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")); From fcfe8d59fded05e378427ced9ba08675b7ead371 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:30:48 +0200 Subject: [PATCH 26/40] fix(ds4): verify-width costs carry over per context range DSpark's width controller started every request from one set of step costs, whatever the context. A step costs more as the context grows (on two boxes q5 is 43 ms on a short prompt and 47-53 ms after a 2K one), so a request that starts from another context's costs picks its widths from the wrong prices until its own steps outweigh them. The costs now live in the shared SpecWidthCostMemory, one set per context range (under 2K, 2-8K, 8-32K, longer), as Qwen's MTP keeps them. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/deepseek4/deepseek4_dspark_spec.cpp | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index a08505ee10..994f1f7e5c 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -1214,13 +1214,14 @@ bool run_deepseek4_dspark_spec_decode( // 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, not to the text: every request starts - // from the costs the earlier ones measured (their warmup hold re-arms). - static std::mutex measured_width_costs_mu; - static std::vector measured_width_costs; + // 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(measured_width_costs_mu); - width_controller.set_measured_costs(measured_width_costs); + std::lock_guard lock(width_costs_mu); + width_costs.load(width_controller, committed); } if (timing && width_controller.enabled()) { std::fprintf(stderr, @@ -2060,8 +2061,8 @@ bool run_deepseek4_dspark_spec_decode( q_cap >= 4 ? width_steps[4] : 0, q_cap >= 5 ? width_steps[5] : 0); { - std::lock_guard lock(measured_width_costs_mu); - measured_width_costs = width_controller.measured_costs(); + 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) { From 190676b855cbe7e28271311f5d8e87731ad4e2fb Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:47:36 +0200 Subject: [PATCH 27/40] fix(cluster): one wire format and one lean rule for every fast-reduce path - The progress thread sized each RDMA write as bf16 whenever _BF16 and _ZC were set, while submit packs bf16 only on the zero-copy path (two ranks, stream flags). With _KERNEL_FLAGS=1 or more than two ranks it sent half of an f32 row. fr_bf16_wire() is now the one rule: the progress thread, submit and the hybrid combine all follow it. - _LEAN drops the slot-reuse wait and the done write on every stream-flag path. Submit kept its wait unless zero-copy was also on, so with _ZC=0 a submit landing on a slot a hybrid exchange last used waited for a done flag nothing writes. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/fast_reduce.cpp | 35 +++++++++++++++++------------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/server/src/cluster/fast_reduce.cpp b/server/src/cluster/fast_reduce.cpp index 8cb0e97349..9e57f31a4d 100644 --- a/server/src/cluster/fast_reduce.cpp +++ b/server/src/cluster/fast_reduce.cpp @@ -213,8 +213,16 @@ static bool fr_zero_copy(int size) { return zc && size == 2; } +// A row crosses the cable as bf16 exactly when the zero-copy path packs it so +// (LUCE_CLUSTER_FAST_REDUCE_BF16); the copy-engine path publishes f32. The +// progress thread sizes every RDMA write by the same rule. +static bool fr_bf16_wire(int size) { + static const bool bf16 = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16"); + return bf16 && fr_zero_copy(size); +} + // LUCE_CLUSTER_FAST_REDUCE_LEAN=1 drops the per-reduction slot-reuse waits and -// done writes on the zero-copy and hybrid paths. Neither carries data: the +// done writes on every stream-flag path. Neither carries data: 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 256 reductions earlier, and the GPU's own stream orders its @@ -703,9 +711,8 @@ bool FastReduce::init(const Config & cfg, std::string * err) { 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); - static const bool bf16_wire = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16") && - fr_env_on("LUCE_CLUSTER_FAST_REDUCE_ZC"); - sge.length = (uint32_t) ((size_t) n * (bf16_wire && !exact_row ? 2u : sizeof(float))); + const bool bf16_row = fr_bf16_wire(s.cfg.size) && !exact_row; + sge.length = (uint32_t) ((size_t) n * (bf16_row ? 2u : sizeof(float))); sge.lkey = s.mr_data->lkey; ibv_send_wr w{}; @@ -829,8 +836,7 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { // third request of a run stop rather than the first. // The previous user of this slot was seq - slots, and it leaves exactly that // value behind. Asking for anything higher is a wait nothing can satisfy. - static const bool zc_lean = fast_reduce_lean() && fr_zero_copy(s.cfg.size); - if (!kernel_flags && !zc_lean && seq > (uint64_t) s.cfg.slots) { + if (!kernel_flags && !fast_reduce_lean() && seq > (uint64_t) s.cfg.slots) { const uint32_t need = (uint32_t) (seq - (uint64_t) s.cfg.slots); if (hipStreamWaitValue32(stm, s.flag_dev + s.done_off(slot), need, hipStreamWaitValueGte, 0xffffffffu) != hipSuccess) { @@ -840,8 +846,7 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { // Zero-copy (discrete GPU): pack straight into the pinned row, publish, wait // for the peer, add its row straight out of pinned memory, mark the slot done. - static const int zc_bf16 = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16") ? 1 : 0; - const int wire_bf16 = exact ? 0 : zc_bf16; + const int wire_bf16 = !exact && fr_bf16_wire(s.cfg.size) ? 1 : 0; if (fr_zero_copy(s.cfg.size)) { float * own_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) s.cfg.rank * s.cfg.max_elems; const int peer = 1 - s.cfg.rank; @@ -858,7 +863,7 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { if (hipStreamWaitValue32(stm, s.flag_dev + s.visible_off(slot, peer), (uint32_t) seq, hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; fast_reduce_launch_add_zc(data, peer_row, (int) n, wire_bf16, stm); - if (zc_lean) return true; + if (fast_reduce_lean()) return true; return hipStreamWriteValue32(stm, s.flag_dev + s.done_off(slot), (uint32_t) seq, 0) == hipSuccess; } @@ -914,7 +919,7 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { fast_reduce_launch_add(data, scratch_row, n_peers, (int) n, s.cfg.max_elems, !kernel_flags ? nullptr : s.flag_dev + s.done_off(slot), s.progress_dev, (uint32_t) seq, stm); - if (!kernel_flags && + if (!kernel_flags && !fast_reduce_lean() && hipStreamWriteValue32(stm, s.flag_dev + s.done_off(slot), (uint32_t) seq, 0) != hipSuccess) { return false; @@ -969,7 +974,7 @@ bool FastReduce::hyb_combine(const float * data, size_t n, void * stream, uint64 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; - static const int zc_bf16 = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16") ? 1 : 0; + const int wire_bf16 = fr_bf16_wire(s.cfg.size) ? 1 : 0; 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; @@ -984,19 +989,19 @@ bool FastReduce::hyb_combine(const float * data, size_t n, void * stream, uint64 uint32_t * ctr_fin = s.pub_counter(stm, slot, 3); if (!ctr_pub || !ctr_fin) return false; fast_reduce_launch_hyb_combine_pub(data, s.hyb_part(0, slot), one_export ? nullptr : s.hyb_part(1, slot), - own_row, (int) n, zc_bf16, + own_row, (int) n, wire_bf16, 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_pub(own_row, peer_row, s.hyb_part(2, slot), (int) n, zc_bf16, + fast_reduce_launch_hyb_final_pub(own_row, peer_row, s.hyb_part(2, slot), (int) n, wire_bf16, s.flag_dev + s.hybf_off(slot), (uint32_t) seq, ctr_fin, stm); return true; } - fast_reduce_launch_hyb_combine(data, s.hyb_part(0, slot), one_export ? nullptr : s.hyb_part(1, slot), own_row, (int) n, zc_bf16, stm); + fast_reduce_launch_hyb_combine(data, s.hyb_part(0, slot), one_export ? nullptr : s.hyb_part(1, slot), own_row, (int) n, wire_bf16, stm); if (hipStreamWriteValue32(stm, s.flag_dev + s.pub_off(slot), (uint32_t) seq, 0) != hipSuccess) return false; 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(2, slot), (int) n, zc_bf16, stm); + fast_reduce_launch_hyb_final(own_row, peer_row, s.hyb_part(2, slot), (int) n, wire_bf16, stm); return hipStreamWriteValue32(stm, s.flag_dev + s.hybf_off(slot), (uint32_t) seq, 0) == hipSuccess; } From e28cf0f656cc90e2c12e0a0e9bd3fcf9a58ea964 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:53:21 +0200 Subject: [PATCH 28/40] fix(cluster): a failed request answers before rank 0 exits - A worker lost mid-request, or a failed request/RequestEnd/report/backend op broadcast, called fail_cluster, which exits on the spot: the client got a reset connection instead of the error. These request-thread paths now use request_fatal_exit, which exits after a grace period, as the watchdog already did. fail_cluster stays for start-up failures. - The head hook no longer samples when the caller passes no token: every caller samples in its decode loop with the request's sampler and RNG, and the unused fallback seeded its RNG once per process. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/cluster_decision_hooks.cpp | 36 ++++--------------- server/src/cluster/cluster_decision_hooks.h | 16 ++------- server/src/cluster/cluster_head_backend.cpp | 17 +++++---- server/src/cluster/cluster_head_backend.h | 3 +- server/test/unit/test_cluster_unit.cpp | 16 ++++----- 5 files changed, 28 insertions(+), 60 deletions(-) diff --git a/server/src/cluster/cluster_decision_hooks.cpp b/server/src/cluster/cluster_decision_hooks.cpp index 6bc4c55168..73056097b1 100644 --- a/server/src/cluster/cluster_decision_hooks.cpp +++ b/server/src/cluster/cluster_decision_hooks.cpp @@ -5,25 +5,11 @@ #include #include -#include namespace luce::cluster { namespace { -int32_t greedy_argmax(const float * logits, int n_vocab) { - if (!logits || n_vocab <= 0) return -1; - int32_t best = 0; - float best_v = logits[0]; - for (int i = 1; i < n_vocab; i++) { - if (logits[i] > best_v) { - best_v = logits[i]; - best = i; - } - } - return best; -} - void set_err(std::string * err, const std::string & msg) { if (err) *err = msg; } @@ -51,24 +37,14 @@ void HeadHooks::reset_counters() { first_mismatch_step_ = 0; } -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*/, +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) { - // Caller did not sample host-side. Only rank 0 owns an RNG. - if (sampler_) { - token = sampler_(logits, n_vocab, cfg); - } else if (cfg.temp > 0.0f) { - static thread_local std::mt19937_64 rng{cfg.seed != 0 ? cfg.seed : 0x5eedULL}; - std::vector no_history; - token = common::sample_logits(logits, n_vocab, cfg, no_history, rng); - } else { - token = greedy_argmax(logits, n_vocab); - } - if (token < 0) { - set_err(err, "head: sampling produced no token"); - return false; - } + set_err(err, "head: no token to broadcast"); + return false; } DecisionMsg msg; diff --git a/server/src/cluster/cluster_decision_hooks.h b/server/src/cluster/cluster_decision_hooks.h index 5639e71e92..70fb9da371 100644 --- a/server/src/cluster/cluster_decision_hooks.h +++ b/server/src/cluster/cluster_decision_hooks.h @@ -34,7 +34,6 @@ #include "common/sampler.h" #include -#include #include #include @@ -79,11 +78,8 @@ struct Ds4ClusterHooks { // AR decode, one call per step on every rank. // // Head / local: `token` and `flags` are INPUTS. The caller sampled - // `token` from `logits` with `cfg` (including any budget-hook override) - // and computed `flags` with make_decision_flags. If the caller passes - // token < 0 the head hook samples itself (greedy argmax, or `cfg` via - // common::sample_logits with an internal RNG when cfg.temp > 0); this - // path exists for callers that do not sample host-side. + // `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). @@ -165,14 +161,9 @@ struct HeadHooksConfig { class HeadHooks final : public Ds4ClusterHooks { public: - // `control` must outlive the hooks. Optional sampler used only when the - // caller passes token < 0. + // `control` must outlive the hooks. HeadHooks(ClusterHeadControl & control, const HeadHooksConfig & cfg); - using Sampler = std::function; - void set_sampler(Sampler s) { sampler_ = std::move(s); } - 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; @@ -195,7 +186,6 @@ class HeadHooks final : public Ds4ClusterHooks { private: ClusterHeadControl & control_; HeadHooksConfig cfg_; - Sampler sampler_; int first_mismatch_rank_ = -1; uint32_t first_mismatch_step_ = 0; }; diff --git a/server/src/cluster/cluster_head_backend.cpp b/server/src/cluster/cluster_head_backend.cpp index c39fbaa199..209d6fcdc6 100644 --- a/server/src/cluster/cluster_head_backend.cpp +++ b/server/src/cluster/cluster_head_backend.cpp @@ -407,9 +407,11 @@ GenerateResult ClusterHeadBackend::run_request(const GenerateRequest & req, cons 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)) { - fail_cluster("request broadcast: " + err, kAbortCodeControl, request_id); + request_fatal_exit("request broadcast: " + err, kAbortCodeControl, request_id); result.fail(GenerateErrorCode::DecodeFailed, "cluster: request broadcast failed: " + err); return result; } @@ -437,9 +439,9 @@ GenerateResult ClusterHeadBackend::run_request(const GenerateRequest & req, cons // 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)))) { - fail_cluster("request " + std::to_string(request_id) + " failed: " + - std::string(result.error_detail()) + (err.empty() ? "" : " / " + err), - kAbortCodeRequest, request_id); + request_fatal_exit("request " + std::to_string(request_id) + " failed: " + + std::string(result.error_detail()) + (err.empty() ? "" : " / " + err), + kAbortCodeRequest, request_id); return result; } @@ -448,12 +450,13 @@ GenerateResult ClusterHeadBackend::run_request(const GenerateRequest & req, cons end.steps = (uint32_t) result.tokens.size(); end.status = result.ok() ? 0 : (int32_t) result.error->code; if (!control_.broadcast(make_frame(end), &err)) { - fail_cluster("RequestEnd broadcast: " + err, kAbortCodeControl, request_id); + 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)) { - fail_cluster("RequestReport gather: " + last_report_.error, kAbortCodeControl, request_id); + request_fatal_exit("RequestReport gather: " + last_report_.error, kAbortCodeControl, + request_id); if (result.ok()) result.fail(GenerateErrorCode::DecodeFailed, "cluster: report gather failed"); return result; } @@ -531,7 +534,7 @@ bool ClusterHeadBackend::broadcast_op(BackendOpKind kind, std::vector a op.args = std::move(args); std::string err; if (!hooks_->backend_op(op, &err)) { - fail_cluster("BackendOp broadcast: " + err, kAbortCodeControl, 0); + request_fatal_exit("BackendOp broadcast: " + err, kAbortCodeControl, 0); return false; } return true; diff --git a/server/src/cluster/cluster_head_backend.h b/server/src/cluster/cluster_head_backend.h index 6b5672a0c1..79f728e265 100644 --- a/server/src/cluster/cluster_head_backend.h +++ b/server/src/cluster/cluster_head_backend.h @@ -161,7 +161,8 @@ class ClusterHeadBackend final : public common::ModelBackend { // 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. + // 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; diff --git a/server/test/unit/test_cluster_unit.cpp b/server/test/unit/test_cluster_unit.cpp index dde9254f2c..a8e1bf8296 100644 --- a/server/test/unit/test_cluster_unit.cpp +++ b/server/test/unit/test_cluster_unit.cpp @@ -1088,15 +1088,13 @@ TEST_CASE(ClusterUnitFixture, HeadAndWorkerHooksOverLoopback) { CHECK(worker.backend_op(got, &w_err)); CHECK_EQUAL(head.counters().backend_ops, 1u); - // Head samples itself when the caller passes token < 0 (greedy). - int32_t sampled = -1; uint8_t sflags = 0; - std::thread wt2([&] { - int32_t t = -1; uint8_t fl = 0; std::string e; - worker.decide_next_token(req, 1, nullptr, 0, scfg, false, t, fl, &e); - }); - REQUIRE(head.decide_next_token(req, 1, logits, 4, scfg, false, sampled, sflags, &h_err)); - CHECK_EQUAL(sampled, 1); - wt2.join(); + // 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") { From 9c537619c6e91b3a29dca4f116d24a67e3e67b3a Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:53:21 +0200 Subject: [PATCH 29/40] fix(ds4): async Engram keys only with the hybrid graph's broadcast - The verify step read the Engram keys on a worker whenever the cluster was attached, but only a graph built for the hybrid cluster path has the broadcast whose callback uploads them. A cluster step without the hybrid graph would have read stale keys; it now reads them in line. - The head-split wo_b column view is taken only from unsliced storage, like wo_a's group view (a sliced rank's head range starts at 0, so the view was a no-op there). - The hybrid exchange counters say what they assume: every rank runs every exchange, and a rank-local step failure ends both ranks. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/deepseek4/deepseek4_cluster.h | 13 +++++++++---- server/src/deepseek4/deepseek4_fused_verify.inc | 5 +++-- server/src/deepseek4/deepseek4_graph.cpp | 5 +++-- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/server/src/deepseek4/deepseek4_cluster.h b/server/src/deepseek4/deepseek4_cluster.h index 85135323ad..3de6c84b6a 100644 --- a/server/src/deepseek4/deepseek4_cluster.h +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -61,7 +61,11 @@ struct Ds4ClusterRuntime { // 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. + // 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_export0 = 0, hyb_export1 = 0, hyb_combine = 0, hyb_import = 0; cluster::ClusterExpertPlacement placement; // Rank-local hot set (owned + replicated experts) for the hybrid storage. @@ -81,11 +85,12 @@ struct Ds4ClusterRuntime { // 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. + // 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): this rank's - // slice of the intermediate axis. count 0 = the shared expert is + // 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; diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index d8a4a4af8a..512bd95d81 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -2277,8 +2277,9 @@ static int ds4_try_fused_verify_step( 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). - Ds4ClusterRuntime * engram_rt = cache.cluster_rt && cache.cluster_rt->size() > 1 + // 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) { diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 2ed50a38eb..3a96ab4a41 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2338,8 +2338,9 @@ static ggml_tensor * build_mla_output_projection( } 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. - ggml_tensor * out_b_w = head_split + // 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)) From e5f6405bf5379bf082a2fd610b57299221f334cc Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 11 Oct 2026 23:53:21 +0200 Subject: [PATCH 30/40] fix(ggml-cuda): masked routes, resident split boundary, exact sink column - moe_combine_masked skips a route masked to -1 instead of reading the validity row before its tensor. - The split HC boundary (LUCE_DS4_HC_BOUNDARY_SPLIT) runs only when the device holds its whole grid at once (occupancy x CUs): its blocks wait for each other at the output barrier, so a wider grid would never end. The unsplit grid stays at most MMVF_MAX_BATCH_SIZE blocks. - The sink-column softmax rounds the scaled score before the mask add on CUDA too (__fmul_rn; nvcc ignores the clang pragma). - The CPU softmax runs the masked sink-column form it used to assert on. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/deps/llama.cpp/ggml/include/ggml.h | 3 +- .../deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp | 5 +-- .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu | 35 +++++++++++++++++-- .../llama.cpp/ggml/src/ggml-cuda/moe-fused.cu | 4 +++ .../llama.cpp/ggml/src/ggml-cuda/softmax.cu | 4 ++- 5 files changed, 45 insertions(+), 6 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 400e03f463..b85fe15994 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2792,7 +2792,8 @@ extern "C" { // 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). + // 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, 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 ca3c22619a..093d2dfe95 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/ds4-hc.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu index 456f2c7d38..d78048808c 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 @@ -1104,6 +1104,34 @@ static __global__ void __launch_bounds__(1024) ds4_hc_boundary_split_kernel( } } +// 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 && @@ -1135,12 +1163,15 @@ void ggml_cuda_ds4_hc_boundary( // 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. + // 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; }(); - if (split_mix && n_tokens * 6 <= 48) { + 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))); 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 f09993b279..32c4f15f8c 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 @@ -407,6 +407,10 @@ static __global__ void moe_combine_masked_kernel( 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) { 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 69b0929714..8df3d3f42f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu @@ -59,8 +59,10 @@ struct soft_max_params { // 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 x*scale + m; + return __fmul_rn(x, scale) + m; } template From 46e9fbe287943887ba335e18d0b564510a70eaa9 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:09:26 +0200 Subject: [PATCH 31/40] refactor(ggml-cuda): remove ggml_backend_cuda_keepalive It kept a per-device stream and pinned host word in statics that were never freed, and its only caller is an opt-in switch that goes away. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../deps/llama.cpp/ggml/include/ggml-cuda.h | 5 -- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 58 ------------------- 2 files changed, 63 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index 3022c149c8..a8b8628698 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -188,11 +188,6 @@ GGML_BACKEND_API void ggml_backend_cuda_copy_batch_async(ggml_backend_t backend, // Graph copy runs issued as one batched launch on the calling thread (see // GGML_CUDA_DISABLE_COPY_BATCH); for tests and profiling. GGML_BACKEND_API size_t ggml_backend_cuda_get_copy_batch_run_count(void); -// Keeps the backend's GPU out of its idle power state: on starts one wave that -// sleeps in a loop on a stream of its own (at most max_ms), off releases it. -// For a process that leaves the GPU idle for milliseconds between launches and -// pays the wake-up (about 1 ms on RDNA4) at the next one. HIP only. -GGML_BACKEND_API void ggml_backend_cuda_keepalive(ggml_backend_t backend, bool on, int max_ms); GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); 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 f12cfdf05a..ae09e39f45 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 @@ -4628,64 +4628,6 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ return true; } -#if defined(GGML_USE_HIP) -static __global__ void ggml_cuda_keepalive_kernel(const volatile unsigned int * stop, unsigned int epoch, - unsigned long long max_ticks) { - const unsigned long long t0 = wall_clock64(); - while (*stop != epoch) { - __builtin_amdgcn_s_sleep(100); - if (wall_clock64() - t0 > max_ticks) break; - } -} -#endif - -void ggml_backend_cuda_keepalive(ggml_backend_t backend, bool on, int max_ms) { -#if defined(GGML_USE_HIP) - if (!ggml_backend_is_cuda(backend)) { - return; - } - struct state { - cudaStream_t stream = nullptr; - unsigned int * flag = nullptr; // pinned host word the kernel polls - unsigned int epoch = 0; - bool running = false; - unsigned long long ticks_per_ms = 0; - }; - static state states[GGML_CUDA_MAX_DEVICES]; - ggml_backend_cuda_context * ctx = (ggml_backend_cuda_context *) backend->context; - state & s = states[ctx->device]; - if (on) { - if (s.running) { - return; - } - ggml_cuda_set_device(ctx->device); - if (!s.stream) { - CUDA_CHECK(cudaStreamCreateWithFlags(&s.stream, cudaStreamNonBlocking)); - CUDA_CHECK(cudaMallocHost((void **) &s.flag, sizeof(unsigned int))); - *s.flag = 0; - int khz = 0; - if (hipDeviceGetAttribute(&khz, hipDeviceAttributeWallClockRate, ctx->device) != hipSuccess || khz <= 0) { - khz = 100000; - } - s.ticks_per_ms = (unsigned long long) khz; - } - s.epoch++; - ggml_cuda_keepalive_kernel<<<1, 32, 0, s.stream>>>(s.flag, s.epoch, - s.ticks_per_ms * (unsigned long long) std::max(1, max_ms)); - CUDA_CHECK(cudaGetLastError()); - s.running = true; - } else { - if (!s.running) { - return; - } - __atomic_store_n(s.flag, s.epoch, __ATOMIC_RELEASE); - s.running = false; - } -#else - GGML_UNUSED(backend); GGML_UNUSED(on); GGML_UNUSED(max_ms); -#endif -} - static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; From 71c78c300fd0ce68dccbb6e43b0c9d6663faa368 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:11:49 +0200 Subject: [PATCH 32/40] fix(ggml-cuda): fuse into MMVQ only where mul_mat would run MMVQ The CONT-into-MMVQ fusion ran MMVQ directly and the rope q8_1 producer assumed MMVQ, each with its own copy of a partial eligibility check. ggml_cuda_mul_mat sends some of those products elsewhere: RDNA4 Q4_K at two or more columns over 2048 or more rows goes to MMQ, MXFP8 to its own matvec, and the qwen4exp BF16 shadow to cuBLAS. The fused graph then ran another kernel than the unfused one (or MMVQ on MXFP8, which it has no kernel for). ggml_cuda_mul_mat now takes its MMVQ eligibility, the RDNA4 Q4_K MMQ preference and the BF16 shadow test from shared helpers, and both fusions decline unless ggml_cuda_mul_mat_takes_mmvq agrees. Also drops the kill switches this PR added for its default paths: LUCE_CUDA_Q4K_NCOL_MMQ, LUCE_DS4_FUSE_HC_BOUNDARY, LUCE_MMID_PREP_MEMO and LUCE_CUDA_MMVQ_Q8_RDNA4_GLU. The tests compare fused against unfused by graph shape, not by these switches. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 135 +++++++++++------- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 20 +-- 2 files changed, 85 insertions(+), 70 deletions(-) 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 ae09e39f45..73d11ac273 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 @@ -3273,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); @@ -3294,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; @@ -3328,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; @@ -3352,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; @@ -3365,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; @@ -3400,17 +3438,7 @@ 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); - // Q4_K at two or more columns over many rows (the DSpark drafter runs - // its dense projections at width five): on gfx1201 MMQ streams these - // at about twice the MMVQ rate (bench_dspark_mmv, width 5: attn_q_b - // 67 vs 137 us, attn_output_b 79 vs 117 us, fc 113 vs 217 us). MMVQ - // keeps the narrow shapes (under 2048 rows), where it is faster. - static const bool q4k_ncol_mmq = [] { - const char * e = std::getenv("LUCE_CUDA_Q4K_NCOL_MMQ"); - return !(e && e[0] == '0'); - }(); - if (q4k_ncol_mmq && use_mul_mat_vec_q && use_mul_mat_q && src0->type == GGML_TYPE_Q4_K && - GGML_CUDA_CC_IS_RDNA4(cc) && src1->ne[1] > 1 && src0->ne[1] >= 2048) { + 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; } } @@ -5445,15 +5473,7 @@ static bool ggml_cuda_node_is_noop(const ggml_tensor * node) { // 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. LUCE_DS4_FUSE_HC_BOUNDARY=0 -// keeps the six launches. -static bool ggml_cuda_ds4_hc_boundary_enabled() { - static const bool on = [] { - const char * v = getenv("LUCE_DS4_FUSE_HC_BOUNDARY"); - return !(v && v[0] == '0' && v[1] == '\0'); - }(); - return on; -} +// 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); @@ -5475,6 +5495,28 @@ 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 @@ -5614,13 +5656,7 @@ static int ggml_cuda_try_cont_into_mmvq(ggml_backend_cuda_context & ctx, ggml_cg walk = walk->src[0]; } if (walk != c || !ggml_is_contiguous(src1) || !ggml_cuda_elidable(cgraph, c, 1)) return -1; - // The plain MMVQ branch of ggml_cuda_mul_mat, nothing else. - 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; - if (!ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX || - src0->type == GGML_TYPE_Q3_1_ROCMFP3_MIX || src0->type == GGML_TYPE_Q2_0_ROCMFP2 || - bad_padding_clear || ggml_backend_buft_is_cuda_split(src0->buffer->buft) || - ggml_mul_mat_is_grouped_src(mm) || src1->ne[1] > ggml_cuda_mmvq_max_ncols()) 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; @@ -5668,13 +5704,7 @@ static bool ggml_cuda_try_rope_q8(ggml_backend_cuda_context & ctx, ggml_cgraph * 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; - // The plain MMVQ branch of ggml_cuda_mul_mat, nothing else. - 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; - if (!ggml_is_quantized(src0->type) || src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX || - src0->type == GGML_TYPE_Q3_1_ROCMFP3_MIX || src0->type == GGML_TYPE_Q2_0_ROCMFP2 || - bad_padding_clear || ggml_backend_buft_is_cuda_split(src0->buffer->buft) || - ggml_mul_mat_is_grouped_src(mm) || src1->ne[1] > ggml_cuda_mmvq_max_ncols()) 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; @@ -5690,7 +5720,6 @@ static bool ggml_cuda_try_rope_q8(ggml_backend_cuda_context & ctx, ggml_cgraph * } static int ggml_cuda_try_ds4_hc_boundary(ggml_backend_cuda_context & ctx, ggml_cgraph * cgraph, int i) { - if (!ggml_cuda_ds4_hc_boundary_enabled()) return -1; const int n = cgraph->n_nodes; // The next node at or after j that launches something. auto next = [&](int j) { 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 9dfc5e9d2b..d23cb4075d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -624,15 +624,6 @@ static bool mmid_grouped_type_ok(ggml_type type) { } } -// [TAG_MMID_GROUPED] LUCE_MMID_PREP_MEMO=0 recomputes the expert grouping for -// every MUL_MAT_ID (see the grouped path in ggml_cuda_mul_mat_vec_q). -static bool mmid_prep_memo_on() { - static const bool on = [] { - const char * e = std::getenv("LUCE_MMID_PREP_MEMO"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); - return on; -} // The memo entry kind of a cached expert grouping (q8_1 entries use a type or -1). #define MMID_PREP_MEMO_KIND (-2) @@ -2938,15 +2929,10 @@ static void mul_mat_vec_q8_0_rdna4_launch( // 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). LUCE_CUDA_MMVQ_Q8_RDNA4_GLU=0 -// keeps three launches. +// 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) - static const bool enabled = []() { - const char * e = std::getenv("LUCE_CUDA_MMVQ_Q8_RDNA4_GLU"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); - if (!enabled || !mmvq_q8_0_rdna4_enabled() || !ggml_cuda_mmvq_batch_invariant()) return false; + 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 || @@ -3911,7 +3897,7 @@ void ggml_cuda_mul_mat_vec_q( // 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 && mmid_prep_memo_on() && src1_q8_1.ptr == nullptr && + 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) { From 617c75e56c2a250d9f5328b0d553d9324d7a275a Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:11:49 +0200 Subject: [PATCH 33/40] fix(ggml-cuda): skinny F16 matvec on HIP RDNA4 only The skinny-K kernel accumulates in F32, so it changes the bits of the single-column F16 products it takes. It had no vendor guard and changed CUDA results too. It now runs only on RDNA4 under HIP, the GPU it was tuned on; CUDA and other AMD archs keep the generic MMVF kernel. Its LUCE_CUDA_MMVF_SKINNY switch and benchmark-log comment are gone. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu | 31 +++++++++---------- 1 file changed, 14 insertions(+), 17 deletions(-) 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 8b8356f287..c06d104dc0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu @@ -677,13 +677,13 @@ 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). The -// generic kernel gives every row its own block, so a K=256 row is a single -// 512-byte trip and the launch is bound by block scheduling (~360 GB/s on -// gfx1201 for the DSpark Markov head, vocab 129280 x rank 256). Here every lane -// loads 16 bytes, K/8 lanes share a row and one wave covers several rows, which -// streams the matrix at ~540 GB/s. Sums accumulate in F32, so values can differ -// from the generic F16-accumulating path in the last bits. +// 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, @@ -726,18 +726,12 @@ static __global__ void __launch_bounds__(256) mul_mat_vec_f16_skinny( } } -static bool mmvf_skinny_enabled() { - static const bool enabled = [] { - const char * value = std::getenv("LUCE_CUDA_MMVF_SKINNY"); - return !(value && value[0] == '0'); - }(); - return enabled; -} - // 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 (src0->type != GGML_TYPE_F16 || ids || !mmvf_skinny_enabled()) { +#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)) { @@ -758,7 +752,6 @@ static bool mul_mat_vec_f16_skinny_try(const ggml_tensor * src0, const ggml_tens } const float * x_bias = fusion && fusion->x_bias ? (const float *) fusion->x_bias->data : nullptr; - const int device = ggml_cuda_get_device(); 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); @@ -773,6 +766,10 @@ static bool mul_mat_vec_f16_skinny_try(const ggml_tensor * src0, const ggml_tens } 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, From baac93983cf7d66643f8792e43237d7de166f92a Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:11:02 +0200 Subject: [PATCH 34/40] refactor(ds4): drop verify input timing, pass the vocab split to the builder The per-verify input timers and the vocab split debug line were research aids. The vocab split is now a builder argument instead of a field kept across Extra::reset(), and the split head logs its rows once. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/deepseek4/deepseek4_fused_verify.inc | 43 +++++++------------ server/src/deepseek4/deepseek4_graph.cpp | 4 +- 2 files changed, 17 insertions(+), 30 deletions(-) diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 512bd95d81..8376e43d0c 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -455,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, @@ -520,7 +523,8 @@ static bool ds4_build_fused_verify_graph( } step_graph_free(fg.sg); fg.reset_nodes(); - { const bool keep_vocab_split = ex.vocab_split; ex.reset(); ex.vocab_split = keep_vocab_split; } + 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) { @@ -1741,6 +1745,12 @@ static bool ds4_build_fused_verify_graph( 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] @@ -1752,9 +1762,9 @@ static bool ds4_build_fused_verify_graph( 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); + 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); @@ -2187,13 +2197,6 @@ static int ds4_try_fused_verify_step( 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); - static bool vocab_split_reported = false; - if (vocab_split_env && !vocab_split_reported) { - vocab_split_reported = true; - std::fprintf(stderr, "[ds4-fused-verify] vocab split %d (hybrid %d argmax_only %d argmax_out %d cluster %d)\n", - vocab_split ? 1 : 0, hybrid ? 1 : 0, hooks->prefer_argmax_only ? 1 : 0, - hooks->argmax_out ? 1 : 0, cache.cluster_rt ? cache.cluster_rt->size() : 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) { @@ -2252,13 +2255,12 @@ static int ds4_try_fused_verify_step( // that state is live can replay kernels against stale allocations. fg->destroy(vc.backend, vc.peer_backend); ex->reset(); - ex->vocab_split = vocab_split; const auto build_t0 = Ds4TimingClock::now(); 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(); @@ -2377,12 +2379,6 @@ static int ds4_try_fused_verify_step( } ds4_fv_set(fg->i64_bundle, i64v.data(), sizeof(int64_t) * i64v.size()); - static const bool set_timing = std::getenv("LUCE_DS4_SET_TIMING") != nullptr; - static double st_inputs = 0, st_mask = 0, st_hash = 0, st_sync = 0; static int st_n = 0; - auto st_now = [] { return std::chrono::steady_clock::now(); }; - auto st_us = [](auto a, auto b) { return std::chrono::duration(b - a).count(); }; - const auto st_t1 = st_now(); - st_inputs += st_us(set_t0, st_t1); // Causal mask values. At long context this tensor is multi-megabyte, // while only the ring slots overwritten by this q-wide batch and the // padded compressed tail change between verifier steps. Build those @@ -2500,8 +2496,6 @@ static int ds4_try_fused_verify_step( } } - const auto st_t2 = st_now(); - st_mask += st_us(st_t1, st_t2); if (token_ids) { for (int il = 0; il < w.n_layer; ++il) { ggml_tensor * hids = fg->hash_ids[(size_t) il]; @@ -2529,14 +2523,7 @@ static int ds4_try_fused_verify_step( // 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. - const auto st_t3 = st_now(); - st_hash += st_us(st_t2, st_t3); if (pinned_io) ggml_backend_synchronize(backend); - st_sync += st_us(st_t3, st_now()); - if (set_timing && ++st_n % 64 == 0) { - std::fprintf(stderr, "[ds4-set-t] per verify over %d: inputs %.0f mask %.0f hash %.0f sync %.0f us (pinned %d)\n", - st_n, st_inputs / st_n, st_mask / st_n, st_hash / st_n, st_sync / st_n, pinned_io ? 1 : 0); - } // 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 diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 3a96ab4a41..e4b85b5e4b 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -6996,7 +6996,7 @@ 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; // the head covers this rank's vocabulary half (argmax only) + 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 @@ -10235,7 +10235,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 " From 2144c3884e4ecd5f935ff405f1ee079d5a8b5a97 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:13:08 +0200 Subject: [PATCH 35/40] refactor(ds4): drop kill switches whose on path is the default Pinned step I/O, the async Engram read, the lane batch compressor and selection, the masked MoE combine and the drafter's device context window are on wherever they run, so their switches and off paths go. The shared expert slice view, sliced dense weights and the prefill device join are cluster-only and on in cluster mode. The worker keepalive was never on. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/common/moe_hybrid_ffn_eval.cpp | 12 +- server/src/deepseek4/deepseek4_backend.cpp | 8 +- .../src/deepseek4/deepseek4_dspark_spec.cpp | 18 -- .../src/deepseek4/deepseek4_fused_verify.inc | 32 +-- server/src/deepseek4/deepseek4_graph.cpp | 227 +++++------------- 5 files changed, 70 insertions(+), 227 deletions(-) diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index 408a3c1c7e..bd376e12e4 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -1067,8 +1067,6 @@ static bool moe_owner_takes_coarse_op(const MoeOwnerGraphSpec & owner, const Moe return split || fused; } -// LUCE_MOE_MASKED_COMBINE=0 keeps the separate route-mask get_rows and mul, -// the weighting mul and the route-sum repeat_back. // LUCE_MOE_Q8_HANDOFF=1: see build_moe_hybrid_ffn_graph. static bool moe_q8_handoff_enabled() { static const bool on = [] { @@ -1078,14 +1076,6 @@ static bool moe_q8_handoff_enabled() { return on; } -static bool moe_masked_combine_enabled() { - static const bool on = [] { - const char * v = std::getenv("LUCE_MOE_MASKED_COMBINE"); - return !(v && v[0] == '0' && v[1] == '\0'); - }(); - return on; -} - static void align_moe_owner_routes( ggml_context * ctx, int n_tokens, @@ -1451,7 +1441,7 @@ bool build_moe_hybrid_ffn_graph( // 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 moe_masked_combine_enabled() && allow_fused_combine && + 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); diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 25cd85be36..d83b41232c 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -3634,11 +3634,9 @@ bool DeepSeek4Backend::init_hybrid_model() { TargetLoadPlan plan; plan.skip_expert_tensors = true; plan.load_ds4_image_bias = !cfg_.mmproj_path.empty(); - // LUCE_CLUSTER_SLICE_WEIGHTS=1: a cluster rank loads only the dense - // slices its head split and shared-expert shard read (about 3.3 GB less - // on the primary GPU for DS4.1 on two ranks, for more hot experts). - if (cluster_ && cluster_->cfg && cluster_->cfg->size > 1 && - env_flag_enabled("LUCE_CLUSTER_SLICE_WEIGHTS")) { + // 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(); diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 994f1f7e5c..e49da8b995 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -1080,16 +1080,6 @@ bool run_deepseek4_dspark_spec_decode( 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; - // LUCE_CLUSTER_WORKER_KEEPALIVE=1: between its verify and the head's next - // draft a worker's GPU is idle for milliseconds, and its first launch after - // that started ~1 ms late (wake-up), which the head then waited for at the - // first exchange. One sleeping wave keeps it awake until the draft arrives. - const bool worker_keepalive = cluster_spec && !cluster_spec_head && !draft_split && - spec_env_flag("LUCE_CLUSTER_WORKER_KEEPALIVE"); - struct KeepaliveOff { - ggml_backend_t backend; bool on; - ~KeepaliveOff() { if (on) ggml_backend_cuda_keepalive(backend, false, 0); } - } keepalive_off{backend, worker_keepalive}; // 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. @@ -1579,11 +1569,6 @@ bool run_deepseek4_dspark_spec_decode( // (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) { - // A worker's keepalive wave may share a hardware queue with the - // compute stream and would hold the head's graph until the wave's - // cap (~40 ms a step): stop it, queue the graph, then queue the - // wave behind it. - if (worker_keepalive) ggml_backend_cuda_keepalive(backend, false, 0); 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)) { @@ -1591,7 +1576,6 @@ bool run_deepseek4_dspark_spec_decode( ok = false; break; } - if (worker_keepalive) ggml_backend_cuda_keepalive(backend, true, 40); } // Cluster lockstep: the head's draft block is authoritative. Every // rank drafts (the drafter forwards run concurrently on all ranks, so @@ -1610,7 +1594,6 @@ bool run_deepseek4_dspark_spec_decode( std::string err; const bool got_draft = cluster_hooks->decide_draft(cluster_req, (uint32_t) steps, pos, draft_tok, &err); - if (worker_keepalive) ggml_backend_cuda_keepalive(backend, false, 0); if (!got_draft) { std::fprintf(stderr, "[ds4-spec] cluster draft missing: %s\n", err.c_str()); ok = false; @@ -1694,7 +1677,6 @@ bool run_deepseek4_dspark_spec_decode( const bool verify_ok = target.verify_batch(draft_tok, pos, verify_last, &tgt_am); tm_verify += spec_ms_since(t0); - if (worker_keepalive) ggml_backend_cuda_keepalive(backend, true, 40); if (probe_inflight) { const SpecClock::time_point probe_t0 = SpecClock::now(); deepseek4_dspark_draft_wait(drafter_backend); diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 8376e43d0c..3299155816 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -593,7 +593,7 @@ static bool ds4_build_fused_verify_graph( 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 (ds4_lane_batch_comp_enabled() && !paged_mode && q > 1) { + if (!paged_mode && q > 1) { ex.pos_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.pos_rows); } @@ -1096,13 +1096,8 @@ static bool ds4_build_fused_verify_graph( } // 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 (LUCE_DS4_LANE_BATCH_SELECT=0 - // keeps one chain per lane and concatenates their selections). - static const bool lane_batch_select = [] { - const char * value = std::getenv("LUCE_DS4_LANE_BATCH_SELECT"); - return !value || std::strcmp(value, "0") != 0; - }(); - const bool external_selection = lane_batch && lane_batch_select && + // 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) { @@ -1121,7 +1116,7 @@ static bool ds4_build_fused_verify_graph( // 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 && ds4_lane_batch_comp_enabled() && attn_lane_q == 1 && + 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 && @@ -2275,7 +2270,7 @@ static int ds4_try_fused_verify_step( // 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 = ds4_pinned_step_io_enabled() && + 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: @@ -2286,8 +2281,7 @@ static int ds4_try_fused_verify_step( bool engram_async = false; if (ex->engram_keys) { const DeepSeek4EngramRuntime * engram = w.engram_runtime.get(); - if (pinned_io && engram_rt && engram && token_ids && ds4_engram_async_enabled() && - ex->engram_keys->buffer) { + 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, @@ -2482,16 +2476,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); - if (pinned_io) { - vc.stage.set(backend, fg->mask_bundle, maskv.data() + range.first, - range.first * sizeof(float), - (range.second - range.first) * sizeof(float)); - } else { - 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)); } } } diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index e4b85b5e4b..79b167849a 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2257,8 +2257,8 @@ static DeepSeek4PreparedProjectedLane build_mla_qkv_projection( // 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 rank holds only its attention heads " - "(LUCE_CLUSTER_SLICE_WEIGHTS) but a graph attends with all of them"); + 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, @@ -2417,19 +2417,6 @@ static bool ds4_lane_batch_enabled() { return enabled; } -// A lane batch runs a ratio-1 kv source's compressor for every lane in one -// width-q pass (LUCE_DS4_LANE_BATCH_COMP=0 keeps one pass per lane). Ratio 1 -// keeps no state: each lane's latent and index key are per-column products, -// and the MMVF block size does not depend on the column count on RDNA4, so -// the batched pass writes the rows the lanes would have written. -static bool ds4_lane_batch_comp_enabled() { - static const bool enabled = [] { - const char * v = std::getenv("LUCE_DS4_LANE_BATCH_COMP"); - return !v || std::strcmp(v, "0") != 0; - }(); - 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 @@ -4709,17 +4696,18 @@ static ggml_tensor * build_shared_ffn( return ggml_mul_mat(ctx, L.ffn_down_shexp, mid_sh); } -// Cluster shared-expert sharding (maikzz32/lucebox-halo-cluster): 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 must add it to the routed partial BEFORE the all-reduce, and the -// reduction then sums both contributions in one message. +// 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, which is a contiguous view. `down` -// contracts over the intermediate axis, so slicing it would need a strided -// view of every row; instead the local intermediate is zero-padded back to -// full width and the full `down` is applied. The zeros contribute nothing. +// 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, @@ -4741,19 +4729,16 @@ static ggml_tensor * build_shared_ffn_slice( 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); - // LUCE_DS4_SHEXP_SLICE_VIEW=1: contract over this rank's slice only, - // through a column view of down (row stride unchanged): half the down - // weights read and no zero padding. Bit-identical to the padded product - // on the RDNA4 Q8_0 matvec, which gives the slice's blocks the lanes they - // take in the whole rows (ggml_backend_cuda_mul_mat_whole_row_lanes); - // other kernels sum the slice in their own order. - static const bool slice_view = [] { - const char * v = std::getenv("LUCE_DS4_SHEXP_SLICE_VIEW"); - return v && *v && std::strcmp(v, "0") != 0; - }(); - if (slice_view && (int64_t) ff_count != n_ff && - ff_begin % ggml_blck_size(L.ffn_down_shexp->type) == 0 && - ff_count % ggml_blck_size(L.ffn_down_shexp->type) == 0) { + 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)); @@ -4761,16 +4746,13 @@ static ggml_tensor * build_shared_ffn_slice( ggml_backend_cuda_mul_mat_whole_row_lanes(out); return out; } - if ((int64_t) ff_count != n_ff) { - 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; - } - mid_sh = padded; + 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, mid_sh); + return ggml_mul_mat(ctx, L.ffn_down_shexp, padded); } bool init_deepseek4_streamed_expert_cache( @@ -6915,25 +6897,6 @@ class Ds4EngramAsync { size_t job_floats_ = 0; }; -// LUCE_DS4_ENGRAM_ASYNC=0 keeps the in-line verify Engram read. -static bool ds4_engram_async_enabled() { - static const bool enabled = [] { - const char * v = std::getenv("LUCE_DS4_ENGRAM_ASYNC"); - return !(v && v[0] == '0' && v[1] == '\0'); - }(); - return enabled; -} - -// LUCE_DS4_PINNED_STEP_IO=0 keeps the blocking per-tensor verify uploads and -// readbacks. -static bool ds4_pinned_step_io_enabled() { - static const bool enabled = [] { - const char * v = std::getenv("LUCE_DS4_PINNED_STEP_IO"); - return !(v && v[0] == '0' && v[1] == '\0'); - }(); - return enabled; -} - 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 @@ -8111,8 +8074,7 @@ static bool eval_ds4_layer_range_hybrid_ffn( // 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() && - ds4_cluster_prefill_device_join_enabled(); + cluster_rt && cluster_rt->size() > 1 && device_outputs && device_outputs->valid(); const bool device_ffn_input = (!cluster_rt || cluster_device_join) && device_input_enabled && layer_storage.n_streamed == 0 && !expert_compute && moe_expert_major_prefill_enabled(n_tokens) && @@ -12300,7 +12262,6 @@ bool deepseek4_step_layer_range( !device_input_env || !*device_input_env || std::strcmp(device_input_env, "0") != 0; const bool ffn_device_join_possible = - (!cache.cluster_rt || ds4_cluster_prefill_device_join_enabled()) && device_input_enabled && moe_expert_major_prefill_enabled(n_tokens) && layer_storage.cold_backend_kind == @@ -13339,7 +13300,7 @@ static ggml_tensor * build_dspark_moe_split(ggml_context * ctx, 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 ggml_add(ctx, routed, shared); + return shared ? ggml_add(ctx, routed, shared) : nullptr; } // Read a small F32 GPU tensor (HC scale, [k]) into host floats. @@ -13360,9 +13321,8 @@ 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. +// 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; @@ -13391,11 +13351,10 @@ struct DsparkContextKvProjector { int end_pos = -1; int valid = 0; - std::vector host_kv; - std::vector projected; - PinnedStage stage; // inputs and the readback of one projection + std::vector host_kv; // the all-zero window before any context + PinnedStage stage; // inputs of one projection - // The window stays on the draft device (LUCE_DS4_DRAFT_DEVICE_CONTEXT), + // 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. @@ -13411,28 +13370,17 @@ struct DsparkContextKvProjector { } }; -// Default on; LUCE_DS4_DRAFT_DEVICE_CONTEXT=0 keeps the host-side window. -static bool dspark_device_context_enabled() { - static const bool on = [] { - const char * v = std::getenv("LUCE_DS4_DRAFT_DEVICE_CONTEXT"); - return !(v && v[0] == '0' && v[1] == '\0'); - }(); - return on; -} - thread_local DsparkContextKvProjector g_dspark_ctx_kv; -// Projects n_cols new feature columns to the drafter's per-layer context K/V. -// out != null: read the [head_dim, n_cols, n_layer] result back (one wait). -// out == null: leave it on the device in *dev_out, queued on backend's stream. +// 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 = nullptr) { + ggml_tensor ** dev_out) { if (!backend || !features || n_cols <= 0) return false; const DeepSeek4Weights & w = d.core; const int n_embd = w.n_embd; @@ -13493,34 +13441,20 @@ static bool dspark_project_context_columns( G.n_cols = n_cols; } - // Inputs and the readback go through pinned slices: one wait in all. + // 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; - const size_t out_bytes = sizeof(float) * (size_t) head_dim * n_cols * w.n_layer; - const bool pinned = ds4_pinned_step_io_enabled() && - P.stage.reserve(backend, in_bytes + out_bytes + sizeof(int32_t) * (size_t) n_cols + 4096); + 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; - if (pinned) { - 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()); - } else { - ggml_backend_tensor_set(G.inp_features, features, 0, in_bytes); - ggml_backend_tensor_set(G.positions, pos.data(), 0, sizeof(int32_t) * pos.size()); - } + 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; } - if (!out) { - if (dev_out) *dev_out = G.out; - return true; - } - out->resize((size_t) head_dim * n_cols * w.n_layer); - const void * src = pinned ? P.stage.get(backend, G.out, 0, out_bytes) : nullptr; - ggml_backend_synchronize(backend); - if (src) std::memcpy(out->data(), src, out_bytes); - else ggml_backend_tensor_get(G.out, out->data(), 0, out_bytes); + *dev_out = G.out; return true; } @@ -13531,13 +13465,13 @@ static bool dspark_update_context_kv_cache( int ctx_len, int committed, const std::vector ** out_host_kv, - ggml_tensor ** out_dev_kv = nullptr) { + 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 (out_dev_kv && dspark_device_context_enabled() && 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 @@ -13574,7 +13508,7 @@ static bool dspark_update_context_kv_cache( 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, - nullptr, &projected) || !projected) { + &projected) || !projected) { return false; } const int keep = std::max(0, std::min(P.valid, n_swa - n_new)); @@ -13606,58 +13540,10 @@ static bool dspark_update_context_kv_cache( *out_host_kv = nullptr; return true; } - if (ctx_len <= 0) { - 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; - 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; } @@ -14013,8 +13899,7 @@ static bool deepseek4_dspark_draft_forward_impl( // ── Set inputs + compute (cached graph) ───────────────────────────── // 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 = ds4_pinned_step_io_enabled() && - C.stage.reserve(backend, + 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) + @@ -14023,9 +13908,9 @@ static bool deepseek4_dspark_draft_forward_impl( (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) { - if (pinned_io) C.stage.set(backend, t, data, 0, bytes); - else ggml_backend_tensor_set(t, data, 0, 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) { @@ -14141,8 +14026,8 @@ static bool deepseek4_dspark_draft_forward_impl( const size_t hidden_bytes = sizeof(float) * (size_t) n_embd * block; out_hidden->resize((size_t) n_embd * block); if (confidence_hidden) confidence_hidden->resize((size_t) n_embd * block); - const void * hidden_src = pinned_io ? C.stage.get(backend, C.out, 0, hidden_bytes) : nullptr; - const void * conf_src = pinned_io && confidence_hidden + 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); From 3e6bf2692db56a568ef95df795c8322491345f42 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:13:37 +0200 Subject: [PATCH 36/40] docs(ds4): drop benchmark numbers and repeated attribution from comments The attribution to the upstream cluster project stays once in deepseek4_cluster.h. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/common/moe_hybrid_ffn_eval.cpp | 3 +-- server/src/deepseek4/deepseek4_backend.cpp | 4 ++-- server/src/deepseek4/deepseek4_backend.h | 2 +- server/src/deepseek4/deepseek4_dspark.h | 2 +- server/src/deepseek4/deepseek4_dspark_spec.cpp | 14 +++++++------- server/src/deepseek4/deepseek4_fused_verify.inc | 2 +- server/src/deepseek4/deepseek4_graph.cpp | 6 +++--- 7 files changed, 16 insertions(+), 17 deletions(-) diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index bd376e12e4..1a4ab683fe 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -1463,8 +1463,7 @@ bool build_moe_hybrid_ffn_graph( // 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 (on an R9700 behind PCIe 3.0 x4 - // the per-layer handoff grows ~0.25 us per KB), and the secondary owner + // 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; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index d83b41232c..f64a69cf97 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1142,7 +1142,7 @@ void log_deepseek4_step_telemetry(const char * phase, DeepSeek4Backend::DeepSeek4Backend(DeepSeek4BackendConfig cfg) : cfg_(std::move(cfg)) {} -// ─── Expert-parallel cluster (after maikzz32/lucebox-halo-cluster) ──────── +// ─── Expert-parallel cluster ─────────────────────────────────────────── bool DeepSeek4Backend::cluster_active() const { return hooks_ != nullptr && hooks_->is_cluster(); @@ -4825,7 +4825,7 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, } } - // Cluster lockstep (maikzz32/lucebox-halo-cluster): rank 0 samples and + // 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(); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index dddf71778f..f4a4a3b237 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -135,7 +135,7 @@ class DeepSeek4Backend : public ModelBackend, public IClusterParticipant { return routing_stats_.get(); } - // ── Expert-parallel cluster (after maikzz32/lucebox-halo-cluster) ── + // ── 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. diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index 1d34cd778f..6015695aa2 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -338,7 +338,7 @@ bool run_deepseek4_dspark_spec_decode( DSparkSpecSampling * sampling = nullptr, // Polled before applying the budget or emitting another token. const std::function & is_cancelled = {}, - // Cluster lockstep hooks (from maikzz32's lucebox-halo-cluster): head/worker + // Cluster lockstep hooks: head/worker // hooks broadcast/receive DraftMsg and AcceptMsg; nullptr = one box. luce::cluster::Ds4ClusterHooks * cluster_hooks = nullptr); diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index e49da8b995..1250d96809 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -337,8 +337,8 @@ bool spec_env_default_on(const char * name) { // 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 verifies q5 for 1.3x the q2 step, not 1.6x) widens on its own. The fused verify cache is -// sized so every width stays resident across the ratio-4 phases. +// 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 @@ -837,7 +837,7 @@ void spec_restore_window_rows(const Ds4VerifyWindowRows & rows, const DeepSeek4W // 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. About 100 MB, copied once per drafter. +// and the rank's rows of markov_w2, copied once per drafter. struct SplitMarkovCopies { ggml_context * ctx = nullptr; ggml_backend_buffer_t buf = nullptr; @@ -1065,7 +1065,7 @@ bool run_deepseek4_dspark_spec_decode( const std::function & is_cancelled, luce::cluster::Ds4ClusterHooks * cluster_hooks) { const int n_embd = target_w.n_embd; - // Cluster lockstep (from maikzz32's lucebox-halo-cluster). Only a real + // 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(); @@ -1303,8 +1303,8 @@ bool run_deepseek4_dspark_spec_decode( const SpecClock::time_point run_t0 = SpecClock::now(); // LUCE_CLUSTER_DRAFT_VOCAB_SPLIT=1 (two ranks): the draft head reads the - // target's 703 MB Q8_0 lm_head every step (~1.1 ms on the head's dGPU, - // while the worker's waits for the block). Split it like the verifier's + // 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. @@ -1501,7 +1501,7 @@ bool run_deepseek4_dspark_spec_decode( // 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: 1.67 -> 1.13 ms for four columns on gfx1201). + // Q8_0 path). const ScopedCudaGraphOverrides head_products( /*disable_graphs=*/false, /*mmvq_max_ncols=*/0, /*skip_property_check=*/false, /*ds4_mix_mmv_max_tokens=*/0, diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 3299155816..950cc662e7 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -536,7 +536,7 @@ 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 (from maikzz32's lucebox-halo-cluster): the + // 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 diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 79b167849a..159910345a 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2105,7 +2105,7 @@ struct DeepSeek4MlaLaneBindings { }; HistoryMode history_mode = HistoryMode::ContiguousRing; - // Expert-parallel cluster head split (maikzz32/lucebox-halo-cluster): + // 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. @@ -6776,8 +6776,8 @@ struct DeepSeek4FusedDecodeCache { // 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: about 0.7 ms a step on the rank that holds the tables, -// and every other rank waits for its broadcast). The first consumer is the +// 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: From 0972e72d6f3501d095a8f489924af6ca2c306516 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:11:58 +0200 Subject: [PATCH 37/40] refactor(cluster): drop dead DS4 cluster helpers and research comments Remove ds4_cluster_allreduce_layer and ds4_cluster_allreduce_fn, which had no callers, and the LUCE_CLUSTER_TRACE_COLL debug print. The plain and exact in-graph broadcasts share one callback body and one node builder. Name the prefill partial threshold. Comments lose benchmark numbers, plan labels, the protocol changelog and a stale note about mixed qtypes; one attribution stays in deepseek4_cluster.h. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/cluster_protocol.h | 4 - server/src/cluster/cluster_telemetry.h | 7 +- server/src/cluster/cluster_worker_main.h | 4 +- server/src/cluster/fast_reduce.cpp | 23 ++-- server/src/cluster/fast_reduce.h | 15 +-- server/src/cluster/fast_reduce_kernels.cu | 20 ++-- server/src/deepseek4/deepseek4_cluster.cpp | 123 ++++++--------------- server/src/deepseek4/deepseek4_cluster.h | 33 ++---- server/src/server/server_main.cpp | 2 +- 9 files changed, 71 insertions(+), 160 deletions(-) diff --git a/server/src/cluster/cluster_protocol.h b/server/src/cluster/cluster_protocol.h index 48feee79a4..b3d3b76f0c 100644 --- a/server/src/cluster/cluster_protocol.h +++ b/server/src/cluster/cluster_protocol.h @@ -29,10 +29,6 @@ namespace luce::cluster { inline constexpr uint32_t kFrameMagic = 0x4C43424Cu; // "LBCL" -// 2: RequestMsg carries the inline snapshot position and separates the -// restore slot from the inline save slot, so the prefix cache can be -// replicated (a request may restore from one slot and check-point into -// another; before this the workers never took the inline snapshot at all). inline constexpr uint16_t kProtocolVersion = 2; inline constexpr size_t kFrameHeaderBytes = 12; inline constexpr size_t kMaxPayloadBytes = 64ull * 1024ull * 1024ull; diff --git a/server/src/cluster/cluster_telemetry.h b/server/src/cluster/cluster_telemetry.h index d8534d4c37..94ce22cf7c 100644 --- a/server/src/cluster/cluster_telemetry.h +++ b/server/src/cluster/cluster_telemetry.h @@ -1,10 +1,9 @@ // 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 (Agent C) and the -// decision hooks / worker main (Agent D) can accumulate into the same struct; -// the head folds the per-rank totals into RequestReportMsg and -// usage.timings.cluster. +// 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 diff --git a/server/src/cluster/cluster_worker_main.h b/server/src/cluster/cluster_worker_main.h index 28f4c46be2..fe23a1a592 100644 --- a/server/src/cluster/cluster_worker_main.h +++ b/server/src/cluster/cluster_worker_main.h @@ -1,5 +1,5 @@ -// cluster_worker_main.h - lockstep worker loop for ranks 1..N-1 -// (maikzz32/lucebox-halo-cluster). server_main builds the backend through the +// 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. diff --git a/server/src/cluster/fast_reduce.cpp b/server/src/cluster/fast_reduce.cpp index 9e57f31a4d..8242ae1cc2 100644 --- a/server/src/cluster/fast_reduce.cpp +++ b/server/src/cluster/fast_reduce.cpp @@ -235,8 +235,9 @@ static bool fast_reduce_lean() { // The kernel that writes an exchange row also publishes its flag (one // system-scope release by its grid's last block) instead of a -// hipStreamWriteValue32 after it: about 0.6 ms less per two-box verify step, -// same bits. LUCE_CLUSTER_FAST_REDUCE_KPUB=0 goes back to the stream writes. +// hipStreamWriteValue32 after it, which saves a stream operation per row and +// gives the same bits. LUCE_CLUSTER_FAST_REDUCE_KPUB=0 goes back to the stream +// writes. static bool fast_reduce_kpub() { static const bool on = [] { const char * v = std::getenv("LUCE_CLUSTER_FAST_REDUCE_KPUB"); @@ -287,10 +288,9 @@ struct FastReduceImpl { uint32_t ** peer_flags = nullptr; // [slots][size-1] uint32_t ** peer_flags_dev = nullptr; - // Device memory the peers' partials are copied into before they are added. - // Adding straight out of the pinned staging buffer is what made the first - // version 250 times slower than the collective it replaces: that memory is - // uncached for the GPU. + // Device memory the peers' partials are copied into before they are added: + // the pinned staging buffer is uncached for the GPU, so adding straight + // out of it is slow. float * scratch = nullptr; // [slots][size-1][max_elems] uint32_t * timed_out_host = nullptr; @@ -817,10 +817,9 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { // A kernel cannot do this cheaply on gfx1151. Raising a flag the host will // see needs a system-scope release, which writes the L2 back; spinning on // one the host raised needs a system-scope acquire, which invalidates it. - // Either way the weights the next layer is about to read are gone, and - // measured that cost 34 ms per reduction -- while the same path with the - // flags removed ran the step in 32.5 ms. Weakening the ordering does not - // help: relaxed is cheap and the host then does not see the flag at all. + // Either way the weights the next layer is about to read are gone. + // Weakening the ordering does not help: relaxed is cheap and the host then + // does not see the flag at all. // // hipStreamWriteValue32 and hipStreamWaitValue32 are handled by the // command processor instead, so the shader cores and their caches are not @@ -886,9 +885,7 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { } // Wait for the peers, then bring their partials into device memory before - // touching them: the staging buffers are uncached for the GPU, and adding - // straight out of them is what made the first version 250 times slower - // than the collective it replaces. + // touching them: the staging buffers are uncached for the GPU. if (kernel_flags) { fast_reduce_launch_wait_flags(s.peer_flags_dev + (size_t) slot * (size_t) n_peers, s.timed_out_dev, s.progress_dev, diff --git a/server/src/cluster/fast_reduce.h b/server/src/cluster/fast_reduce.h index 48f389acc8..e7ea3db86d 100644 --- a/server/src/cluster/fast_reduce.h +++ b/server/src/cluster/fast_reduce.h @@ -1,16 +1,9 @@ // fast_reduce.h - a lean all-reduce for the decode path's tiny payloads. // -// WHY THIS EXISTS. A qwen4exp decode step on two nodes computes in 29.6 ms of -// GPU work and takes 41.0. The difference is ninety-seven all-reduces of -// 10 KiB each -- 117 us apiece. That is not the wire: ib_write_lat puts the -// same payload at 13.65 us, and the round trip through a pinned flag between -// the GPU and a host thread measures 0.79 us. It is what a general-purpose -// collective costs per call, ninety-seven times, and it is the single reason -// adding ranks to this model buys nothing: every axis worth sharding pays for -// its bytes with reductions at about 11.9 MB each, and every axis this model -// has sits on that line. -// -// So this replaces the collective for exactly the case it is bad at: a few +// 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 diff --git a/server/src/cluster/fast_reduce_kernels.cu b/server/src/cluster/fast_reduce_kernels.cu index e4bab2fff7..ff0d71801e 100644 --- a/server/src/cluster/fast_reduce_kernels.cu +++ b/server/src/cluster/fast_reduce_kernels.cu @@ -11,16 +11,15 @@ // set_flag, the kernel form of the flags, kept for comparison // wait_flags (LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS=1). // -// WHY THE COPY PATH MOVES THE PAYLOAD BY DMA. A kernel moving it cost 33 ms per -// reduction against RCCL's 117 us. Pinned host memory is mapped uncached for -// the GPU, so a kernel reading or writing bulk data there crawls; the copy -// engine does not care, because it is not reading through the GPU's caches. +// WHY THE COPY PATH MOVES THE PAYLOAD BY DMA. Pinned host memory is mapped +// uncached for the GPU, so a kernel reading or writing bulk data there crawls; +// the copy engine does not care, because it is not reading through the GPU's +// caches. // // WHY THE FLAGS ARE PINNED AND NOT MANAGED. 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, so they need -// no fault to be reached -- and the round trip through one measures 0.79 us -// (server/test/cluster_flag_latency.cu), which is what makes this worth doing. +// no fault to be reached. // // 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. @@ -38,9 +37,7 @@ constexpr int kBlock = 256; // Ordering costs cache, and on this part it costs all of it. A system-scope // release writes the L2 back and a system-scope acquire invalidates it, so a // flag raised or spun on with the strong form evicts the weights the next layer -// is about to read. Four of them per reduction, ninety-seven reductions per -// step: measured, that alone was 34 ms per reduction against RCCL's 117 us, -// while the same path with the announcements removed ran the step in 32.5 ms. +// is about to read, several times per reduction. // // So each flag gets the weakest ordering that is still correct: // @@ -164,9 +161,8 @@ __device__ __forceinline__ float fr_bf16_to_f32(uint16_t h) { // LUCE_CLUSTER_FAST_REDUCE_KPUB: 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. Measured against a hipStreamWriteValue32 after the kernel: -// publish + host ack 9-11 us instead of 16-18 on gfx1201, 3.4 instead of 6.3 on -// gfx1151. A null flag publishes nothing. counter is this device's, zeroed, and +// stores the flag, which is cheaper than a hipStreamWriteValue32 after the +// kernel. A null flag publishes nothing. 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) { if (!flag) return; diff --git a/server/src/deepseek4/deepseek4_cluster.cpp b/server/src/deepseek4/deepseek4_cluster.cpp index d1351ebcf7..65d3a52a45 100644 --- a/server/src/deepseek4/deepseek4_cluster.cpp +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -20,15 +20,6 @@ namespace luce::common { -// LUCE_CLUSTER_TRACE_COLL=1: one stderr line per collective (rank, sequence, kind, elements, stream), so the -// two ranks' collective sequences can be diffed when a step hangs. -static void trace_coll(int rank, const char * kind, size_t n, const void * stream) { - static const bool on = std::getenv("LUCE_CLUSTER_TRACE_COLL") != nullptr; - static unsigned long long seq = 0; - if (!on) return; - std::fprintf(stderr, "[coll] r%d #%llu %s n=%zu stream=%p\n", rank, seq++, kind, n, stream); -} - namespace { using ClusterClock = std::chrono::steady_clock; @@ -47,11 +38,9 @@ void set_err(std::string * err, const std::string & msg) { if (err) *err = msg; } -// Mixed ROCmFP qtypes need decode tables registered per resident tensor -// (register_deepseek4_moe_hybrid_mix_tables), which today requires a -// materialized GPU cold owner. Refuse instead of silently loading garbage. -// Supporting them needs register_deepseek4_moe_hybrid_mix_tables to accept -// hot-only storage (cold owner None). +// 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[] = { @@ -105,7 +94,6 @@ bool allreduce_device(Ds4ClusterRuntime & rt, // enqueues on (cuda_ctx->stream()), so the collective is ordered after the // producing kernels. cluster::DeviceStream stream = ggml_backend_cuda_get_stream(backend); - trace_coll(rt.rank(), "host-allreduce", n, (const void *) stream); const uint64_t bytes_f32 = (uint64_t) n * sizeof(float); const cluster::AllreduceDType dtype = rt.cfg->allreduce_dtype; @@ -551,7 +539,11 @@ void ds4_cluster_mask_routes(const Ds4ClusterRuntime & rt, // ─── All-reduce ───────────────────────────────────────────────────────── -// Path 3b: invoked by GGML_MOE_FUSED_CLUSTER_ALLREDUCE while the graph runs. +// 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, @@ -573,9 +565,8 @@ static void ds4_cluster_allreduce_graph_callback(void * user, void * data, const char * v = std::getenv("LUCE_CLUSTER_PREFILL_Q8"); return v && *v && std::strcmp(v, "0") != 0; }(); - if (prefill_q8 && n >= (size_t) 262144 && n % cluster::kQ8ExchangeBlock == 0 && rt->scratch_f32 && + 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)) { - trace_coll(rt->rank(), "graph-allreduce-q8", n, stream); std::string errq; if (!rt->comm->allreduce_sum_q8_exchange(data, n, rt->scratch_f32->data, (cluster::DeviceStream) stream, &errq)) { @@ -586,8 +577,7 @@ static void ds4_cluster_allreduce_graph_callback(void * user, void * data, rt->telemetry.allreduce_bytes += (uint64_t) cluster::q8_exchange_slot_bytes(n); return; } - if (prefill_bf16 && n >= (size_t) 262144 && rt->scratch_bf16 && rt->scratch_elems >= n) { - trace_coll(rt->rank(), "graph-allreduce-bf16", n, stream); + 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)) { @@ -599,12 +589,10 @@ static void ds4_cluster_allreduce_graph_callback(void * user, void * data, return; } if (rt->fast && rt->fast->ok() && rt->fast->submit((float *) data, n, stream)) { - trace_coll(rt->rank(), "fast-allreduce", n, stream); rt->telemetry.allreduce_calls += 1; rt->telemetry.allreduce_bytes += (uint64_t) n * sizeof(float); return; } - trace_coll(rt->rank(), "graph-allreduce", n, stream); 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 @@ -689,10 +677,6 @@ ggml_tensor * ds4_cluster_hybrid_exchange_nodes(ggml_context * ctx, return imp; } -ggml_cluster_allreduce_fn ds4_cluster_allreduce_fn() { - return &ds4_cluster_allreduce_graph_callback; -} - ggml_tensor * ds4_cluster_allreduce_node(ggml_context * ctx, ggml_tensor * partial, Ds4ClusterRuntime & rt, @@ -712,28 +696,6 @@ ggml_tensor * ds4_cluster_allreduce_node(ggml_context * ctx, &ds4_cluster_allreduce_graph_callback, &rt); } -bool ds4_cluster_allreduce_layer(Ds4ClusterRuntime & rt, - ggml_backend_t backend, - ggml_tensor * partial, - int layer, - DeepSeek4StepTelemetry * telemetry, - std::string * err) { - if (!partial) { - set_err(err, "cluster all-reduce: null partial tensor"); - return false; - } - if (partial->type != GGML_TYPE_F32 || !ggml_is_contiguous(partial)) { - set_err(err, "cluster all-reduce: partial must be contiguous F32"); - return false; - } - const auto t0 = ClusterClock::now(); - const size_t n = (size_t) ggml_nelements(partial); - const bool ok = allreduce_device(rt, backend, partial->data, n, layer, - /*force_wait=*/false, telemetry, err); - if (telemetry) telemetry->cluster_allreduce_us += us_since(t0); - return ok; -} - bool ds4_cluster_allreduce_layer_host(Ds4ClusterRuntime & rt, ggml_backend_t backend, float * partial_host, @@ -767,18 +729,20 @@ bool ds4_cluster_allreduce_layer_host(Ds4ClusterRuntime & rt, namespace { -void broadcast_graph_callback(void * user, void * data, size_t n, void * stream) { +// 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 (rt->before_broadcast && !rt->before_broadcast(data, n, stream)) { + 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)) { - trace_coll(rt->rank(), "fast-broadcast", n, stream); + rt->fast->submit_from_root((float *) data, n, stream, /*root=*/0, exact)) { return; } - trace_coll(rt->rank(), "graph-broadcast", n, stream); std::string err; if (!rt->comm->broadcast_bytes(data, n * sizeof(float), /*root=*/0, (cluster::DeviceStream) stream, &err)) { @@ -788,31 +752,18 @@ void broadcast_graph_callback(void * user, void * data, size_t n, void * stream) } } -// Rank 0's values bit for bit: f32 on the wire even where the fast reduce -// rounds its payloads to bf16. No upload hook (that belongs to the Engram keys). -void broadcast_exact_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; - if (rt->fast && rt->fast->ok() && - rt->fast->submit_from_root((float *) data, n, stream, /*root=*/0, /*exact=*/true)) { - trace_coll(rt->rank(), "fast-broadcast-exact", n, stream); - return; - } - trace_coll(rt->rank(), "graph-broadcast", n, stream); - 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); } -} // namespace +void broadcast_exact_graph_callback(void * user, void * data, size_t n, void * stream) { + broadcast_in_graph(user, data, n, stream, /*exact=*/true); +} -ggml_tensor * ds4_cluster_broadcast_exact_node(ggml_context * ctx, - ggml_tensor * value, - Ds4ClusterRuntime & rt) { +// 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) { @@ -820,22 +771,21 @@ ggml_tensor * ds4_cluster_broadcast_exact_node(ggml_context * ctx, return nullptr; } if (!ggml_is_contiguous(value)) value = ggml_cont(ctx, value); - return ggml_cluster_allreduce(ctx, value, &broadcast_exact_graph_callback, &rt); + 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) { - 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); - // The in-graph collective node runs its callback in place on a copy of - // `value`; this callback is a broadcast from rank 0. - return ggml_cluster_allreduce(ctx, value, &broadcast_graph_callback, &rt); + return broadcast_node(ctx, value, rt, &broadcast_graph_callback); } bool ds4_cluster_broadcast_host(Ds4ClusterRuntime & rt, @@ -850,7 +800,6 @@ bool ds4_cluster_broadcast_host(Ds4ClusterRuntime & rt, } if (!rt.ensure_scratch(backend, n, err)) return false; cluster::DeviceStream stream = ggml_backend_cuda_get_stream(backend); - trace_coll(rt.rank(), "host-broadcast", n, (const void *) stream); 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; diff --git a/server/src/deepseek4/deepseek4_cluster.h b/server/src/deepseek4/deepseek4_cluster.h index 3de6c84b6a..2ac85e761a 100644 --- a/server/src/deepseek4/deepseek4_cluster.h +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -1,4 +1,5 @@ -// deepseek4_cluster.h - DeepSeek4 expert-parallel cluster runtime. +// deepseek4_cluster.h - DeepSeek4 expert-parallel cluster runtime, after +// maikzz32's lucebox-halo-cluster. // // Every rank loads the dense model (with LUCE_CLUSTER_SLICE_WEIGHTS, only the // attention heads and shared-expert slice it computes) and only its shard of @@ -78,7 +79,7 @@ struct Ds4ClusterRuntime { int slot_validity_n_slots = 0; cluster::ClusterStepTelemetry telemetry; bool trace = false; - // Path 3b: a graph node cannot return an error, so the in-graph + // 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; @@ -165,18 +166,7 @@ void ds4_cluster_mask_routes(const Ds4ClusterRuntime & rt, int route_width, int n_tokens); -// All-reduce a device-resident F32 partial [n_embd, n_tokens] in place on the -// backend stream. Stays stream-ordered (no host sync) unless rt.trace is set, -// in which case it waits with cfg->timeout_ms and logs the layer. Updates -// rt.telemetry and, when given, telemetry->cluster_allreduce_*. -bool ds4_cluster_allreduce_layer(Ds4ClusterRuntime & rt, - ggml_backend_t backend, - ggml_tensor * partial, - int layer, - DeepSeek4StepTelemetry * telemetry, - std::string * err); - -// Path 3b: an all-reduce as an ordinary graph node. `partial` is summed +// 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 @@ -188,14 +178,6 @@ ggml_tensor * ds4_cluster_allreduce_node(ggml_context * ctx, Ds4ClusterRuntime & rt, bool inplace = false); -// The same sum-all-reduce as a raw callback, for graphs built outside the -// deepseek4 layer. server/src/common/dspark_head.cpp needs it to combine the -// vocabulary-sliced logits of the DSpark head, and must not gain a cluster -// dependency to do so, so it takes an ggml_cluster_allreduce_fn and a void* -// rather than the runtime itself. Pass `&rt` as the user pointer: errors land -// in rt.node_error and byte counts in the usual telemetry. -ggml_cluster_allreduce_fn ds4_cluster_allreduce_fn(); - // 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 @@ -221,10 +203,9 @@ bool ds4_cluster_attention_parallel_enabled(); bool ds4_cluster_prefill_device_join_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 (two boxes, 2K prompt in one band: 649-669 -// -> 710-727 tok/s); smaller bands split into parts too short to pay for it -// (512-row bands measured slower). LUCE_CLUSTER_PREFILL_PIPELINE=1 forces it -// on, =0 off. A pure function of the bands, which every rank agrees on. +// 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); diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index b663b8d3d9..df4c49b733 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -608,7 +608,7 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, : bargs.expert_placement; dst = path; } else if (std::strncmp(argv[i], "--cluster-", 10) == 0) { - // Expert-parallel cluster flags (maikzz32/lucebox-halo-cluster). + // Expert-parallel cluster flags. const char * flag = argv[i]; const auto int_value = [&](int & out) { if (i + 1 >= argc) return false; From de9b18516d5149a1f20657da90592159ff7b3abf Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:13:39 +0200 Subject: [PATCH 38/40] refactor(cluster): remove the --cluster-verify-hash probe Nothing outside the unit tests ever called hash_probe, so the flag only travelled in Welcome and was printed. Remove the flag, its validation, the Welcome field, HashProbeMsg, hash_floats/hash_tokens, the mismatch counters and the head hooks' unused config. Welcome changes on the wire, so the protocol version goes to 3. The gather timeout test now calls the control channel directly. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/cluster_config.cpp | 3 - server/src/cluster/cluster_config.h | 3 - server/src/cluster/cluster_control.h | 4 +- server/src/cluster/cluster_decision_hooks.cpp | 96 +------------------ server/src/cluster/cluster_decision_hooks.h | 52 +--------- server/src/cluster/cluster_head_backend.cpp | 17 +--- server/src/cluster/cluster_head_backend.h | 4 - server/src/cluster/cluster_protocol.cpp | 19 +--- server/src/cluster/cluster_protocol.h | 20 +--- server/src/cluster/cluster_worker_main.cpp | 3 - server/src/server/server_main.cpp | 6 -- server/test/unit/test_cluster_unit.cpp | 79 ++++----------- 12 files changed, 33 insertions(+), 273 deletions(-) diff --git a/server/src/cluster/cluster_config.cpp b/server/src/cluster/cluster_config.cpp index 0bbae99e85..a7fa79148b 100644 --- a/server/src/cluster/cluster_config.cpp +++ b/server/src/cluster/cluster_config.cpp @@ -89,9 +89,6 @@ std::string ClusterConfig::validate() const { if (timeout_ms == 0) { return "--cluster-timeout-ms must be positive"; } - if (verify_hash_every < 0) { - return "--cluster-verify-hash must be non-negative, got " + std::to_string(verify_hash_every); - } return {}; } diff --git a/server/src/cluster/cluster_config.h b/server/src/cluster/cluster_config.h index 33b90a0604..4e28676c6b 100644 --- a/server/src/cluster/cluster_config.h +++ b/server/src/cluster/cluster_config.h @@ -75,9 +75,6 @@ struct ClusterConfig { AllreduceDType allreduce_dtype = AllreduceDType::Auto; uint32_t timeout_ms = kClusterDefaultTimeoutMs; // --cluster-timeout-ms - // Debug: cross-rank hidden-state hash every n steps (0 = off). - int verify_hash_every = 0; - // --cluster-selftest: bootstrap, run the collective self-test, exit. bool selftest = false; diff --git a/server/src/cluster/cluster_control.h b/server/src/cluster/cluster_control.h index 3673b3fbe8..d1a3eeea4b 100644 --- a/server/src/cluster/cluster_control.h +++ b/server/src/cluster/cluster_control.h @@ -73,8 +73,8 @@ class ClusterHeadControl { // 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. all - // HashProbe or RequestReport messages); out[rank-1] is filled. + // 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; diff --git a/server/src/cluster/cluster_decision_hooks.cpp b/server/src/cluster/cluster_decision_hooks.cpp index 73056097b1..8ec94fea71 100644 --- a/server/src/cluster/cluster_decision_hooks.cpp +++ b/server/src/cluster/cluster_decision_hooks.cpp @@ -3,8 +3,6 @@ #include "cluster/cluster_decision_hooks.h" #include "cluster/cluster_control.h" -#include -#include namespace luce::cluster { @@ -16,26 +14,9 @@ void set_err(std::string * err, const std::string & msg) { } // namespace -uint64_t hash_floats(const float * data, size_t n) { - if (!data || n == 0) return fnv1a64(nullptr, 0); - return fnv1a64(data, n * sizeof(float)); -} - -uint64_t hash_tokens(const std::vector & tokens) { - if (tokens.empty()) return fnv1a64(nullptr, 0); - return fnv1a64(tokens.data(), tokens.size() * sizeof(int32_t)); -} - // ─── HeadHooks ───────────────────────────────────────────────────────── -HeadHooks::HeadHooks(ClusterHeadControl & control, const HeadHooksConfig & cfg) - : control_(control), cfg_(cfg) {} - -void HeadHooks::reset_counters() { - Ds4ClusterHooks::reset_counters(); - first_mismatch_rank_ = -1; - first_mismatch_step_ = 0; -} +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*/, @@ -92,64 +73,6 @@ bool HeadHooks::decide_accept(uint64_t req, uint32_t step, int32_t & accept, int return true; } -bool HeadHooks::hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, - std::string * err) { - std::vector frames; - const uint64_t t0 = monotonic_us(); - const bool ok = control_.gather(MsgType::HashProbe, cfg_.timeout_ms, frames, err); - counters_.ctrl_wait_us += monotonic_us() - t0; - if (!ok) return false; - counters_.hash_probes++; - - bool mismatch = false; - for (size_t i = 0; i < frames.size(); i++) { - HashProbeMsg probe; - if (!parse_frame(frames[i], probe)) { - set_err(err, "head: malformed HashProbe from rank " + std::to_string(i + 1)); - return false; - } - if (probe.request_id != req || probe.step != step) { - set_err(err, "head: HashProbe from rank " + std::to_string(probe.rank) + - " is for req=" + std::to_string(probe.request_id) + - " step=" + std::to_string(probe.step) + - ", expected req=" + std::to_string(req) + - " step=" + std::to_string(step)); - return false; - } - const bool hc_diff = probe.hc_state_hash != hc_hash; - const bool tok_diff = probe.token_hash != tok_hash; - if (hc_diff || tok_diff) { - mismatch = true; - counters_.hash_mismatches++; - if (first_mismatch_rank_ < 0) { - first_mismatch_rank_ = probe.rank; - first_mismatch_step_ = step; - } - std::fprintf(stderr, - "[cluster] HASH MISMATCH req=%llu step=%u rank=%d: hc %016llx vs head %016llx%s, " - "tokens %016llx vs head %016llx%s%s\n", - (unsigned long long) req, step, probe.rank, - (unsigned long long) probe.hc_state_hash, (unsigned long long) hc_hash, - hc_diff ? " (DIFF)" : "", - (unsigned long long) probe.token_hash, (unsigned long long) tok_hash, - tok_diff ? " (DIFF)" : "", - probe.first_divergent_layer >= 0 - ? (" first_divergent_layer=" + std::to_string(probe.first_divergent_layer)).c_str() - : ""); - } else if (cfg_.trace) { - std::fprintf(stderr, "[cluster] hash ok req=%llu step=%u rank=%d hc=%016llx tok=%016llx\n", - (unsigned long long) req, step, probe.rank, - (unsigned long long) hc_hash, (unsigned long long) tok_hash); - } - } - if (mismatch && cfg_.strict_hash) { - set_err(err, "hash mismatch at step " + std::to_string(step) + - " (first divergent rank " + std::to_string(first_mismatch_rank_) + ")"); - return false; - } - 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); @@ -274,21 +197,4 @@ bool WorkerHooks::decide_accept(uint64_t req, uint32_t step, int32_t & accept, i return true; } -bool WorkerHooks::hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, - std::string * err) { - HashProbeMsg msg; - msg.request_id = req; - msg.step = step; - msg.rank = cfg_.rank; - msg.hc_state_hash = hc_hash; - msg.token_hash = tok_hash; - msg.first_divergent_layer = -1; // per-layer attribution is a LUCE_CLUSTER_TRACE follow-up - const uint64_t t0 = monotonic_us(); - const bool ok = control_.send(make_frame(msg), err); - counters_.ctrl_wait_us += monotonic_us() - t0; - if (!ok) return false; - counters_.hash_probes++; - return true; -} - } // namespace luce::cluster diff --git a/server/src/cluster/cluster_decision_hooks.h b/server/src/cluster/cluster_decision_hooks.h index 70fb9da371..0d951b7696 100644 --- a/server/src/cluster/cluster_decision_hooks.h +++ b/server/src/cluster/cluster_decision_hooks.h @@ -17,9 +17,7 @@ // 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 and, every -// verify_hash_every steps, gathers HashProbe frames from the -// workers and compares them with its own hash. +// 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 @@ -48,8 +46,6 @@ struct HookCounters { uint64_t decisions = 0; // Decision frames sent (head) / received (worker) uint64_t drafts = 0; uint64_t accepts = 0; - uint64_t hash_probes = 0; - uint64_t hash_mismatches = 0; // head only uint64_t backend_ops = 0; uint64_t ctrl_wait_us = 0; // time blocked on the control channel @@ -101,13 +97,6 @@ struct Ds4ClusterHooks { int32_t & accept, int32_t & bonus, uint8_t & flags, std::string * err) = 0; - // Determinism probe (--cluster-verify-hash n). Worker sends its hashes; - // head gathers one HashProbe per worker and compares against its own. - // A mismatch is logged (first divergent rank and step) and counted; it - // is fatal only when strict mode is enabled on the head hooks. - virtual bool hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, - 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; @@ -117,11 +106,8 @@ struct Ds4ClusterHooks { // early-return semantics when no cluster is configured. virtual bool is_cluster() const = 0; - // Steps between hash probes; 0 = off. - virtual int verify_hash_every() const { return 0; } - - virtual const HookCounters & counters() const { return counters_; } - virtual void reset_counters() { counters_.reset(); } + 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 @@ -145,24 +131,16 @@ class LocalHooks final : public Ds4ClusterHooks { bool decide_accept(uint64_t, uint32_t, int32_t &, int32_t &, uint8_t &, std::string *) override { return true; } - bool hash_probe(uint64_t, uint32_t, uint64_t, uint64_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 ───────────────────────────────────────────────────────── -struct HeadHooksConfig { - uint32_t timeout_ms = 30000; // deadline for gather() - int verify_hash_every = 0; - bool strict_hash = false; // mismatch -> return false - bool trace = false; // LUCE_CLUSTER_TRACE -}; - class HeadHooks final : public Ds4ClusterHooks { public: // `control` must outlive the hooks. - HeadHooks(ClusterHeadControl & control, const HeadHooksConfig & cfg); + 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, @@ -171,30 +149,17 @@ class HeadHooks final : public Ds4ClusterHooks { 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 hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, - 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; } - int verify_hash_every() const override { return cfg_.verify_hash_every; } - - // First mismatch seen since reset_counters(); rank -1 = none. - int first_mismatch_rank() const { return first_mismatch_rank_; } - uint32_t first_mismatch_step() const { return first_mismatch_step_; } - void reset_counters() override; private: ClusterHeadControl & control_; - HeadHooksConfig cfg_; - int first_mismatch_rank_ = -1; - uint32_t first_mismatch_step_ = 0; }; // ─── WorkerHooks ─────────────────────────────────────────────────────── struct WorkerHooksConfig { - int rank = -1; - uint32_t timeout_ms = 30000; // deadline for every blocking recv - int verify_hash_every = 0; + uint32_t timeout_ms = 30000; // deadline for every blocking recv }; class WorkerHooks final : public Ds4ClusterHooks { @@ -208,12 +173,9 @@ class WorkerHooks final : public Ds4ClusterHooks { 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 hash_probe(uint64_t req, uint32_t step, uint64_t hc_hash, uint64_t tok_hash, - 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; } - int verify_hash_every() const override { return cfg_.verify_hash_every; } // When a recv expected a Decision/Draft/Accept but got a different frame // (RequestEnd because the head failed the request, Abort, Shutdown), the @@ -237,8 +199,4 @@ class WorkerHooks final : public Ds4ClusterHooks { std::string abort_reason_; }; -// FNV-1a over a float buffer / token vector; shared by backend and tests. -uint64_t hash_floats(const float * data, size_t n); -uint64_t hash_tokens(const std::vector & tokens); - } // namespace luce::cluster diff --git a/server/src/cluster/cluster_head_backend.cpp b/server/src/cluster/cluster_head_backend.cpp index 209d6fcdc6..30305d2b75 100644 --- a/server/src/cluster/cluster_head_backend.cpp +++ b/server/src/cluster/cluster_head_backend.cpp @@ -192,7 +192,6 @@ bool ClusterHeadBackend::init() { welcome.timeout_ms = cfg_.timeout_ms; welcome.allreduce_dtype = (uint8_t) cfg_.allreduce_dtype; welcome.shared_expert = (uint8_t) cfg_.shared_expert; - welcome.verify_hash_every = cfg_.verify_hash_every; welcome.rank_hostnames.resize((size_t) cfg_.size); welcome.rank_hostnames[0] = identity_.hostname; for (const HelloMsg & h : hellos_) { @@ -236,12 +235,7 @@ bool ClusterHeadBackend::init() { std::exit(ok ? 0 : 1); } - HeadHooksConfig hcfg; - hcfg.timeout_ms = cfg_.timeout_ms; - hcfg.verify_hash_every = cfg_.verify_hash_every; - hcfg.strict_hash = false; - hcfg.trace = cluster_env_trace(); - hooks_ = std::make_unique(control_, hcfg); + hooks_ = std::make_unique(control_); // Second set_cluster call: attaches the communicator to the // runtime the factory created before init(); placement-defining fields @@ -265,9 +259,8 @@ bool ClusterHeadBackend::init() { void ClusterHeadBackend::print_ready_banner() const { inner_->print_ready_banner(); - std::printf("[cluster] head ready: size=%d workers=%d rccl=%s verify_hash_every=%d\n", - cfg_.size, control_.n_workers(), rccl_version_string().c_str(), - cfg_.verify_hash_every); + 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); } @@ -482,10 +475,6 @@ bool ClusterHeadBackend::gather_reports(uint64_t request_id, uint32_t steps) { 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; - last_report_.hash_probes = hooks_->counters().hash_probes; - last_report_.hash_mismatches = hooks_->counters().hash_mismatches; - last_report_.first_mismatch_rank = hooks_->first_mismatch_rank(); - last_report_.first_mismatch_step = hooks_->first_mismatch_step(); std::vector frames; std::string err; diff --git a/server/src/cluster/cluster_head_backend.h b/server/src/cluster/cluster_head_backend.h index 79f728e265..60b8e8c25a 100644 --- a/server/src/cluster/cluster_head_backend.h +++ b/server/src/cluster/cluster_head_backend.h @@ -46,10 +46,6 @@ struct ClusterRequestReport { bool complete = false; // every worker reported std::vector ranks; uint64_t head_ctrl_wait_us = 0; - uint64_t hash_probes = 0; - uint64_t hash_mismatches = 0; - int first_mismatch_rank = -1; - uint32_t first_mismatch_step = 0; std::string error; // non-empty when gathering failed }; diff --git a/server/src/cluster/cluster_protocol.cpp b/server/src/cluster/cluster_protocol.cpp index e6b059337f..a7e176daca 100644 --- a/server/src/cluster/cluster_protocol.cpp +++ b/server/src/cluster/cluster_protocol.cpp @@ -230,7 +230,6 @@ const char * msg_type_name(MsgType type) { case MsgType::Decision: return "Decision"; case MsgType::Draft: return "Draft"; case MsgType::Accept: return "Accept"; - case MsgType::HashProbe: return "HashProbe"; case MsgType::RequestReport: return "RequestReport"; case MsgType::BackendOp: return "BackendOp"; case MsgType::RequestEnd: return "RequestEnd"; @@ -264,7 +263,6 @@ void WelcomeMsg::encode(ByteWriter & w) const { w.u32(timeout_ms); w.u8(allreduce_dtype); w.u8(shared_expert); - w.i32(verify_hash_every); w.u32((uint32_t)rank_hostnames.size()); for (const std::string & h : rank_hostnames) w.str(h); } @@ -276,7 +274,7 @@ bool WelcomeMsg::decode(ByteReader & r, WelcomeMsg & out) { 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) || !r.i32(out.verify_hash_every)) { + !r.u8(out.shared_expert)) { return false; } uint32_t n = 0; @@ -384,21 +382,6 @@ bool AcceptMsg::decode(ByteReader & r, AcceptMsg & out) { r.i32(out.bonus) && r.u8(out.flags); } -void HashProbeMsg::encode(ByteWriter & w) const { - w.u64(request_id); - w.u32(step); - w.i32(rank); - w.u64(hc_state_hash); - w.u64(token_hash); - w.i32(first_divergent_layer); -} - -bool HashProbeMsg::decode(ByteReader & r, HashProbeMsg & out) { - return r.u64(out.request_id) && r.u32(out.step) && r.i32(out.rank) && - r.u64(out.hc_state_hash) && r.u64(out.token_hash) && - r.i32(out.first_divergent_layer); -} - void RequestReportMsg::encode(ByteWriter & w) const { w.u64(request_id); w.i32(rank); diff --git a/server/src/cluster/cluster_protocol.h b/server/src/cluster/cluster_protocol.h index b3d3b76f0c..5943d46f85 100644 --- a/server/src/cluster/cluster_protocol.h +++ b/server/src/cluster/cluster_protocol.h @@ -29,7 +29,7 @@ namespace luce::cluster { inline constexpr uint32_t kFrameMagic = 0x4C43424Cu; // "LBCL" -inline constexpr uint16_t kProtocolVersion = 2; +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 @@ -43,7 +43,6 @@ enum class MsgType : uint16_t { 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 - HashProbe = 9, // worker -> head (and head local), determinism check RequestReport = 10, // worker -> head, per-request telemetry BackendOp = 11, // head -> workers, snapshot/park/etc. replicated call RequestEnd = 12, // head -> workers, generation finished (normal path) @@ -148,7 +147,6 @@ struct WelcomeMsg { uint32_t timeout_ms = 0; uint8_t allreduce_dtype = 0; // AllreduceDType as u8 uint8_t shared_expert = 0; // SharedExpertMode as u8 - int32_t verify_hash_every = 0; std::vector rank_hostnames; // index = rank static MsgType type() { return MsgType::Welcome; } @@ -271,20 +269,6 @@ struct AcceptMsg { static bool decode(ByteReader & r, AcceptMsg & out); }; -// Determinism probe (--cluster-verify-hash n). -struct HashProbeMsg { - uint64_t request_id = 0; - uint32_t step = 0; - int32_t rank = -1; - uint64_t hc_state_hash = 0; // FNV-1a over the final-layer HC state bytes - uint64_t token_hash = 0; // FNV-1a over accepted token ids so far - int32_t first_divergent_layer = -1; // filled only under LUCE_CLUSTER_TRACE - - static MsgType type() { return MsgType::HashProbe; } - void encode(ByteWriter & w) const; - static bool decode(ByteReader & r, HashProbeMsg & out); -}; - // Per-request telemetry sent by every worker after RequestEnd; the head // merges it into usage.timings.cluster. struct RequestReportMsg { @@ -364,7 +348,7 @@ bool parse_frame(const Frame & f, Msg & out) { return r.at_end(); } -// FNV-1a 64-bit, used for placement hashes and determinism probes. +// 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_worker_main.cpp b/server/src/cluster/cluster_worker_main.cpp index 7a9240fcb0..a741a16bb3 100644 --- a/server/src/cluster/cluster_worker_main.cpp +++ b/server/src/cluster/cluster_worker_main.cpp @@ -356,7 +356,6 @@ int run_cluster_worker(const ClusterConfig & cfg, // 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.verify_hash_every = welcome.verify_hash_every; 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 " @@ -398,9 +397,7 @@ int run_cluster_worker(const ClusterConfig & cfg, } WorkerHooksConfig hcfg; - hcfg.rank = st.cfg.rank; hcfg.timeout_ms = st.cfg.timeout_ms; - hcfg.verify_hash_every = st.cfg.verify_hash_every; st.hooks = std::make_unique(st.control, hcfg); // Second set_cluster call: attach the communicator to the // runtime the factory created before init(). diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index df4c49b733..1d76bc53b6 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -201,7 +201,6 @@ static void print_usage(const char * prog) { " --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-verify-hash Debug: cross-rank hidden-state hash every n steps\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" @@ -671,11 +670,6 @@ static int parse_model_options(int argc, char ** argv, ModelOptions & model, } } else if (std::strcmp(flag, "--cluster-selftest") == 0) { bargs.cluster.selftest = true; - } else if (std::strcmp(flag, "--cluster-verify-hash") == 0) { - if (!int_value(bargs.cluster.verify_hash_every) || bargs.cluster.verify_hash_every < 0) { - std::fprintf(stderr, "[server] --cluster-verify-hash must be non-negative\n"); - return 2; - } } else { std::fprintf(stderr, "[server] unknown or incomplete cluster flag %s\n", flag); return 2; diff --git a/server/test/unit/test_cluster_unit.cpp b/server/test/unit/test_cluster_unit.cpp index a8e1bf8296..418c745ef1 100644 --- a/server/test/unit/test_cluster_unit.cpp +++ b/server/test/unit/test_cluster_unit.cpp @@ -242,23 +242,22 @@ TEST_CASE(ClusterUnitFixture, Fnv1a64KnownVectors) { // 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)); - - const std::vector toks = {1, 2, 3}; - CHECK_EQUAL(hash_tokens(toks), fnv1a64(toks.data(), toks.size() * sizeof(int32_t))); - CHECK_EQUAL(hash_tokens({}), fnv1a64(nullptr, 0)); - const float fl[2] = {1.0f, -2.0f}; - CHECK_EQUAL(hash_floats(fl, 2), fnv1a64(fl, sizeof(fl))); - CHECK_EQUAL(hash_floats(nullptr, 0), fnv1a64(nullptr, 0)); } 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 (uint16_t t = 1; t <= 13; t++) { - const char * n = msg_type_name((MsgType) t); + 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(), (size_t) 13); + CHECK_EQUAL(names.size(), sizeof(types) / sizeof(types[0])); } // ─── Protocol: messages ───────────────────────────────────────────────── @@ -294,7 +293,6 @@ TEST_CASE(ClusterUnitFixture, HelloWelcomeRoundTrip) { w.timeout_ms = 30000; w.allreduce_dtype = 2; w.shared_expert = 1; - w.verify_hash_every = 8; w.rank_hostnames = {"halo-0", "halo-1", "", "halo-3"}; WelcomeMsg wo; check_roundtrip_and_truncation(w, wo); @@ -304,7 +302,6 @@ TEST_CASE(ClusterUnitFixture, HelloWelcomeRoundTrip) { CHECK_EQUAL(wo.timeout_ms, 30000u); CHECK_EQUAL(wo.allreduce_dtype, 2); CHECK_EQUAL(wo.shared_expert, 1); - CHECK_EQUAL(wo.verify_hash_every, 8); 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")); @@ -468,21 +465,7 @@ TEST_CASE(ClusterUnitFixture, DecisionDraftAcceptRoundTrip) { CHECK_EQUAL(ac_out.flags, (uint8_t) kDecisionStop); } -TEST_CASE(ClusterUnitFixture, HashProbeReportBackendOpRequestEndRoundTrip) { - HashProbeMsg hp; - hp.request_id = 1; - hp.step = 8; - hp.rank = 2; - hp.hc_state_hash = 0xAAAAAAAAAAAAAAAAull; - hp.token_hash = 0x5555555555555555ull; - hp.first_divergent_layer = 17; - HashProbeMsg hp_out; - check_roundtrip_and_truncation(hp, hp_out); - CHECK_EQUAL(hp_out.rank, 2); - CHECK_EQUAL(hp_out.hc_state_hash, hp.hc_state_hash); - CHECK_EQUAL(hp_out.token_hash, hp.token_hash); - CHECK_EQUAL(hp_out.first_divergent_layer, 17); - +TEST_CASE(ClusterUnitFixture, ReportBackendOpRequestEndRoundTrip) { RequestReportMsg rr; rr.request_id = 1; rr.rank = 1; @@ -886,7 +869,6 @@ TEST_CASE(ClusterUnitFixture, LocalHooksAreTransparent) { LocalHooks h; CHECK(h.is_head()); CHECK_FALSE(h.is_cluster()); - CHECK_EQUAL(h.verify_hash_every(), 0); int32_t token = 123; uint8_t flags = kDecisionEos; luce::common::SamplerCfg cfg; @@ -902,7 +884,6 @@ TEST_CASE(ClusterUnitFixture, LocalHooksAreTransparent) { CHECK(h.decide_accept(1, 0, accept, bonus, flags, &err)); CHECK_EQUAL(accept, 2); CHECK_EQUAL(bonus, 9); - CHECK(h.hash_probe(1, 0, 1, 2, &err)); CHECK(h.backend_op(BackendOpMsg{}, &err)); h.set_current_request(77); CHECK_EQUAL(h.current_request(), 77u); @@ -950,7 +931,6 @@ struct LoopbackPair { WelcomeMsg send; send.size = 2; send.timeout_ms = 5000; - send.verify_hash_every = 4; send.rank_hostnames = {"head", "worker"}; std::atomic w_ok{false}; std::thread wt([&] { @@ -975,7 +955,7 @@ TEST_CASE(ClusterUnitFixture, ControlHandshakeAcceptsMatchingIdentity) { 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.verify_hash_every, 4); + CHECK_EQUAL(lp.welcome.timeout_ms, 5000u); CHECK_EQUAL(lp.head.n_workers(), 1); CHECK(lp.head.all_alive()); CHECK(lp.worker.alive()); @@ -1013,18 +993,12 @@ TEST_CASE(ClusterUnitFixture, HeadAndWorkerHooksOverLoopback) { if (lp.port == 0) SKIP("no loopback port available"); REQUIRE(lp.handshake(lp.worker_id)); - HeadHooksConfig hc; - hc.timeout_ms = 5000; - hc.verify_hash_every = 4; - HeadHooks head(lp.head, hc); + HeadHooks head(lp.head); WorkerHooksConfig wc; - wc.rank = 1; wc.timeout_ms = 5000; - wc.verify_hash_every = 4; WorkerHooks worker(lp.worker, wc); CHECK(head.is_head()); CHECK(head.is_cluster()); CHECK_FALSE(worker.is_head()); CHECK(worker.is_cluster()); - CHECK_EQUAL(worker.verify_hash_every(), 4); const uint64_t req = 42; head.set_current_request(req); @@ -1032,18 +1006,16 @@ TEST_CASE(ClusterUnitFixture, HeadAndWorkerHooksOverLoopback) { luce::common::SamplerCfg scfg; const float logits[4] = {0.1f, 0.9f, 0.2f, 0.3f}; - SECTION("decision, draft, accept, hash probe and backend op") { + 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, w_hp = false, w_hp2 = false; + 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); - w_hp = worker.hash_probe(req, 4, 0xAAAA, 0xBBBB, &w_err); // matches - w_hp2 = worker.hash_probe(req, 8, 0xAAAA, 0xCCCC, &w_err); // token hash differs }); std::string h_err; @@ -1054,25 +1026,15 @@ TEST_CASE(ClusterUnitFixture, HeadAndWorkerHooksOverLoopback) { 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)); - REQUIRE(head.hash_probe(req, 4, 0xAAAA, 0xBBBB, &h_err)); - CHECK_EQUAL(head.counters().hash_mismatches, 0u); - CHECK_EQUAL(head.first_mismatch_rank(), -1); - REQUIRE(head.hash_probe(req, 8, 0xAAAA, 0xBBBB, &h_err)); // non-strict: logged, not fatal - CHECK_EQUAL(head.counters().hash_mismatches, 1u); - CHECK_EQUAL(head.first_mismatch_rank(), 1); - CHECK_EQUAL(head.first_mismatch_step(), 8u); 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(w_hp); CHECK(w_hp2); CHECK_EQUAL(head.counters().decisions, 1u); CHECK_EQUAL(head.counters().drafts, 1u); CHECK_EQUAL(head.counters().accepts, 1u); - CHECK_EQUAL(head.counters().hash_probes, 2u); CHECK_EQUAL(worker.counters().decisions, 1u); - CHECK_EQUAL(worker.counters().hash_probes, 2u); // Backend op reaches the worker as a plain frame (the worker main // loop executes it; WorkerHooks::backend_op is a no-op). @@ -1162,9 +1124,8 @@ TEST_CASE(ClusterUnitFixture, EmptyDraftBlockIsTheEndMarker) { if (lp.port == 0) SKIP("no loopback port available"); REQUIRE(lp.handshake(lp.worker_id)); - HeadHooksConfig hc; hc.timeout_ms = 5000; - HeadHooks head(lp.head, hc); - WorkerHooksConfig wc; wc.rank = 1; wc.timeout_ms = 5000; + HeadHooks head(lp.head); + WorkerHooksConfig wc; wc.timeout_ms = 5000; WorkerHooks worker(lp.worker, wc); const uint64_t req = 9; @@ -1204,18 +1165,16 @@ TEST_CASE(ClusterUnitFixture, EmptyDraftMsgRoundTrip) { CHECK(got.tokens.empty()); } -TEST_CASE(ClusterUnitFixture, HeadHooksFailWhenWorkerGone) { +TEST_CASE(ClusterUnitFixture, HeadGatherFailsWhenWorkerGone) { LoopbackPair lp; if (lp.port == 0) SKIP("no loopback port available"); REQUIRE(lp.handshake(lp.worker_id)); lp.worker.close(); - HeadHooksConfig hc; - hc.timeout_ms = 500; - HeadHooks head(lp.head, hc); + std::vector frames; std::string err; // A gather with no live worker must return false within the deadline. - CHECK_FALSE(head.hash_probe(1, 0, 1, 2, &err)); + CHECK_FALSE(lp.head.gather(MsgType::RequestReport, 500, frames, &err)); CHECK_FALSE(err.empty()); lp.head.close(); } From effc0c8d0c37c3e77ea38038734af6c5e5a72f58 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:15:48 +0200 Subject: [PATCH 39/40] refactor(cluster): drop unused head state, the HTTP view header and stale switches The head backend's introspection accessors, its fault reason and its exit-on-abort toggle had no readers or callers, and common/cluster_view.h declared types nothing used. Cluster ranks always slice their dense weights and keep the device join in prefill now, so LUCE_CLUSTER_SLICE_WEIGHTS, LUCE_DS4_SHEXP_SLICE_VIEW and LUCE_CLUSTER_PREFILL_DEVICE_JOIN go from the cluster defaults and the code. Plan labels and a reference to a launcher script that does not exist are gone from the comments. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/cluster_head_backend.cpp | 25 +++----- server/src/cluster/cluster_head_backend.h | 29 ++------- server/src/common/cluster_view.h | 70 --------------------- server/src/deepseek4/deepseek4_cluster.cpp | 10 +-- server/src/deepseek4/deepseek4_cluster.h | 8 +-- server/src/server/launch_profiles.h | 2 - 6 files changed, 14 insertions(+), 130 deletions(-) delete mode 100644 server/src/common/cluster_view.h diff --git a/server/src/cluster/cluster_head_backend.cpp b/server/src/cluster/cluster_head_backend.cpp index 30305d2b75..89178568d6 100644 --- a/server/src/cluster/cluster_head_backend.cpp +++ b/server/src/cluster/cluster_head_backend.cpp @@ -160,8 +160,8 @@ bool ClusterHeadBackend::init() { 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()); - // Launcher contract (scripts/cluster/launch_cluster.sh waits for this - // exact stdout prefix before it starts the workers). + // 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); @@ -269,10 +269,6 @@ void ClusterHeadBackend::print_ready_banner() const { void ClusterHeadBackend::fail_cluster(const std::string & reason, int code, uint64_t request_id) { if (fatal_) return; fatal_ = true; - { - std::lock_guard lk(fault_mu_); - if (fault_reason_.empty()) fault_reason_ = reason; - } std::fprintf(stderr, "[cluster] head: FATAL (%d): %s\n", code, reason.c_str()); AbortMsg abort; abort.rank = 0; @@ -282,24 +278,18 @@ void ClusterHeadBackend::fail_cluster(const std::string & reason, int code, uint control_.broadcast(make_frame(abort), nullptr); if (comm_) comm_->abort(); control_.close(); - if (exit_on_abort_) { - std::fprintf(stderr, "[cluster] head: exiting so the supervisor restarts all ranks\n"); - std::fflush(nullptr); - std::_Exit(3); - } + std::fprintf(stderr, "[cluster] head: exiting so the supervisor restarts all ranks\n"); + std::fflush(nullptr); + std::_Exit(3); } -// ─── M4: peer-failure watchdog ────────────────────────────────────────── +// ─── 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; - { - std::lock_guard lk(fault_mu_); - if (fault_reason_.empty()) fault_reason_ = reason; - } // 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. @@ -310,7 +300,6 @@ void ClusterHeadBackend::request_fatal_exit(const std::string & reason, int code abort.request_id = request_id; abort.reason = reason; control_.broadcast(make_frame(abort), nullptr); - if (!exit_on_abort_) return; // 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. @@ -569,7 +558,7 @@ common::ModelBackend::SnapshotRef ClusterHeadBackend::snapshot_ref(int slot) con 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 (M4: ship it). + // 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; diff --git a/server/src/cluster/cluster_head_backend.h b/server/src/cluster/cluster_head_backend.h index 60b8e8c25a..1ba9438176 100644 --- a/server/src/cluster/cluster_head_backend.h +++ b/server/src/cluster/cluster_head_backend.h @@ -7,12 +7,11 @@ // 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 into `last_report()`. +// one RequestReport per worker. // -// Failure policy (M1): any control-channel or collective failure broadcasts -// Abort, aborts the communicator and, unless exit_on_abort is false, ends -// the process with a non-zero status so a supervisor restarts all ranks -// together. Communicator re-initialization is M4. +// 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 @@ -21,8 +20,6 @@ #include "cluster/cluster_control.h" #include "cluster/cluster_decision_hooks.h" #include "cluster/cluster_protocol.h" -#include "common/cluster_view.h" -#include "common/model_backend.h" #include "common/cluster_participant.h" #include "common/model_backend.h" @@ -116,18 +113,6 @@ class ClusterHeadBackend final : public common::ModelBackend { bool spark_bootstrap_finalize(const std::string & profile_path) override; void shutdown() override; - // ── Cluster introspection ──────────────────────────────────────── - // Plain snapshots for the HTTP layer (common/cluster_view.h). - // telemetry describes the request that just finished and is read on the - // request thread; props is constant after init(). - - const ClusterRequestReport & last_report() const { return last_report_; } - const ClusterConfig & config() const { return cfg_; } - const std::vector & worker_identities() const { return hellos_; } - IClusterComm * comm() { return comm_.get(); } - bool fatal() const { return fatal_; } - void set_exit_on_abort(bool v) { exit_on_abort_ = v; } - // 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. @@ -180,11 +165,6 @@ class ClusterHeadBackend final : public common::ModelBackend { std::atomic watchdog_stop_{false}; std::atomic request_in_flight_{false}; std::atomic exit_scheduled_{false}; - // Why the cluster died, set by whichever thread noticed. Read by - // cluster_request_telemetry so the failed request's response carries it - // in usage.timings.cluster.error instead of looking like an empty answer. - mutable std::mutex fault_mu_; - std::string fault_reason_; ClusterHeadControl control_; std::unique_ptr comm_; std::unique_ptr hooks_; @@ -194,7 +174,6 @@ class ClusterHeadBackend final : public common::ModelBackend { ClusterRequestReport last_report_; bool initialized_ = false; bool fatal_ = false; - bool exit_on_abort_ = true; bool shut_down_ = false; }; diff --git a/server/src/common/cluster_view.h b/server/src/common/cluster_view.h deleted file mode 100644 index 4cd54cac38..0000000000 --- a/server/src/common/cluster_view.h +++ /dev/null @@ -1,70 +0,0 @@ -// cluster_view.h - dependency-free snapshots of cluster state for the HTTP -// layer. -// -// server/src/server/ has no business knowing about RCCL, the control channel -// or the DeepSeek4 runtime, and it must keep compiling with -// -DLUCE_CLUSTER=OFF. So the cluster backend copies what the API should -// expose into these plain structs, exactly the way get_routing_stats() hands -// out routing telemetry, and ModelBackend answers "no cluster" by default. -// -// Threading: cluster_request_telemetry() is filled by the rank-0 decorator at -// the end of a request and must be read on the same thread that ran it (the -// request worker), which is where GenTimings is assembled. cluster_props() is -// constant after bootstrap and is safe to read from any thread. - -#pragma once - -#include -#include -#include - -namespace luce::common { - -// One rank's contribution to the request that just finished. -struct ClusterRankTiming { - int rank = 0; - uint32_t steps = 0; - uint64_t compute_us = 0; // local prefill+decode wall time - uint64_t allreduce_calls = 0; - uint64_t allreduce_bytes = 0; - uint64_t allreduce_wait_us = 0; // host time blocked on collectives; 0 with - // the in-graph all-reduce - uint64_t ctrl_wait_us = 0; // host time blocked on the control channel - uint64_t peak_device_bytes = 0; // high-water mark, sampled at request ends -}; - -// Per-request cluster telemetry. `ranks[0]` is the head. -struct ClusterTelemetryView { - bool active = false; // false: not a cluster run - int size = 0; - uint64_t request_id = 0; - uint32_t steps = 0; - bool complete = false; // every worker reported - uint64_t head_ctrl_wait_us = 0; - uint64_t hash_probes = 0; - uint64_t hash_mismatches = 0; - int first_mismatch_rank = -1; - uint32_t first_mismatch_step = 0; - std::string error; // non-empty when the gather failed - std::vector ranks; -}; - -// Constant after bootstrap: what this cluster is, for /props. -struct ClusterPropsView { - bool active = false; - int size = 0; - int rank = 0; - std::string ifname; - std::string ib_hca; - int gid_index = 0; - std::string placement; // uniform | balanced | - uint64_t placement_hash = 0; - int replicate_hot = 0; - std::string shared_expert; // replicate | shard | rank0 - std::string allreduce_dtype; // f32 | bf16 | auto - bool ingraph_allreduce = false; // collectives run inside the graph - uint64_t resident_expert_bytes = 0; // this rank's expert shard - uint32_t timeout_ms = 0; -}; - -} // namespace luce::common diff --git a/server/src/deepseek4/deepseek4_cluster.cpp b/server/src/deepseek4/deepseek4_cluster.cpp index 65d3a52a45..2cacfaf79a 100644 --- a/server/src/deepseek4/deepseek4_cluster.cpp +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -459,7 +459,7 @@ bool ds4_cluster_ingraph_allreduce_enabled() { } bool ds4_cluster_prefill_pipeline_enabled(const std::vector & bands) { - if (!ds4_cluster_prefill_device_join_enabled() || bands.empty()) return false; + 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; @@ -469,14 +469,6 @@ bool ds4_cluster_prefill_pipeline_enabled(const std::vector & bands) { [](int rows) { return rows >= kDs4ClusterPipelineMinBand; }); } -bool ds4_cluster_prefill_device_join_enabled() { - static const bool enabled = [] { - const char * v = std::getenv("LUCE_CLUSTER_PREFILL_DEVICE_JOIN"); - return !v || !*v || std::strcmp(v, "0") != 0; - }(); - return enabled; -} - bool ds4_cluster_attention_parallel_enabled() { // Off unless set; cluster mode sets it (cluster_launch_env), and =0 keeps // attention replicated. diff --git a/server/src/deepseek4/deepseek4_cluster.h b/server/src/deepseek4/deepseek4_cluster.h index 2ac85e761a..8e526c5b9c 100644 --- a/server/src/deepseek4/deepseek4_cluster.h +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -1,9 +1,8 @@ // deepseek4_cluster.h - DeepSeek4 expert-parallel cluster runtime, after // maikzz32's lucebox-halo-cluster. // -// Every rank loads the dense model (with LUCE_CLUSTER_SLICE_WEIGHTS, only the -// attention heads and shared-expert slice it computes) and only its shard of -// routed experts (owned + replicated, see cluster_expert_placement.h). Per MoE +// 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:: @@ -198,9 +197,6 @@ ggml_tensor * ds4_cluster_hybrid_exchange_nodes(ggml_context * ctx, // (on in cluster mode; =0 keeps attention replicated). Cached after the first // call. bool ds4_cluster_attention_parallel_enabled(); -// LUCE_CLUSTER_PREFILL_DEVICE_JOIN (default on): cluster prefill keeps the -// device-resident hot+cold join and reduces in the HC-post graph. -bool ds4_cluster_prefill_device_join_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 diff --git a/server/src/server/launch_profiles.h b/server/src/server/launch_profiles.h index 9b408ba1b9..f2356ab852 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -276,7 +276,6 @@ inline const std::vector & cluster_launch_env() { {"LUCE_CLUSTER_DRAFT_VOCAB_SPLIT", "1"}, {"LUCE_CLUSTER_DRAFT_SPLIT", "1"}, {"LUCE_CLUSTER_WORKER_SKIP_DRAFT", "1"}, - {"LUCE_CLUSTER_SLICE_WEIGHTS", "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 @@ -289,7 +288,6 @@ inline const std::vector & cluster_launch_env() { {"LUCE_DS4_SCHED_DEDICATED_COPIES", "1"}, {"LUCE_DS4_HC_BOUNDARY_SPLIT", "1"}, {"LUCE_DS4_ENGRAM_PRECOMPUTE", "2560"}, - {"LUCE_DS4_SHEXP_SLICE_VIEW", "1"}, {"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 From 1996c8168903328b0c0d7601e1e73a75a921865b Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:22:11 +0200 Subject: [PATCH 40/40] refactor(cluster): one decode exchange path in the fast reduce Cluster mode always ran the zero-copy, bf16-wire, lean exchange with in-kernel publishes, one hybrid export and a DMA import. The other paths behind LUCE_CLUSTER_FAST_REDUCE_ZC, _BF16, _LEAN, _KPUB, _KERNEL_FLAGS, LUCE_CLUSTER_HYBRID_ONE_EXPORT and LUCE_CLUSTER_HYBRID_IMPORT_KERNEL were never taken by default, so they go with their switches: the copy-engine add and its scratch, the kernel-form flags and their words, the done flags and the second hybrid export. The fast reduce now refuses any size but two, so other sizes keep RCCL as documented. The MTU and stall settings are read once at init instead of in function-local statics, and the accessors nothing read are gone. The default path's arithmetic is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) --- server/src/cluster/fast_reduce.cpp | 431 ++++-------------- server/src/cluster/fast_reduce.h | 44 +- server/src/cluster/fast_reduce_kernels.cu | 256 ++--------- server/src/deepseek4/deepseek4_cluster.cpp | 39 +- server/src/deepseek4/deepseek4_cluster.h | 5 +- .../src/deepseek4/deepseek4_fused_verify.inc | 11 +- server/src/server/launch_profiles.h | 4 - 7 files changed, 171 insertions(+), 619 deletions(-) diff --git a/server/src/cluster/fast_reduce.cpp b/server/src/cluster/fast_reduce.cpp index 8242ae1cc2..e851b85c0b 100644 --- a/server/src/cluster/fast_reduce.cpp +++ b/server/src/cluster/fast_reduce.cpp @@ -33,32 +33,15 @@ namespace luce::cluster { -// The kernels live in fast_reduce_kernels.cu (LANGUAGE HIP). -void fast_reduce_launch_pack(const float * src, void * dst, int n, int bf16, hipStream_t stream); -void fast_reduce_launch_pack_pub(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_hyb_combine_pub(const float * cold, const float * hot, const float * shared, void * own_row, - int n, int bf16, uint32_t * flag, uint32_t value, uint32_t * counter, - hipStream_t stream); -void fast_reduce_launch_hyb_final_pub(const void * own_row, const void * peer_row, float * dst, int n, int bf16, - uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream); -void fast_reduce_launch_add_zc(float * dst, const void * peer, int n, int bf16, hipStream_t stream); -void fast_reduce_launch_hyb_combine(const float * cold, const float * hot, const float * shared, void * own_row, - int n, int bf16, hipStream_t stream); -void fast_reduce_launch_hyb_final(const void * own_row, const void * peer_row, float * dst, int n, int bf16, - hipStream_t stream); -void fast_reduce_launch_set_flag(uint32_t * pub_flag, uint32_t * slot_done, - uint32_t * progress, uint32_t seq, - uint32_t slot_span, uint64_t spin_limit, - hipStream_t stream); -void fast_reduce_launch_wait_flags(uint32_t * const * peer_flags, - uint32_t * timed_out, uint32_t * progress, - uint32_t seq, int n_peers, - uint64_t spin_limit, hipStream_t stream); -void fast_reduce_launch_add(float * dst, const float * scratch, int n_peers, - int n, int stride, uint32_t * slot_done, - uint32_t * progress, uint32_t seq, - hipStream_t stream); +// 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 { @@ -196,56 +179,12 @@ bool exchange(const FastReduce::Config & cfg, const Endpoint & mine, } // namespace -// A switch is on when set to anything but empty or "0". Cluster mode sets the -// measured configuration's switches by default (launch_profiles.h), so "0" is -// how an operator turns one off. +// 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; } -// Two ranks with LUCE_CLUSTER_FAST_REDUCE_ZC=1 (and the stream flags) take the -// zero-copy path: pack into the pinned row, add the peer's row straight out of -// pinned memory. It needs no device scratch. -static bool fr_zero_copy(int size) { - static const bool zc = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_ZC") && - !fr_env_on("LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS"); - return zc && size == 2; -} - -// A row crosses the cable as bf16 exactly when the zero-copy path packs it so -// (LUCE_CLUSTER_FAST_REDUCE_BF16); the copy-engine path publishes f32. The -// progress thread sizes every RDMA write by the same rule. -static bool fr_bf16_wire(int size) { - static const bool bf16 = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_BF16"); - return bf16 && fr_zero_copy(size); -} - -// LUCE_CLUSTER_FAST_REDUCE_LEAN=1 drops the per-reduction slot-reuse waits and -// done writes on every stream-flag path. Neither carries data: 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 256 reductions earlier, and the GPU's own stream orders its -// writes behind its reads. On a discrete GPU each of these stream operations -// is a blit kernel with its own dispatch gap, four per layer. -static bool fast_reduce_lean() { - static const bool lean = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_LEAN"); - return lean; -} - -// The kernel that writes an exchange row also publishes its flag (one -// system-scope release by its grid's last block) instead of a -// hipStreamWriteValue32 after it, which saves a stream operation per row and -// gives the same bits. LUCE_CLUSTER_FAST_REDUCE_KPUB=0 goes back to the stream -// writes. -static bool fast_reduce_kpub() { - static const bool on = [] { - const char * v = std::getenv("LUCE_CLUSTER_FAST_REDUCE_KPUB"); - return !(v && v[0] == '0' && v[1] == '\0'); - }(); - return on; -} - 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 @@ -279,26 +218,9 @@ struct FastReduceImpl { // 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] done + // | [slots] pub | [slots] hybrid export | [slots] hybrid final uint32_t * flag_dev = nullptr; - // Device-visible pointer table for the flags, one row per slot. The - // payloads need no such table: they are moved by the copy engine, which - // takes its addresses from the host. - uint32_t ** peer_flags = nullptr; // [slots][size-1] - uint32_t ** peer_flags_dev = nullptr; - - // Device memory the peers' partials are copied into before they are added: - // the pinned staging buffer is uncached for the GPU, so adding straight - // out of it is slow. - float * scratch = nullptr; // [slots][size-1][max_elems] - - uint32_t * timed_out_host = nullptr; - uint32_t * timed_out_dev = nullptr; - // seq*4 + {1 published, 2 waiting, 3 done}: which reduction the GPU is in - // and what it is doing there. - uint32_t * progress_host = nullptr; - uint32_t * progress_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 @@ -312,7 +234,7 @@ struct FastReduceImpl { // 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 = hot | part 1 = shared | 2 = final][slots][hyb_elems]. + // 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; @@ -339,17 +261,13 @@ struct FastReduceImpl { size_t pub_off(int slot) const { return 2u * (size_t) cfg.slots * (size_t) cfg.size + (size_t) slot; } - size_t done_off(int slot) const { - return pub_off(0) + (size_t) cfg.slots + (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); } - uint32_t * done_flag(int slot) { return flag_host + done_off(slot); } - size_t hybh_off(int slot) const { return done_off(0) + (size_t) cfg.slots + (size_t) 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 + 4u * (size_t) cfg.slots; + return 2u * (size_t) cfg.size * (size_t) cfg.slots + 3u * (size_t) cfg.slots; } }; @@ -376,15 +294,15 @@ uint32_t * FastReduceImpl::pub_counter(hipStream_t stm, int slot, int kind) { FastReduce::~FastReduce() { shutdown(); } bool FastReduce::ok() const { return p_ && p_->up; } -uint64_t FastReduce::submitted() const { return p_ ? p_->seq.load() : 0; } -uint64_t FastReduce::timed_out() const { - return (p_ && p_->timed_out_host) ? *p_->timed_out_host : 0; -} 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; @@ -426,20 +344,10 @@ bool FastReduce::init(const Config & cfg, std::string * err) { return bail("hipHostMalloc for the flag region failed"); } std::memset(s.flag_host, 0, s.flag_count() * sizeof(uint32_t)); - if (hipHostMalloc((void **) &s.timed_out_host, sizeof(uint32_t) * 16, - hipHostMallocDefault) != hipSuccess) { - return bail("hipHostMalloc for the timeout word failed"); - } - *s.timed_out_host = 0; - if (hipHostMalloc((void **) &s.progress_host, sizeof(uint32_t) * 16, - hipHostMallocDefault) != hipSuccess) { - return bail("hipHostMalloc for the progress word failed"); - } - *s.progress_host = 0; - if (fr_env_on("LUCE_CLUSTER_HYBRID_EXCHANGE") && cfg.size == 2) { + 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) 3 * (size_t) cfg.slots * (size_t) s.hyb_elems * sizeof(float); + 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"); @@ -453,9 +361,7 @@ bool FastReduce::init(const Config & cfg, std::string * err) { } if (hipHostGetDevicePointer((void **) &s.data_dev, s.data_host, 0) != hipSuccess || - hipHostGetDevicePointer((void **) &s.flag_dev, s.flag_host, 0) != hipSuccess || - hipHostGetDevicePointer((void **) &s.timed_out_dev, s.timed_out_host, 0) != hipSuccess || - hipHostGetDevicePointer((void **) &s.progress_dev, s.progress_host, 0) != hipSuccess) { + hipHostGetDevicePointer((void **) &s.flag_dev, s.flag_host, 0) != hipSuccess) { return bail("hipHostGetDevicePointer failed"); } @@ -517,15 +423,16 @@ bool FastReduce::init(const Config & cfg, std::string * err) { 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; - // LUCE_CLUSTER_FAST_REDUCE_MTU=1024|2048|4096 (default 4096: the direct ConnectX link runs 4096). - { - static const int mtu_env = std::getenv("LUCE_CLUSTER_FAST_REDUCE_MTU") ? std::atoi(std::getenv("LUCE_CLUSTER_FAST_REDUCE_MTU")) : 4096; - a.path_mtu = mtu_env >= 4096 ? IBV_MTU_4096 : (mtu_env >= 2048 ? IBV_MTU_2048 : IBV_MTU_1024); - } + 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; @@ -555,36 +462,16 @@ bool FastReduce::init(const Config & cfg, std::string * err) { } } - // ── pointer tables, one row per slot ────────────────────────────────── - const int n_peers = cfg.size - 1; - const size_t rows = (size_t) cfg.slots * (size_t) n_peers; - if (hipHostMalloc((void **) &s.peer_flags, rows * sizeof(uint32_t *), - hipHostMallocDefault) != hipSuccess) { - return bail("hipHostMalloc for the flag table failed"); - } - for (int slot = 0; slot < cfg.slots; ++slot) { - int k = 0; - for (int r = 0; r < cfg.size; ++r) { - if (r == cfg.rank) continue; - s.peer_flags[(size_t) slot * (size_t) n_peers + (size_t) k] = - s.flag_dev + s.visible_off(slot, r); - ++k; - } - } - if (hipHostGetDevicePointer((void **) &s.peer_flags_dev, s.peer_flags, 0) != hipSuccess) { - return bail("hipHostGetDevicePointer for the flag table failed"); - } - if (!fr_zero_copy(cfg.size) && - hipMalloc((void **) &s.scratch, - rows * (size_t) cfg.max_elems * sizeof(float)) != hipSuccess) { - return bail("hipMalloc for the peer scratch 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]() { + 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, @@ -607,16 +494,15 @@ bool FastReduce::init(const Config & cfg, std::string * err) { // 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(); - bool drained = false; // the lean paths write no done flags: wait the full time - while (!drained && std::chrono::duration(std::chrono::steady_clock::now() - drain_t0).count() < 15.0) { + 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); // late submits too - drained = *(volatile uint32_t *) s.done_flag((int) (last % (uint64_t) s.cfg.slots)) == (uint32_t) last; - if (!drained) std::this_thread::sleep_for(std::chrono::milliseconds(20)); + 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: %s, exiting\n", s.cfg.rank, - drained ? "GPU drained" : "GPU released"); + std::fprintf(stderr, "[fast-reduce] rank %d: GPU released, exiting\n", s.cfg.rank); std::fflush(nullptr); std::_Exit(4); }; @@ -636,8 +522,6 @@ bool FastReduce::init(const Config & cfg, std::string * err) { // server waits here between requests and that time is no stall. bool handed_out = false; auto handed_t0 = wait_t0; - static const double stall_s = std::getenv("LUCE_CLUSTER_FAST_REDUCE_STALL_S") - ? std::atof(std::getenv("LUCE_CLUSTER_FAST_REDUCE_STALL_S")) : 20.0; while (*pub != (uint32_t) next) { if (s.stop.load(std::memory_order_relaxed)) return; if (++spins < 200000ull) continue; @@ -661,20 +545,16 @@ bool FastReduce::init(const Config & cfg, std::string * err) { } if (next <= 1) continue; // The GPU published reduction p = next-1 and is waiting on the - // peers' flags for it. Done long ago is the normal idle case; - // published but not done means a peer never delivered. + // 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); - if (*(volatile uint32_t *) s.done_flag(ps) == (uint32_t) p) continue; - if (fast_reduce_lean()) { - bool arrived = true; - for (int r = 0; r < s.cfg.size; ++r) { - if (r == s.cfg.rank) continue; - const uint32_t v = *(volatile uint32_t *) s.visible_flag(ps, r); - if (v != (uint32_t) p) arrived = false; - } - if (arrived) continue; + 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) { @@ -687,7 +567,7 @@ bool FastReduce::init(const Config & cfg, std::string * err) { } std::fprintf(stderr, "[fast-reduce] rank %d: reduction %llu published %.1f s ago, " - "not done; peer flags%s (want %llu); host submitted %llu\n", + "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()); @@ -711,7 +591,7 @@ bool FastReduce::init(const Config & cfg, std::string * err) { 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 = fr_bf16_wire(s.cfg.size) && !exact_row; + const bool bf16_row = !exact_row; sge.length = (uint32_t) ((size_t) n * (bf16_row ? 2u : sizeof(float))); sge.lkey = s.mr_data->lkey; @@ -779,11 +659,8 @@ bool FastReduce::init(const Config & cfg, std::string * err) { // 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), " - "gpu timeouts up to seq %u\n", - (unsigned long long) next, n, - *(volatile uint32_t *) s.timed_out_host); + std::fprintf(stderr, "[fast-reduce] seq %llu delivered (%d floats)\n", + (unsigned long long) next, n); } ++next; } @@ -804,123 +681,31 @@ bool FastReduce::submit(float * data, size_t n, void * stream, bool exact) { // 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; - const int n_peers = s.cfg.size - 1; - auto stm = (hipStream_t) stream; - - // Loads, not cycles: roughly 100 ms of system-scope reads. What it waits - // on takes about fourteen microseconds, so this only fires when something - // is wrong -- and when it does it has to fire, not spin forever. - constexpr uint64_t kSpin = 200000ull; - - // Publish and wait are stream operations, not kernels. + // 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. // - // A kernel cannot do this cheaply on gfx1151. Raising a flag the host will - // see needs a system-scope release, which writes the L2 back; spinning on - // one the host raised needs a system-scope acquire, which invalidates it. - // Either way the weights the next layer is about to read are gone. - // Weakening the ordering does not help: relaxed is cheap and the host then - // does not see the flag at all. + // 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. // - // hipStreamWriteValue32 and hipStreamWaitValue32 are handled by the - // command processor instead, so the shader cores and their caches are not - // involved at all -- which is what this needs and what a kernel cannot be. - // LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS=1 restores the kernel form for - // comparison. - static const bool kernel_flags = fr_env_on("LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS"); - - // Do not write into a slot the progress thread has not finished with. The - // kernel form waited on this; the stream form has to as well, or the GPU -- - // which runs a whole step's reductions ahead of the host -- laps the ring - // and overwrites a payload that has not been sent. That is what made the - // third request of a run stop rather than the first. - // The previous user of this slot was seq - slots, and it leaves exactly that - // value behind. Asking for anything higher is a wait nothing can satisfy. - if (!kernel_flags && !fast_reduce_lean() && seq > (uint64_t) s.cfg.slots) { - const uint32_t need = (uint32_t) (seq - (uint64_t) s.cfg.slots); - if (hipStreamWaitValue32(stm, s.flag_dev + s.done_off(slot), need, - hipStreamWaitValueGte, 0xffffffffu) != hipSuccess) { - return false; - } - } - - // Zero-copy (discrete GPU): pack straight into the pinned row, publish, wait - // for the peer, add its row straight out of pinned memory, mark the slot done. - const int wire_bf16 = !exact && fr_bf16_wire(s.cfg.size) ? 1 : 0; - if (fr_zero_copy(s.cfg.size)) { - float * own_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) s.cfg.rank * s.cfg.max_elems; - const int peer = 1 - s.cfg.rank; - const float * peer_row = s.data_dev + (size_t) slot * s.slot_stride() + (size_t) peer * s.cfg.max_elems; - if (fast_reduce_kpub()) { - uint32_t * ctr = s.pub_counter(stm, slot, 0); - if (!ctr) return false; - fast_reduce_launch_pack_pub(data, own_row, (int) n, wire_bf16, - s.flag_dev + s.pub_off(slot), (uint32_t) seq, ctr, stm); - } else { - fast_reduce_launch_pack(data, own_row, (int) n, wire_bf16, stm); - if (hipStreamWriteValue32(stm, s.flag_dev + s.pub_off(slot), (uint32_t) seq, 0) != hipSuccess) return false; - } - if (hipStreamWaitValue32(stm, s.flag_dev + s.visible_off(slot, peer), (uint32_t) seq, - hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; - fast_reduce_launch_add_zc(data, peer_row, (int) n, wire_bf16, stm); - if (fast_reduce_lean()) return true; - return hipStreamWriteValue32(stm, s.flag_dev + s.done_off(slot), (uint32_t) seq, 0) == hipSuccess; - } - - // Publish: the copy engine moves the partial into the pinned staging - // buffer, then the flag is written. Both are on the stream, so the flag - // cannot precede the data it announces. - if (hipMemcpyAsync(s.data_host + (size_t) slot * s.slot_stride() + - (size_t) s.cfg.rank * s.cfg.max_elems, - data, n * sizeof(float), hipMemcpyDefault, stm) != hipSuccess) { - return false; - } - if (kernel_flags) { - fast_reduce_launch_set_flag(s.flag_dev + s.pub_off(slot), - s.flag_dev + s.done_off(slot), - s.progress_dev, (uint32_t) seq, - (uint32_t) (s.cfg.slots - 2), kSpin, stm); - } else if (hipStreamWriteValue32(stm, s.flag_dev + s.pub_off(slot), - (uint32_t) seq, 0) != hipSuccess) { - return false; - } - - // Wait for the peers, then bring their partials into device memory before - // touching them: the staging buffers are uncached for the GPU. - if (kernel_flags) { - fast_reduce_launch_wait_flags(s.peer_flags_dev + (size_t) slot * (size_t) n_peers, - s.timed_out_dev, s.progress_dev, - (uint32_t) seq, n_peers, kSpin, stm); - } else { - for (int r = 0; r < s.cfg.size; ++r) { - if (r == s.cfg.rank) continue; - if (hipStreamWaitValue32(stm, s.flag_dev + s.visible_off(slot, r), - (uint32_t) seq, hipStreamWaitValueEq, - 0xffffffffu) != hipSuccess) { - return false; - } - } - } - float * scratch_row = s.scratch + - ((size_t) slot * (size_t) n_peers) * (size_t) s.cfg.max_elems; - int k = 0; - for (int r = 0; r < s.cfg.size; ++r) { - if (r == s.cfg.rank) continue; - if (hipMemcpyAsync(scratch_row + (size_t) k * (size_t) s.cfg.max_elems, - s.data_host + (size_t) slot * s.slot_stride() + - (size_t) r * s.cfg.max_elems, - n * sizeof(float), hipMemcpyDefault, stm) != hipSuccess) { - return false; - } - ++k; - } - fast_reduce_launch_add(data, scratch_row, n_peers, (int) n, s.cfg.max_elems, - !kernel_flags ? nullptr : s.flag_dev + s.done_off(slot), - s.progress_dev, (uint32_t) seq, stm); - if (!kernel_flags && !fast_reduce_lean() && - hipStreamWriteValue32(stm, s.flag_dev + s.done_off(slot), - (uint32_t) seq, 0) != hipSuccess) { - return false; - } + // 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; } @@ -939,31 +724,17 @@ static uint64_t hyb_seq_for(FastReduceImpl & s, uint64_t k, size_t n) { return seq; } -bool FastReduce::hyb_export(const float * data, size_t n, void * stream, int part, uint64_t k) { +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 || part < 0 || part > 1) return false; + 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 (part == 0 && !fast_reduce_lean() && seq > (uint64_t) s.cfg.slots) { - const uint32_t need = (uint32_t) (seq - (uint64_t) s.cfg.slots); - if (hipStreamWaitValue32(stm, s.flag_dev + s.done_off(slot), need, - hipStreamWaitValueGte, 0xffffffffu) != hipSuccess) return false; - } - static const bool one_export = fr_env_on("LUCE_CLUSTER_HYBRID_ONE_EXPORT"); - const bool last_part = part == 1 || one_export; - if (last_part && fast_reduce_kpub()) { - uint32_t * ctr = s.pub_counter(stm, slot, 1); - if (!ctr) return false; - fast_reduce_launch_pack_pub(data, s.hyb_part(part, slot), (int) n, 0, - s.flag_dev + s.hybh_off(slot), (uint32_t) seq, ctr, stm); - return true; - } - fast_reduce_launch_pack(data, s.hyb_part(part, slot), (int) n, 0, stm); - if (last_part) { - return hipStreamWriteValue32(stm, s.flag_dev + s.hybh_off(slot), (uint32_t) seq, 0) == hipSuccess; - } + 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; } @@ -971,7 +742,6 @@ bool FastReduce::hyb_combine(const float * data, size_t n, void * stream, uint64 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 int wire_bf16 = fr_bf16_wire(s.cfg.size) ? 1 : 0; 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; @@ -980,26 +750,16 @@ bool FastReduce::hyb_combine(const float * data, size_t n, void * stream, uint64 auto stm = (hipStream_t) stream; if (hipStreamWaitValue32(stm, s.flag_dev + s.hybh_off(slot), (uint32_t) seq, hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; - static const bool one_export = fr_env_on("LUCE_CLUSTER_HYBRID_ONE_EXPORT"); - if (fast_reduce_kpub()) { - 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_pub(data, s.hyb_part(0, slot), one_export ? nullptr : s.hyb_part(1, slot), - own_row, (int) n, wire_bf16, - 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_pub(own_row, peer_row, s.hyb_part(2, slot), (int) n, wire_bf16, - s.flag_dev + s.hybf_off(slot), (uint32_t) seq, ctr_fin, stm); - return true; - } - fast_reduce_launch_hyb_combine(data, s.hyb_part(0, slot), one_export ? nullptr : s.hyb_part(1, slot), own_row, (int) n, wire_bf16, stm); - if (hipStreamWriteValue32(stm, s.flag_dev + s.pub_off(slot), (uint32_t) seq, 0) != 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(2, slot), (int) n, wire_bf16, stm); - return hipStreamWriteValue32(stm, s.flag_dev + s.hybf_off(slot), (uint32_t) seq, 0) == hipSuccess; + 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) { @@ -1011,14 +771,7 @@ bool FastReduce::hyb_import(float * data, size_t n, void * stream, uint64_t k) { auto stm = (hipStream_t) stream; if (hipStreamWaitValue32(stm, s.flag_dev + s.hybf_off(slot), (uint32_t) seq, hipStreamWaitValueEq, 0xffffffffu) != hipSuccess) return false; - static const bool import_dma = !fr_env_on("LUCE_CLUSTER_HYBRID_IMPORT_KERNEL"); - if (import_dma) { - if (hipMemcpyAsync(data, s.hyb_part(2, slot), n * sizeof(float), hipMemcpyHostToDevice, stm) != hipSuccess) return false; - } else { - fast_reduce_launch_pack(s.hyb_part(2, slot), data, (int) n, 0, stm); - } - if (fast_reduce_lean()) return true; - return hipStreamWriteValue32(stm, s.flag_dev + s.done_off(slot), (uint32_t) seq, 0) == hipSuccess; + 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) { @@ -1054,12 +807,8 @@ void FastReduce::shutdown() { } } s.pub_counters.clear(); - if (s.peer_flags) { (void) hipHostFree(s.peer_flags); s.peer_flags = nullptr; } - if (s.scratch) { (void) hipFree(s.scratch); s.scratch = nullptr; } 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.timed_out_host) { (void) hipHostFree(s.timed_out_host); s.timed_out_host = nullptr; } - if (s.progress_host) { (void) hipHostFree(s.progress_host); s.progress_host = nullptr; } if (s.hyb_host) { (void) hipHostFree(s.hyb_host); s.hyb_host = nullptr; } } @@ -1084,11 +833,9 @@ 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 *, int, uint64_t) { 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; } -uint64_t FastReduce::submitted() const { return 0; } -uint64_t FastReduce::timed_out() const { return 0; } } // namespace luce::cluster diff --git a/server/src/cluster/fast_reduce.h b/server/src/cluster/fast_reduce.h index e7ea3db86d..1b3f0adb3f 100644 --- a/server/src/cluster/fast_reduce.h +++ b/server/src/cluster/fast_reduce.h @@ -9,15 +9,16 @@ // 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. Pinned host memory that both the NIC and the GPU reach: the -// NIC writes each peer's row into it, and the GPU reads it there or after a DMA -// copy, so no GPUDirect is needed. It also assumes the ranks run +// 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: sum of f32, small sizes, no -// in-place aliasing games, no groups. Anything outside that keeps using RCCL. +// 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" @@ -50,17 +51,16 @@ class FastReduce { 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 has - // to be deep: the guard against lapping it is a stream wait, and a - // stream wait that fires stops the GPU running ahead, which is where - // this path's speed comes from. Prefill's much larger reductions fall - // back to RCCL, where they cost nothing that matters. + // 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: registers the buffers, exchanges endpoints over a + // 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. @@ -73,9 +73,9 @@ class FastReduce { // 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. `exact` keeps this payload in f32 on - // the wire and in the add when LUCE_CLUSTER_FAST_REDUCE_BF16 rounds the - // others, for values that must arrive bit for bit. + // 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 @@ -83,24 +83,20 @@ class FastReduce { // 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, zero-copy two-rank mode): - // a rank whose partial is split between a dGPU (hot experts, shared - // expert) and the iGPU that owns its cold experts. The dGPU exports its - // parts to pinned memory while it waits anyway; the iGPU adds its part to - // them in the order the dGPU used to, packs and publishes the rank's row, + // 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, int part, uint64_t k); + 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); - // Reductions issued, and flag waits of the kernel form that timed out, - // for the telemetry line. - uint64_t submitted() const; - uint64_t timed_out() const; private: std::unique_ptr p_; diff --git a/server/src/cluster/fast_reduce_kernels.cu b/server/src/cluster/fast_reduce_kernels.cu index ff0d71801e..93c5ecf471 100644 --- a/server/src/cluster/fast_reduce_kernels.cu +++ b/server/src/cluster/fast_reduce_kernels.cu @@ -1,29 +1,21 @@ // fast_reduce_kernels.cu - the device half of the lean decode-path reduction. // -// The flags normally go by stream memory operations (hipStreamWriteValue32 and -// hipStreamWaitValue32, in fast_reduce.cpp). The kernels here are the rest: +// The waits are stream memory operations (hipStreamWaitValue32, in +// fast_reduce.cpp); these kernels move and add the rows: // -// pack, add_zc the zero-copy path: pack this rank's row into pinned memory -// (f32 or bf16) and add the peer's row straight from it. -// add the copy path: add the peers' partials, by then copied into -// device memory. -// hyb_* the two-rank hybrid exchange (FastReduce::hyb_combine). -// set_flag, the kernel form of the flags, kept for comparison -// wait_flags (LUCE_CLUSTER_FAST_REDUCE_KERNEL_FLAGS=1). +// 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 // -// WHY THE COPY PATH MOVES THE PAYLOAD BY DMA. Pinned host memory is mapped -// uncached for the GPU, so a kernel reading or writing bulk data there crawls; -// the copy engine does not care, because it is not reading through the GPU's -// caches. +// 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. // -// WHY THE FLAGS ARE PINNED AND NOT MANAGED. 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, so they need -// no fault to be reached. -// -// 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. -// The kernel form's spin is bounded and counted (FastReduce::timed_out). +// 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 @@ -33,122 +25,17 @@ namespace luce::cluster { namespace { constexpr int kBlock = 256; +constexpr int kMaxGrid = 128; -// Ordering costs cache, and on this part it costs all of it. A system-scope -// release writes the L2 back and a system-scope acquire invalidates it, so a -// flag raised or spun on with the strong form evicts the weights the next layer -// is about to read, several times per reduction. -// -// So each flag gets the weakest ordering that is still correct: -// -// relaxed system for anything the host reads or writes. The payload is -// ordered by the stream -- the copy is a separate -// operation that completes before the kernel starts -- so -// no release is needed to publish it. -// one acquire after a spin succeeds, not once per iteration. The -// invalidate has to happen; it does not have to happen a -// million times. -// agent scope for slot_done, which only this GPU's own later kernels -// read. Telling the whole system about it buys nothing. -__device__ __forceinline__ uint32_t sys_load_relaxed(const uint32_t * p) { - return __hip_atomic_load(p, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM); -} - -__device__ __forceinline__ void sys_store_relaxed(uint32_t * p, uint32_t v) { - __hip_atomic_store(p, v, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_SYSTEM); -} - -__device__ __forceinline__ void sys_acquire(const uint32_t * p) { - (void) __hip_atomic_load(p, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_SYSTEM); -} - -__device__ __forceinline__ uint32_t dev_load(const uint32_t * p) { - return __hip_atomic_load(p, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); -} - -__device__ __forceinline__ void dev_store(uint32_t * p, uint32_t v) { - __hip_atomic_store(p, v, __ATOMIC_RELAXED, __HIP_MEMORY_SCOPE_AGENT); +int grid_for(int n) { + int grid = (n + kBlock - 1) / kBlock; + if (grid < 1) grid = 1; + if (grid > kMaxGrid) grid = kMaxGrid; + return grid; } } // namespace -__global__ void fast_reduce_set_flag_kernel(uint32_t * pub_flag, - uint32_t * slot_done, - uint32_t * progress, - uint32_t seq, - uint32_t slot_span, - uint64_t spin_limit) { - if (threadIdx.x || blockIdx.x) return; - // Do not overwrite a slot the progress thread has not finished with. The - // GPU can be a step's worth of reductions ahead, so the ring is defended - // rather than assumed wide enough. - if (seq > slot_span) { - uint64_t spins = 0; - while (dev_load(slot_done) + slot_span < seq) { - if (++spins > spin_limit) break; - } - } - sys_store_relaxed(progress, seq * 4u + 1u); - sys_store_relaxed(pub_flag, seq); -} - -__global__ void fast_reduce_wait_flags_kernel(uint32_t * const * __restrict__ peer_flags, - uint32_t * timed_out, - uint32_t * progress, - uint32_t seq, - int n_peers, - uint64_t spin_limit) { - if (threadIdx.x || blockIdx.x) return; - sys_store_relaxed(progress, seq * 4u + 2u); - for (int p = 0; p < n_peers; ++p) { - uint64_t spins = 0; - while (sys_load_relaxed(peer_flags[p]) != seq) { - if (++spins > spin_limit) { - sys_store_relaxed(timed_out, seq); - return; - } - } - } - // One invalidate, now that every peer has arrived, so the copies queued - // behind this read what the NIC wrote and not a cached line. - sys_acquire(peer_flags[0]); -} - -// The peers' partials are in device memory by now -- the copy engine moved them -// there after the wait -- so this is an ordinary cached add. -__global__ void fast_reduce_add_kernel(float * __restrict__ dst, - const float * __restrict__ scratch, - int n_peers, - int n, - int stride, - uint32_t * slot_done, - uint32_t * progress, - uint32_t seq) { - const int i = blockIdx.x * blockDim.x + threadIdx.x; - for (int k = i; k < n; k += blockDim.x * gridDim.x) { - float acc = dst[k]; - for (int p = 0; p < n_peers; ++p) { - acc += scratch[(size_t) p * (size_t) stride + (size_t) k]; - } - dst[k] = acc; - } - // The slot is free once its data has been consumed, not once it arrived. - // Saying so from one thread is safe only because the whole grid's writes - // are complete at the kernel boundary -- which is why the flag is raised - // here and not by a fence inside the loop. - // - // A null slot_done means the caller wants the arithmetic without the - // announcement: a system-scope release writes the L2 back, and the probe - // needs to price that separately from the copies. - if (i == 0 && slot_done) { - dev_store(slot_done, seq); - sys_store_relaxed(progress, seq * 4u + 3u); - } -} - -// ── Zero-copy path for discrete GPUs (LUCE_CLUSTER_FAST_REDUCE_ZC=1) ────── -// The partial goes straight into the pinned row the NIC sends from, and the -// peer's row is added straight out of pinned memory: no copy-engine hops. -// bf16 payloads round both operands so every rank computes the same sum. +// 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); @@ -158,14 +45,12 @@ __device__ __forceinline__ float fr_bf16_to_f32(uint16_t h) { return __uint_as_float(((uint32_t) h) << 16); } -// LUCE_CLUSTER_FAST_REDUCE_KPUB: 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. A null flag publishes nothing. counter is this device's, zeroed, and -// left zeroed for the next grid on the same stream. +// 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) { - if (!flag) return; __threadfence(); __syncthreads(); if (threadIdx.x == 0) { @@ -187,7 +72,7 @@ __global__ void fast_reduce_pack_kernel(const float * __restrict__ src, void * d fr_publish_last_block(flag, value, counter); } -__global__ void fast_reduce_add_zc_kernel(float * __restrict__ dst, const void * peer, int n, int bf16) { +__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])); @@ -198,95 +83,46 @@ __global__ void fast_reduce_add_zc_kernel(float * __restrict__ dst, const void * } } -// Hybrid exchange, iGPU side. own = (cold + hot) + shared in f32, the order -// the dGPU's graph adds them in, then the wire row; the final sum is -// bf16(own) + peer exactly as fast_reduce_add_zc_kernel forms it. +// 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, - const float * __restrict__ shared, void * own_row, int n, int bf16, + 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) { - const float own = shared ? (cold[k] + hot[k]) + shared[k] : cold[k] + hot[k]; - if (bf16) ((uint16_t *) own_row)[k] = fr_f32_to_bf16(own); - else ((float *) own_row)[k] = own; + ((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, int bf16, uint32_t * flag, uint32_t value, uint32_t * counter) { + 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) { - if (bf16) dst[k] = fr_bf16_to_f32(((const uint16_t *) own_row)[k]) + fr_bf16_to_f32(((const uint16_t *) peer_row)[k]); - else dst[k] = ((const float *) own_row)[k] + ((const float *) peer_row)[k]; + 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_hyb_combine_pub(const float * cold, const float * hot, const float * shared, void * own_row, - int n, int bf16, uint32_t * flag, uint32_t value, uint32_t * counter, - hipStream_t stream) { - int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; - hipLaunchKernelGGL(fast_reduce_hyb_combine_kernel, dim3(grid), dim3(kBlock), 0, stream, cold, hot, shared, own_row, n, bf16, +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_hyb_combine(const float * cold, const float * hot, const float * shared, void * own_row, - int n, int bf16, hipStream_t stream) { - fast_reduce_launch_hyb_combine_pub(cold, hot, shared, own_row, n, bf16, nullptr, 0, nullptr, stream); +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_final_pub(const void * own_row, const void * peer_row, float * dst, int n, int bf16, - uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream) { - int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; - hipLaunchKernelGGL(fast_reduce_hyb_final_kernel, dim3(grid), dim3(kBlock), 0, stream, own_row, peer_row, dst, n, bf16, - flag, value, counter); -} - -void fast_reduce_launch_hyb_final(const void * own_row, const void * peer_row, float * dst, int n, int bf16, - hipStream_t stream) { - fast_reduce_launch_hyb_final_pub(own_row, peer_row, dst, n, bf16, nullptr, 0, nullptr, stream); -} - -void fast_reduce_launch_pack_pub(const float * src, void * dst, int n, int bf16, - uint32_t * flag, uint32_t value, uint32_t * counter, hipStream_t stream) { - int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; - hipLaunchKernelGGL(fast_reduce_pack_kernel, dim3(grid), dim3(kBlock), 0, stream, src, dst, n, bf16, - flag, value, counter); +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_pack(const float * src, void * dst, int n, int bf16, hipStream_t stream) { - fast_reduce_launch_pack_pub(src, dst, n, bf16, nullptr, 0, nullptr, stream); -} - -void fast_reduce_launch_add_zc(float * dst, const void * peer, int n, int bf16, hipStream_t stream) { - int grid = (n + kBlock - 1) / kBlock; if (grid < 1) grid = 1; if (grid > 128) grid = 128; - hipLaunchKernelGGL(fast_reduce_add_zc_kernel, dim3(grid), dim3(kBlock), 0, stream, dst, peer, n, bf16); -} - -void fast_reduce_launch_set_flag(uint32_t * pub_flag, uint32_t * slot_done, - uint32_t * progress, uint32_t seq, - uint32_t slot_span, uint64_t spin_limit, - hipStream_t stream) { - hipLaunchKernelGGL(fast_reduce_set_flag_kernel, dim3(1), dim3(1), 0, stream, - pub_flag, slot_done, progress, seq, slot_span, spin_limit); -} - -void fast_reduce_launch_wait_flags(uint32_t * const * peer_flags, - uint32_t * timed_out, uint32_t * progress, - uint32_t seq, int n_peers, - uint64_t spin_limit, hipStream_t stream) { - hipLaunchKernelGGL(fast_reduce_wait_flags_kernel, dim3(1), dim3(1), 0, stream, - peer_flags, timed_out, progress, seq, n_peers, spin_limit); -} - -void fast_reduce_launch_add(float * dst, const float * scratch, int n_peers, - int n, int stride, uint32_t * slot_done, - uint32_t * progress, uint32_t seq, - hipStream_t stream) { - int grid = (n + kBlock - 1) / kBlock; - if (grid < 1) grid = 1; - if (grid > 256) grid = 256; - hipLaunchKernelGGL(fast_reduce_add_kernel, dim3(grid), dim3(kBlock), 0, stream, - dst, scratch, n_peers, n, stride, slot_done, progress, seq); +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/deepseek4/deepseek4_cluster.cpp b/server/src/deepseek4/deepseek4_cluster.cpp index 2cacfaf79a..3c1b896971 100644 --- a/server/src/deepseek4/deepseek4_cluster.cpp +++ b/server/src/deepseek4/deepseek4_cluster.cpp @@ -605,13 +605,9 @@ static void ds4_cluster_allreduce_graph_callback(void * user, void * data, 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_export0_cb(void * user, void * data, size_t n, void * stream) { +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, 0, rt->hyb_export0++)) ds4_hyb_fail(rt, "export (hot)"); -} -static void ds4_hyb_export1_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, 1, rt->hyb_export1++)) ds4_hyb_fail(rt, "export (shared)"); + 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); @@ -637,35 +633,20 @@ ggml_tensor * ds4_cluster_hybrid_exchange_nodes(ggml_context * ctx, ggml_tensor * cold, Ds4ClusterRuntime & rt, ggml_tensor ** combine_node, - ggml_tensor ** export_hot_node, - ggml_tensor ** export_shared_node) { - if (!ctx || !hot || !shared || !cold || !combine_node || !export_hot_node || !export_shared_node) return nullptr; + 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. - // LUCE_CLUSTER_HYBRID_ONE_EXPORT=1: export hot + shared as one partial - // (half the dGPU link bytes; the iGPU adds cold + (hot + shared)). - static const bool one_export = [] { - const char * v = std::getenv("LUCE_CLUSTER_HYBRID_ONE_EXPORT"); - return v && *v && std::strcmp(v, "0") != 0; - }(); - ggml_tensor * ex0 = nullptr; - ggml_tensor * ex1 = nullptr; - if (one_export) { - ggml_tensor * hs = ggml_add(ctx, hot, shared); - ex1 = ggml_cluster_allreduce_inplace(ctx, hs, &ds4_hyb_export0_cb, &rt); - ex0 = ex1; - } else { - ex0 = ggml_cluster_allreduce_inplace(ctx, hot, &ds4_hyb_export0_cb, &rt); - ex1 = ggml_cluster_allreduce_inplace(ctx, shared, &ds4_hyb_export1_cb, &rt); - } + // 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); - // The import overwrites the exported shared partial with the reduced sum. - ggml_tensor * imp = ggml_cluster_allreduce_inplace(ctx, ex1, &ds4_hyb_import_cb, &rt); + ggml_tensor * imp = ggml_cluster_allreduce_inplace(ctx, ex, &ds4_hyb_import_cb, &rt); *combine_node = comb; - *export_hot_node = ex0; - *export_shared_node = ex1; + *export_node = ex; return imp; } diff --git a/server/src/deepseek4/deepseek4_cluster.h b/server/src/deepseek4/deepseek4_cluster.h index 8e526c5b9c..f64b5f4790 100644 --- a/server/src/deepseek4/deepseek4_cluster.h +++ b/server/src/deepseek4/deepseek4_cluster.h @@ -66,7 +66,7 @@ struct Ds4ClusterRuntime { // 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_export0 = 0, hyb_export1 = 0, hyb_combine = 0, hyb_import = 0; + 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; @@ -190,8 +190,7 @@ ggml_tensor * ds4_cluster_hybrid_exchange_nodes(ggml_context * ctx, ggml_tensor * cold, Ds4ClusterRuntime & rt, ggml_tensor ** combine_node, - ggml_tensor ** export_hot_node, - ggml_tensor ** export_shared_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 diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 950cc662e7..676a3f986a 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -1638,19 +1638,16 @@ static bool ds4_build_fused_verify_graph( // imports the sum (no device join, one link crossing // each way on the dGPU). ggml_tensor * combine = nullptr; - ggml_tensor * export_hot = nullptr; - ggml_tensor * export_shared = 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, &export_hot, &export_shared); + *cluster_rt, &combine, &exported); if (!sum) return false; - ggml_build_forward_expand(gf, export_hot); - ggml_build_forward_expand(gf, export_shared); - cluster_nodes.push_back(export_shared); + 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); - cluster_nodes.push_back(export_hot); // 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; diff --git a/server/src/server/launch_profiles.h b/server/src/server/launch_profiles.h index f2356ab852..77f9dfdcbc 100644 --- a/server/src/server/launch_profiles.h +++ b/server/src/server/launch_profiles.h @@ -258,12 +258,8 @@ inline const std::vector & cluster_launch_env() { // 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_ZC", "1"}, - {"LUCE_CLUSTER_FAST_REDUCE_BF16", "1"}, - {"LUCE_CLUSTER_FAST_REDUCE_LEAN", "1"}, {"LUCE_CLUSTER_FAST_REDUCE_STALL_S", "15"}, {"LUCE_CLUSTER_HYBRID_EXCHANGE", "1"}, - {"LUCE_CLUSTER_HYBRID_ONE_EXPORT", "1"}, // Prefill-sized reductions stay on RCCL: 8-bit blocks, or bf16 with // LUCE_CLUSTER_PREFILL_Q8=0. {"LUCE_CLUSTER_PREFILL_Q8", "1"},