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 2569d0828..448ef93ec 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -497,7 +497,12 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } -template +static bool is_gfx1151(const int cc) { + return cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151; +} + +template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst, @@ -519,9 +524,16 @@ static __global__ void mul_mat_vec_q( const int tid = warp_size*threadIdx.y + threadIdx.x; const int row0 = rows_per_cuda_block*blockIdx.x; - const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_row_x = (fixed_ncols_x > 0 ? fixed_ncols_x : ncols_x) / qk; constexpr int blocks_per_iter = vdr * nwarps*warp_size / qi; + static_assert(fixed_ncols_x == 0 || + type == GGML_TYPE_Q3_0_ROCMFPX || + type == GGML_TYPE_Q2_0_ROCMFP2, + "fixed-K decode specialization is limited to profiled ROCmFP2/FP3 types"); + static_assert(!unroll_k_loop_2 || type == GGML_TYPE_Q4_0_ROCMFP4_FAST, + "partial K-loop unrolling is limited to the profiled FP4FAST path"); + const uint32_t channel_dst = blockIdx.y; const bool has_ids = ids != nullptr; @@ -601,23 +613,76 @@ static __global__ void mul_mat_vec_q( const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y; - for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { - const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx + if constexpr (fixed_ncols_x > 0) { + static_assert(fixed_ncols_x % qk == 0, "fixed K must contain whole quant blocks"); + constexpr int fixed_blocks = fixed_ncols_x/qk; + constexpr int fixed_iters = (fixed_blocks + blocks_per_iter - 1) / blocks_per_iter; +#pragma unroll + for (int iter = 0; iter < fixed_iters; ++iter) { + const int kbx = tid / (qi/vdr) + iter*blocks_per_iter; + if constexpr (fixed_blocks % blocks_per_iter != 0) { + if (kbx >= fixed_blocks) { + continue; + } + } + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (tid % (qi/vdr)); - // x block quant index when casting the quants to int - const int kqs = vdr * (tid % (qi/vdr)); +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } + } + } + } + } else if constexpr (unroll_k_loop_2) { +#pragma unroll 2 + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (tid % (qi/vdr)); #pragma unroll - for (int j = 0; j < ncols_dst; ++j) { - const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; + for (int j = 0; j < ncols_dst; ++j) { + const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; #pragma unroll - for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp[j][i] += vec_dot_q_cuda( - vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); - if constexpr (has_fusion) { - if (use_gate) { - tmp_gate[j][i] += vec_dot_q_cuda( - vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } + } + } + } + } else { + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (tid % (qi/vdr)); + +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } } } } @@ -1138,7 +1203,8 @@ static std::pair calc_launch_params( return {block_nums, block_dims}; } -template +template static void mul_mat_vec_q_switch_fusion( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, @@ -1150,19 +1216,70 @@ static void mul_mat_vec_q_switch_fusion( const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; if (has_fusion) { - mul_mat_vec_q<<>> + mul_mat_vec_q<<>> (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, 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, ids_stride, ids_tokenwise_samples); return; } - mul_mat_vec_q<<>> + mul_mat_vec_q<<>> (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, 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, ids_stride, ids_tokenwise_samples); } +template +static void mul_mat_vec_rocmfpx_fixed_k_launch( + const void * vx, const void * vy, const int32_t * ids, + const ggml_cuda_mm_fusion_args_device fusion, float * dst, + const uint32_t ncols_x, const uint32_t nrows_x, const uint3 nchannels_y, + const uint32_t nchannels_dst, const uint32_t stride_row_x, + const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, + const uint32_t nsamples_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, + const uint32_t stride_sample_dst, const uint32_t ids_stride, + const int warp_size, cudaStream_t stream) { + static_assert(type == GGML_TYPE_Q3_0_ROCMFPX || + type == GGML_TYPE_Q2_0_ROCMFP2, + "fixed-K helper is limited to profiled ROCmFP2/FP3 kernels"); + static_assert(fixed_ncols_x == 2048 || fixed_ncols_x == 4096, + "fixed-K helper compiled an unexpected DS4 shape"); + + const dim3 block_nums(nrows_x, nchannels_dst, nsamples_dst); + const dim3 block_dims(warp_size, 1, 1); + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, 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, block_nums, block_dims, 0, + ids_stride, stream); +} + +static void mul_mat_vec_rocmfp4_unroll2_launch( + const void * vx, const void * vy, const int32_t * ids, + const ggml_cuda_mm_fusion_args_device fusion, float * dst, + const uint32_t ncols_x, const uint32_t nrows_x, const uint3 nchannels_y, + const uint32_t nchannels_dst, const uint32_t stride_row_x, + const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, + const uint32_t nsamples_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, + const uint32_t stride_sample_dst, const uint32_t ids_stride, + const int warp_size, cudaStream_t stream) { + const dim3 block_nums(nrows_x, nchannels_dst, nsamples_dst); + const dim3 block_dims(warp_size, 1, 1); + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, 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, block_nums, block_dims, 0, + ids_stride, stream); +} + template static void mul_mat_vec_q_moe_launch( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, @@ -1333,6 +1450,60 @@ static void mul_mat_vec_q_switch_ncols_dst( return; } + // Decode-only gfx1151 specializations. Each one preserves the original + // per-lane K traversal and accumulation order, so exact validated shapes + // can use the faster kernel without a serving-time tuning flag. + if constexpr (type == GGML_TYPE_Q4_0_ROCMFP4_FAST) { + if (is_gfx1151(cc) && ncols_dst == 1) { + mul_mat_vec_rocmfp4_unroll2_launch( + vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, + nchannels_dst, stride_row_x, stride_col_y, stride_col_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, + ids_stride, warp_size, stream); + return; + } + } + + if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX) { + if (is_gfx1151(cc) && ncols_dst == 1 && ncols_x == 2048) { + mul_mat_vec_rocmfpx_fixed_k_launch( + vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, + nchannels_dst, stride_row_x, stride_col_y, stride_col_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, + ids_stride, warp_size, stream); + return; + } + } + + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2) { + if (is_gfx1151(cc) && ncols_dst == 1) { + if (ncols_x == 4096) { + mul_mat_vec_rocmfpx_fixed_k_launch( + vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, + nchannels_dst, stride_row_x, stride_col_y, stride_col_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, + ids_stride, warp_size, stream); + return; + } + if (ncols_x == 2048) { + mul_mat_vec_rocmfpx_fixed_k_launch( + vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, + nchannels_dst, stride_row_x, stride_col_y, stride_col_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, + ids_stride, warp_size, stream); + return; + } + } + } + switch (ncols_dst) { case 1: { constexpr int c_ncols_dst = 1; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh index bb75e7977..07a28864f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh @@ -401,6 +401,29 @@ static __device__ __forceinline__ int rocmfpx_pack4_fp6_bits24_vec_cuda(const ui return *((const int *) &v); } +static __device__ __forceinline__ int rocmfpx_pack4_fp3_bits12_vec_cuda(const uint32_t bits12) { +#if defined(GGML_USE_HIP) + // Byte tables for codes {0, 1, 2, 4} and {0, -1, -2, -4}. + // v_perm_b32 selects a byte with each 3-bit selector, replacing four + // scalar code decodes without changing the packed int8 values. + constexpr uint32_t values_low = 0x04020100u; + constexpr uint32_t values_high = 0xFCFEFF00u; + const uint32_t selectors = + ((bits12 >> 0) & 7u) | + (((bits12 >> 3) & 7u) << 8) | + (((bits12 >> 6) & 7u) << 16) | + (((bits12 >> 9) & 7u) << 24); + return (int) __builtin_amdgcn_perm(values_high, values_low, selectors); +#else + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp3_code_vec_cuda(bits12 & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 3) & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 6) & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 9) & 7u)); + return *((const int *) &v); +#endif +} + static __device__ __forceinline__ int rocmfpx_pack4_fp3_vec_cuda(const uint8_t * qs, const int base) { const char4 v = make_char4( (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 0)*3, 3)), @@ -532,12 +555,7 @@ static __device__ __forceinline__ float vec_dot_rocmfpx_fp3_q8_1( const uint32_t val_high = qs[reg_idx + 1]; const uint32_t bits12 = (reg_shift == 0) ? (val_low & 0xFFFu) : (((val_low >> reg_shift) | (val_high << (32 - reg_shift))) & 0xFFFu); - const char4 v = make_char4( - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(bits12 & 7u), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 3) & 7u), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 6) & 7u), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 9) & 7u)); - const int val_packed = *((const int *) &v); + const int val_packed = rocmfpx_pack4_fp3_bits12_vec_cuda(bits12); const int u = get_int_b4(bq8_1->qs, iqs + i); diff --git a/server/docs/DS4.md b/server/docs/DS4.md index c3e32012f..ca83f8555 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -1,6 +1,8 @@ # DeepSeek V4 Flash — DFlash Integration -This document describes the current DeepSeek V4 Flash implementation in DFlash. Today DeepSeek4 runs through the layer-split path: a local CUDA prefix shard plus either a local follow-on shard or a remote Halo/HIP target shard. +This document describes the current DeepSeek V4 Flash implementation in +DFlash. DeepSeek4 supports a monolithic HIP backend for single-device Strix +Halo systems and a layer-split backend for local or mixed-device deployments. ## Model Architecture @@ -23,7 +25,7 @@ DeepSeek V4 Flash is a 43-layer MoE model with: | Area | Files | |------|-------| -| Backend selection / init | `src/common/backend_factory.cpp`, `src/deepseek4/deepseek4_layer_split_adapter.{h,cpp}` | +| Backend selection / init | `src/common/backend_factory.cpp`, `src/deepseek4/deepseek4_backend.{h,cpp}`, `src/deepseek4/deepseek4_layer_split_adapter.{h,cpp}` | | Per-shard forward graph | `src/deepseek4/deepseek4_graph.cpp` | | Model weights and metadata | `src/deepseek4/deepseek4_internal.h`, `src/deepseek4/deepseek4_loader.cpp` | | HC pre/post CUDA kernel | `src/deepseek4/deepseek4_hc_cuda.cu`, `.h` | @@ -45,6 +47,32 @@ The production DeepSeek4 path does **not** use the retired per-expert worker spl ## Execution Modes +### Monolithic HIP + +Single-device HIP launches use `DeepSeek4Backend`. Two explicit serving +options are available: + +- `--ds4-fused-decode` enables the cached single-graph decode path. It keeps + HC, attention, MoE, and the output projection on the GPU and avoids + per-layer host round trips. On HIP this option requests a monolithic model + load because the fused graph must reference every expert tensor directly. + If that allocation fails, the backend logs the fallback and continues with + hybrid expert placement and layered decode. +- `--ds4-expert-top-k N` keeps the highest-ranked `N` routed experts and + renormalizes their weights. `0` uses the model default. Reducing this value is an + approximate inference policy and must be quality-validated for the target + workload. + +These options currently apply only to the monolithic HIP backend. For the +validated Strix Halo profile: + +```bash +./server/build-hip/dflash_server /opt/models/DeepSeek-V4-Flash.gguf \ + --target-device hip:0 \ + --ds4-fused-decode \ + --ds4-expert-top-k 4 +``` + ### Local single-shard If the adapter decides all 43 layers fit on one CUDA GPU, it loads a single shard locally and no IPC daemon is involved. @@ -148,6 +176,6 @@ Explicit mixed-backend split using the generic target-shard flags: | Target | Backend | Purpose | |--------|---------|---------| -| `dflash_server` | CUDA | Production server / CUDA parent | +| `dflash_server` | CUDA or HIP | Production server | | `backend_ipc_daemon` | HIP | Remote Halo target shard for mixed-backend layer split | | `test_deepseek4_unit` | CUDA | Unit tests (no model files needed) | diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 1383c2166..df96a1410 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -232,6 +232,8 @@ std::unique_ptr create_backend(const BackendArgs & args) { cfg.stream_fd = args.stream_fd; cfg.max_ctx = args.device.max_ctx; cfg.chunk = args.chunk; + cfg.expert_top_k = args.ds4_expert_top_k; + cfg.fused_decode = args.ds4_fused_decode; auto backend = std::make_unique(cfg); if (!backend->init()) { diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index 73ed8698c..75c2053f7 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -45,6 +45,10 @@ struct BackendArgs { // Chunked prefill int chunk = 512; + // deepseek4-specific decode options + int ds4_expert_top_k = 0; // 0 = model default + bool ds4_fused_decode = false; + // qwen35-specific speculative decode options int fa_window = 0; // 0 = full attention. qwen3.6 full-attn layers must see the whole context; a finite window drops the system prompt/tools -> breaks tool calls. int kq_stride_pad = 32; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index e75a49ea9..eb6799b58 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -69,6 +69,10 @@ static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTeleme dst.output_us += src.output_us; dst.sample_us += src.sample_us; dst.emit_us += src.emit_us; + dst.full_graph_build_us += src.full_graph_build_us; + dst.full_graph_set_us += src.full_graph_set_us; + dst.full_graph_compute_us += src.full_graph_compute_us; + dst.full_graph_read_us += src.full_graph_read_us; dst.hot_selected += src.hot_selected; dst.cold_selected += src.cold_selected; } @@ -92,6 +96,7 @@ static void log_step_tel(const char * phase, "ffn_hot_graph_build=%llu ffn_hot_graph_hit=%llu ffn_cold_graph_build=%llu ffn_cold_graph_hit=%llu " "hc_pre=%.1fms hc_pre_build=%.1fms hc_pre_input=%.1fms hc_pre_compute=%.1fms " "hc_post=%.1fms output=%.1fms sample=%.1fms emit=%.1fms " + "full_build=%.1fms full_set=%.1fms full_compute=%.1fms full_read=%.1fms " "hot_sel=%d cold_sel=%d\n", phase, tokens, steps, wall_s, tok_s, ms(t.total_us), ms(t.embed_us), ms(t.attn_build_us), ms(t.attn_compute_us), ms(t.attn_read_us), @@ -107,6 +112,8 @@ static void log_step_tel(const char * phase, ms(t.hc_pre_compute_us), ms(t.hc_post_attn_us + t.hc_post_ffn_us), ms(t.output_us), ms(t.sample_us), ms(t.emit_us), + ms(t.full_graph_build_us), ms(t.full_graph_set_us), ms(t.full_graph_compute_us), + ms(t.full_graph_read_us), t.hot_selected, t.cold_selected); } @@ -309,25 +316,30 @@ DeepSeek4Backend::~DeepSeek4Backend() { shutdown(); } -bool DeepSeek4Backend::init() { - backend_ = ggml_backend_cuda_init(cfg_.device.gpu); - if (!backend_) { - std::fprintf(stderr, "[deepseek4] failed to create CUDA backend (gpu=%d)\n", - cfg_.device.gpu); - return false; - } - - snap_backend_ = ggml_backend_init_by_name("cpu", nullptr); - +bool DeepSeek4Backend::load_model() { const PlacementBackend target_backend = cfg_.device.backend == PlacementBackend::Auto ? compiled_placement_backend() : cfg_.device.backend; - // HIP single-device launches should avoid the monolithic full-model load: - // a managed ~80 GiB allocation can stall or be killed on integrated UMA - // systems before we ever reach the existing OOM fallback path. - if (target_backend == PlacementBackend::Hip) { + // The fused graph references every expert tensor directly, so it cannot + // run on the hybrid expert representation. Honor an explicit fused-decode + // request by trying the monolithic load first. Large HIP models otherwise + // keep using the hybrid path: a managed allocation can stall or be killed + // on integrated UMA systems before an OOM fallback is possible. + if (target_backend == PlacementBackend::Hip && cfg_.fused_decode) { + std::fprintf(stderr, + "[deepseek4] fused decode requested; loading monolithic HIP model\n"); + if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { + std::fprintf(stderr, + "[deepseek4] monolithic HIP load failed; trying hybrid mode\n"); + if (!init_hybrid_model()) { + std::fprintf(stderr, "[deepseek4] hybrid mode also failed: %s\n", + cfg_.model_path); + return false; + } + } + } else if (target_backend == PlacementBackend::Hip) { std::fprintf(stderr, "[deepseek4] HIP target detected; using hybrid expert load path\n"); if (!init_hybrid_model()) { @@ -342,6 +354,36 @@ bool DeepSeek4Backend::init() { } } + if (cfg_.expert_top_k < 0 || cfg_.expert_top_k > w_.n_expert_used) { + std::fprintf(stderr, + "[deepseek4] expert top-k must be in [0,%d], got %d\n", + w_.n_expert_used, cfg_.expert_top_k); + return false; + } + w_.routed_expert_top_k = cfg_.expert_top_k; + w_.fused_decode = cfg_.fused_decode && !moe_hybrid_; + if (cfg_.fused_decode && moe_hybrid_) { + std::fprintf(stderr, + "[deepseek4] fused decode unavailable with hybrid expert placement; " + "using layered decode\n"); + } + return true; +} + +bool DeepSeek4Backend::init() { + backend_ = ggml_backend_cuda_init(cfg_.device.gpu); + if (!backend_) { + std::fprintf(stderr, "[deepseek4] failed to create CUDA backend (gpu=%d)\n", + cfg_.device.gpu); + return false; + } + + snap_backend_ = ggml_backend_init_by_name("cpu", nullptr); + + if (!load_model()) { + return false; + } + const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { std::fprintf(stderr, "[deepseek4] failed to allocate KV cache (ctx=%d)\n", max_ctx); @@ -353,8 +395,13 @@ bool DeepSeek4Backend::init() { // The DeepSeek4Backend single-GPU path now runs all experts locally. } - std::fprintf(stderr, "[deepseek4] initialized: %d layers, ctx=%d, %d experts (%d used)%s\n", - w_.n_layer, max_ctx, w_.n_expert, w_.n_expert_used, + const int active_experts = + w_.routed_expert_top_k > 0 ? w_.routed_expert_top_k : w_.n_expert_used; + std::fprintf(stderr, + "[deepseek4] initialized: %d layers, ctx=%d, %d experts " + "(%d/%d routed), fused_decode=%s%s\n", + w_.n_layer, max_ctx, w_.n_expert, active_experts, w_.n_expert_used, + w_.fused_decode ? "on" : "off", moe_hybrid_ ? " [hybrid]" : ""); return true; } @@ -487,30 +534,13 @@ bool DeepSeek4Backend::unpark(const std::string & what) { const bool want_target = (what.empty() || what == "all" || what == "target"); if (!want_target || !parked_) return true; - const PlacementBackend target_backend = - cfg_.device.backend == PlacementBackend::Auto - ? compiled_placement_backend() - : cfg_.device.backend; - - if (target_backend == PlacementBackend::Hip) { - if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] unpark: failed to restore hybrid mode\n"); - free_deepseek4_weights(w_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; - return false; - } - } else if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { - std::fprintf(stderr, "[deepseek4] unpark: full model reload failed, trying hybrid mode...\n"); - if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); - free_deepseek4_weights(w_); - stream_engine_.destroy(); - moe_hybrid_.reset(); - moe_placement_ = {}; - return false; - } + if (!load_model()) { + std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); + free_deepseek4_weights(w_); + stream_engine_.destroy(); + moe_hybrid_.reset(); + moe_placement_ = {}; + return false; } const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; @@ -533,11 +563,20 @@ bool DeepSeek4Backend::unpark(const std::string & what) { int DeepSeek4Backend::do_prefill(const std::vector & tokens, const DaemonIO & io, int kv_offset) { - // Hybrid currently implements HC for single-token steps only; keep prefill - // token-by-token so the first sampled token is seeded from the correct HC state. - const int chunk = moe_hybrid_ ? 1 : (cfg_.chunk > 0 ? cfg_.chunk : 512); + // The current batched graph updates only the chunk's final compressor + // window. Until batched compressor state is sequentially updated, process + // prefill one token at a time to preserve long-prompt correctness. + constexpr int chunk = 1; const int n_total = (int)tokens.size(); int pos = kv_offset; + // New sequence: clear the cache buffer so compressor state double-buffers + // and compressed-KV rows start from zeros, exactly like a fresh server. + // Without this, the first flush windows of a request pool over the + // previous request's leftover state rows and outputs from the 2nd/3rd + // request on can drift by a token or two. + if (kv_offset == 0) { + reset_deepseek4_cache(cache_); + } last_logits_.clear(); const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); const auto phase_t0 = Clock::now(); @@ -556,7 +595,6 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, DeepSeek4StepTelemetry step_tel; if (timing) step_tel.embed_us = elapsed_us(embed_t0, Clock::now()); - // Run forward pass std::vector logits; if (!deepseek4_step(backend_, w_, cache_, embed.data(), n_tok, pos, logits, moe_hybrid_.get(), tokens.data() + i, diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 508637867..537d3ba44 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -87,6 +87,7 @@ class DeepSeek4Backend : public ModelBackend { const BudgetHook & budget_hook = {}, bool * forced_close_out = nullptr); + bool load_model(); bool init_hybrid_model(); bool compute_uniform_hybrid_placement(const DeepSeek4Weights & w, int max_ctx, diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 06fbe8a98..9879fee36 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -25,12 +25,18 @@ #include #include #include +#include #include -#include #include #include +#include +#include #include +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) +#include +#endif + namespace dflash::common { namespace { @@ -41,9 +47,39 @@ static uint64_t ds4_elapsed_us(Ds4TimingClock::time_point start, return (uint64_t)std::chrono::duration_cast(end - start).count(); } -static bool ds4_env_flag(const char * name) { - const char * value = std::getenv(name); - return value && value[0] && std::strcmp(value, "0") != 0; +static int ds4_effective_expert_count(const DeepSeek4Weights & w) { + const int requested = w.routed_expert_top_k; + if (requested > 0 && requested < w.n_expert_used) { + return requested; + } + return w.n_expert_used; +} + +static size_t ds4_attn_step_meta_size(int n_tokens) { + size_t arena_size = 48 * 1024 * 1024; + if (n_tokens >= 512) { + arena_size += (size_t)n_tokens * 32 * 1024; + } + return arena_size; +} + +static size_t ds4_attn_step_graph_size(int n_tokens) { + if (n_tokens <= 1) { + return 2048; + } + if (n_tokens <= 512) { + return 32768; + } + if (n_tokens <= 1024) { + return 131072; + } + if (n_tokens <= 2048) { + return 262144; + } + if (n_tokens <= 4096) { + return 524288; + } + return 1048576; } template @@ -99,6 +135,30 @@ static void add_ffn_telemetry(DeepSeek4StepTelemetry * dst, } // namespace +int deepseek4_previous_raw_ring_spans( + int kv_start, + int n_swa, + DeepSeek4RawRingSpan spans[2]) { + if (!spans || kv_start <= 0 || n_swa <= 0) { + return 0; + } + if (kv_start < n_swa) { + spans[0] = {0, kv_start}; + return 1; + } + + const int current_row = kv_start % n_swa; + int count = 0; + const int tail_count = n_swa - current_row - 1; + if (tail_count > 0) { + spans[count++] = {current_row + 1, tail_count}; + } + if (current_row > 0) { + spans[count++] = {0, current_row}; + } + return count; +} + struct DeepSeek4I32InputBinding { ggml_tensor * tensor = nullptr; int32_t value = 0; @@ -136,6 +196,7 @@ struct DeepSeek4CachedDecodeFfnGraph { ggml_backend_t backend = nullptr; int layer_idx = -1; int n_tokens = 0; + int n_expert_used = 0; bool hash_routed = false; StepGraph sg; ggml_tensor * hash_ids = nullptr; @@ -153,6 +214,7 @@ struct DeepSeek4CachedDecodeFfnGraph { backend = nullptr; layer_idx = -1; n_tokens = 0; + n_expert_used = 0; hash_routed = false; } }; @@ -188,6 +250,10 @@ struct DeepSeek4AttentionGraphInputs { ggml_tensor * index_state_rows = nullptr; ggml_tensor * index_comp_rows = nullptr; ggml_tensor * index_comp_pos = nullptr; + // Fused-decode stable-KV path only: additive score mask over + // [n_swa raw rows ++ padded comp rows]; 0 for valid, -1e30 for padding. + ggml_tensor * attn_row_mask = nullptr; + int padded_comp = 0; // padded compressed-row count (>= n_comp) }; struct DeepSeek4CachedDecodeAttnGraph { @@ -198,8 +264,11 @@ struct DeepSeek4CachedDecodeAttnGraph { int n_raw = 0; int n_comp_attn = 0; int n_index_comp = 0; + bool attn_flush = false; + bool index_flush = false; bool compressed = false; bool indexed = false; + bool uses_shared_inputs = false; StepGraph sg; DeepSeek4AttentionGraphInputs inputs; @@ -224,8 +293,11 @@ struct DeepSeek4CachedDecodeAttnGraph { n_raw = 0; n_comp_attn = 0; n_index_comp = 0; + attn_flush = false; + index_flush = false; compressed = false; indexed = false; + uses_shared_inputs = false; } }; @@ -315,14 +387,14 @@ static bool build_cached_decode_ffn_graph( ggml_tensor * ffn_normed = build_rms_norm(out.sg.ctx, out.sg.inp_embed, L.ffn_norm, w.rms_eps); ggml_tensor * ffn_out = nullptr; if (hash_routed) { - out.hash_ids = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I32, w.n_expert_used, n_tokens); + const int n_used = ds4_effective_expert_count(w); + out.hash_ids = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I32, n_used, n_tokens); ggml_set_input(out.hash_ids); ggml_tensor * shared_out = build_shared_ffn(out.sg.ctx, ffn_normed, w, L); ggml_tensor * logits = ggml_mul_mat(out.sg.ctx, L.ffn_gate_inp, ffn_normed); ggml_tensor * probs = ggml_sqrt(out.sg.ctx, ggml_softplus(out.sg.ctx, logits)); - const int n_used = w.n_expert_used; const int n_ff_exp = w.n_ff_exp; ggml_tensor * cur_3d = ggml_reshape_3d(out.sg.ctx, ffn_normed, w.n_embd, 1, n_tokens); ggml_tensor * gate_e = ggml_mul_mat_id(out.sg.ctx, L.ffn_gate_exps, cur_3d, out.hash_ids); @@ -345,7 +417,8 @@ static bool build_cached_decode_ffn_graph( ggml_tensor * weights_3d = ggml_reshape_3d(out.sg.ctx, weights, 1, n_used, n_tokens); ggml_tensor * routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); - routed_out = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_cont( + out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); routed_out = ggml_sum_rows(out.sg.ctx, routed_out); routed_out = ggml_reshape_2d(out.sg.ctx, routed_out, w.n_embd, n_tokens); @@ -373,6 +446,7 @@ static bool build_cached_decode_ffn_graph( out.backend = backend; out.layer_idx = layer_idx; out.n_tokens = n_tokens; + out.n_expert_used = ds4_effective_expert_count(w); out.hash_routed = hash_routed; return true; } @@ -428,12 +502,14 @@ static ggml_tensor * build_clamped_swiglu(ggml_context * ctx, ggml_tensor * gate, ggml_tensor * up, float clamp) { - // DS4 clamps only the upper side of gate, but both sides of up. - gate = ggml_clamp(ctx, gate, -INFINITY, clamp); - up = ggml_clamp(ctx, up, -clamp, clamp); - // silu(gate) * up - gate = ggml_silu(ctx, gate); - return ggml_mul(ctx, gate, up); + return ggml_swiglu_ds4_split(ctx, gate, up, clamp); +} + +static ggml_tensor * ds4_cast_if_needed( + ggml_context * ctx, + ggml_tensor * x, + ggml_type type) { + return x->type == type ? x : ggml_cast(ctx, x, type); } // ─── Helper: Partial RoPE (tail rotation) ─────────────────────────────── @@ -526,7 +602,8 @@ static void build_compressor_step( ggml_tensor * comp_rows_inp, ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, - std::vector & i32_array_inputs) { + std::vector & i32_array_inputs, + ggml_tensor ** comp_cache_source_out = nullptr) { if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; @@ -541,6 +618,9 @@ static void build_compressor_step( ggml_tensor * kv_cur = ggml_mul_mat(ctx, kv_proj, cur_last); ggml_tensor * sc_cur = ggml_mul_mat(ctx, gate_proj, cur_last); + ggml_tensor * state_kv_source = state.state_kv; + ggml_tensor * state_score_source = state.state_score; + ggml_tensor * comp_cache_source = comp_cache; ggml_tensor * ape_col = nullptr; if (ape_row_inp) { @@ -554,8 +634,10 @@ static void build_compressor_step( sc_cur = ggml_add(ctx, sc_cur, ape_col); if (state_rows_inp) { - ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_kv, kv_cur, state_rows_inp)); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_score, sc_cur, state_rows_inp)); + state_kv_source = ggml_set_rows(ctx, state.state_kv, kv_cur, state_rows_inp); + state_score_source = ggml_set_rows(ctx, state.state_score, sc_cur, state_rows_inp); + ggml_build_forward_expand(gf, state_kv_source); + ggml_build_forward_expand(gf, state_score_source); } else { ggml_tensor * kv_slot = ggml_view_2d( ctx, state.state_kv, comp_width, 1, state.state_kv->nb[1], @@ -586,26 +668,26 @@ static void build_compressor_step( ggml_tensor * sv_sc = nullptr; int n_state_rows = ratio; if (ratio == 4) { - const size_t hi_off_kv = (size_t)ratio * state.state_kv->nb[1] + - (size_t)head_dim * state.state_kv->nb[0]; - const size_t hi_off_sc = (size_t)ratio * state.state_score->nb[1] + - (size_t)head_dim * state.state_score->nb[0]; - ggml_tensor * prev_kv = ggml_view_2d(ctx, state.state_kv, head_dim, ratio, - state.state_kv->nb[1], 0); - ggml_tensor * cur_kv_hi = ggml_view_2d(ctx, state.state_kv, head_dim, ratio, - state.state_kv->nb[1], hi_off_kv); - ggml_tensor * prev_sc = ggml_view_2d(ctx, state.state_score, head_dim, ratio, - state.state_score->nb[1], 0); - ggml_tensor * cur_sc_hi = ggml_view_2d(ctx, state.state_score, head_dim, ratio, - state.state_score->nb[1], hi_off_sc); + const size_t hi_off_kv = (size_t)ratio * state_kv_source->nb[1] + + (size_t)head_dim * state_kv_source->nb[0]; + const size_t hi_off_sc = (size_t)ratio * state_score_source->nb[1] + + (size_t)head_dim * state_score_source->nb[0]; + ggml_tensor * prev_kv = ggml_view_2d(ctx, state_kv_source, head_dim, ratio, + state_kv_source->nb[1], 0); + ggml_tensor * cur_kv_hi = ggml_view_2d(ctx, state_kv_source, head_dim, ratio, + state_kv_source->nb[1], hi_off_kv); + ggml_tensor * prev_sc = ggml_view_2d(ctx, state_score_source, head_dim, ratio, + state_score_source->nb[1], 0); + ggml_tensor * cur_sc_hi = ggml_view_2d(ctx, state_score_source, head_dim, ratio, + state_score_source->nb[1], hi_off_sc); sv_kv = ggml_concat(ctx, prev_kv, cur_kv_hi, 1); sv_sc = ggml_concat(ctx, prev_sc, cur_sc_hi, 1); n_state_rows = 2 * ratio; } else { - sv_kv = ggml_view_2d(ctx, state.state_kv, comp_width, n_state_rows, - state.state_kv->nb[1], 0); - sv_sc = ggml_view_2d(ctx, state.state_score, comp_width, n_state_rows, - state.state_score->nb[1], 0); + sv_kv = ggml_view_2d(ctx, state_kv_source, comp_width, n_state_rows, + state_kv_source->nb[1], 0); + sv_sc = ggml_view_2d(ctx, state_score_source, comp_width, n_state_rows, + state_score_source->nb[1], 0); } // Transpose to [n_state_rows, comp_width] so softmax operates per-dimension ggml_tensor * sc_T = ggml_cont(ctx, ggml_transpose(ctx, sv_sc)); @@ -643,7 +725,8 @@ static void build_compressor_step( } if (comp_rows_inp) { - ggml_build_forward_expand(gf, ggml_set_rows(ctx, comp_cache, pooled, comp_rows_inp)); + comp_cache_source = ggml_set_rows(ctx, comp_cache, pooled, comp_rows_inp); + ggml_build_forward_expand(gf, comp_cache_source); } else { ggml_tensor * comp_slot = ggml_view_2d( ctx, comp_cache, head_dim, 1, comp_cache->nb[1], @@ -651,6 +734,10 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, pooled_f16, comp_slot)); } + if (comp_cache_source_out) { + *comp_cache_source_out = comp_cache_source; + } + if (ratio == 4) { for (int r = 0; r < ratio; ++r) { ggml_tensor * src_kv = ggml_view_2d(ctx, state.state_kv, comp_width, 1, @@ -715,7 +802,8 @@ static void build_indexer_compressor_step( comp_rows_inp, comp_pos_inp, i64_array_inputs, - i32_array_inputs); + i32_array_inputs, + nullptr); } static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int ratio, int token_pos) { @@ -726,6 +814,18 @@ static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int return std::min(n_cached + grew_this_step, (int) comp_cache->ne[1]); } +// Round the live compressed-row count up to a fixed stride so the fused decode +// graph topology repeats across steps (enabling CUDA/HIP graph replay). The +// rows in [n_comp, padded) are masked to -1e30 in the score matrix, which +// underflows to exactly 0 in softmax, so a padded read is bit-identical to an +// unpadded read of the first n_comp rows. +static constexpr int DS4_COMP_PAD_STRIDE = 16; +static int ds4_padded_comp_rows(int n_comp, int cap) { + if (n_comp <= 0) return 0; + const int padded = ((n_comp + DS4_COMP_PAD_STRIDE - 1) / DS4_COMP_PAD_STRIDE) * DS4_COMP_PAD_STRIDE; + return padded < cap ? padded : cap; +} + static ggml_tensor * build_indexer_score( ggml_context * ctx, ggml_tensor * qr_norm_last, // [n_lora_q, 1] @@ -887,9 +987,11 @@ static ggml_tensor * build_mla_attention( // ── Store ALL KV rows in the raw SWA ring ───────────────────── // For decode (n_tokens=1): write single row. For prefill: write all rows. + ggml_tensor * raw_kv_source = lc.raw_kv; if (cached_inputs && cached_inputs->raw_kv_rows) { ggml_tensor * kv_f32 = ggml_is_contiguous(kv) ? kv : ggml_cont(ctx, kv); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, lc.raw_kv, kv_f32, cached_inputs->raw_kv_rows)); + raw_kv_source = ggml_set_rows(ctx, lc.raw_kv, kv_f32, cached_inputs->raw_kv_rows); + ggml_build_forward_expand(gf, raw_kv_source); } else { for (int ti = 0; ti < n_tokens; ti++) { const int pos_ti = kv_start + ti; @@ -908,6 +1010,7 @@ static ggml_tensor * build_mla_attention( ctx, cur, n_embd, 1, cur->nb[1], (size_t)(n_tokens - 1) * cur->nb[1]); ggml_tensor * qr_last = ggml_view_2d( ctx, qr, n_lora_q, 1, qr->nb[1], (size_t)(n_tokens - 1) * qr->nb[1]); + ggml_tensor * comp_kv_source = lc.comp_kv; if (ratio > 0 && L.attn_compressor_kv) { build_compressor_step(ctx, gf, cur_last, L.attn_compressor_ape, @@ -931,7 +1034,8 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->attn_comp_rows : nullptr, cached_inputs ? cached_inputs->attn_comp_pos : nullptr, i64_array_inputs, - i32_array_inputs); + i32_array_inputs, + &comp_kv_source); } if (ratio == 4 && L.indexer_compressor_kv) { @@ -950,45 +1054,73 @@ static ggml_tensor * build_mla_attention( // raw_kv: [head_dim, n_swa] F16 persistent ring buffer (single KV head, shared) // comp_kv: [head_dim, comp_cap] F16 compressed rows. // n_raw = min(kv_start + n_tokens, n_swa) - const int n_raw = std::min(kv_start + n_tokens, w.n_swa); - const int n_comp_attn = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + const bool masked_kv = (n_tokens == 1) && cached_inputs && cached_inputs->attn_row_mask; + const int n_comp_live = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + // Stable path reads the full physical ring (masking not-yet-written slots) + // and a padded compressed-row span; the plain path reads only valid rows. + const int n_raw = masked_kv ? w.n_swa : std::min(kv_start + n_tokens, w.n_swa); + const int n_comp_attn = masked_kv ? cached_inputs->padded_comp : n_comp_live; + const int n_valid_raw = std::min(kv_start + n_tokens, w.n_swa); const int n_attn = n_raw + n_comp_attn; const float kq_scale = 1.0f / sqrtf((float)head_dim); - // Get valid KV rows. For single-token decode/prefill, include the current - // in-graph KV row directly; otherwise attention can race the side-effecting - // cache write and see the previous contents of the raw KV slot. - ggml_tensor * kv_f32 = nullptr; - if (n_tokens == 1) { - ggml_tensor * cur_kv_f32 = ggml_cast(ctx, kv, GGML_TYPE_F32); + // Get valid KV rows. For single-token decode, include the current in-graph + // KV row directly; otherwise attention can race the side-effecting cache + // write and see the previous contents of the raw KV slot. + auto raw_kv_view = [&](int row, int count) -> ggml_tensor * { + ggml_tensor * view = ggml_view_2d( + ctx, lc.raw_kv, head_dim, count, lc.raw_kv->nb[1], + (size_t)row * lc.raw_kv->nb[1]); + return ds4_cast_if_needed(ctx, view, GGML_TYPE_F32); + }; + + ggml_tensor * kv_attn = nullptr; + if (masked_kv) { + // Fused stable-KV path: read the full physical ring; rows not yet + // written are masked to -1e30 in the score matrix (exact 0 after + // softmax). Only the fused decode graph sets attn_row_mask. Read + // through the set_rows result so the current row's in-graph write is + // ordered before this read. + ggml_tensor * ring = ggml_view_2d( + ctx, raw_kv_source, head_dim, w.n_swa, raw_kv_source->nb[1], 0); + kv_attn = ds4_cast_if_needed(ctx, ring, GGML_TYPE_F32); + } else if (n_tokens == 1) { + ggml_tensor * cur_kv = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); if (n_raw > 1) { - ggml_tensor * prev_f32 = ggml_cast(ctx, - ggml_view_2d(ctx, lc.raw_kv, head_dim, n_raw - 1, - lc.raw_kv->nb[1], 0), - GGML_TYPE_F32); - kv_f32 = ggml_concat(ctx, prev_f32, cur_kv_f32, 1); + ggml_tensor * prev = nullptr; + DeepSeek4RawRingSpan spans[2]; + const int n_spans = + deepseek4_previous_raw_ring_spans(kv_start, w.n_swa, spans); + for (int i = 0; i < n_spans; ++i) { + ggml_tensor * span = raw_kv_view(spans[i].row, spans[i].count); + prev = prev ? ggml_concat(ctx, prev, span, 1) : span; + } + kv_attn = prev ? ggml_concat(ctx, prev, cur_kv, 1) : cur_kv; } else { - kv_f32 = cur_kv_f32; + kv_attn = cur_kv; } } else { - kv_f32 = ggml_cast(ctx, ggml_view_2d(ctx, lc.raw_kv, head_dim, n_raw, - lc.raw_kv->nb[1], 0), GGML_TYPE_F32); + kv_attn = raw_kv_view(0, n_raw); } - if (n_comp_attn > 0 && lc.comp_kv) { - ggml_tensor * comp_f32 = ggml_cast(ctx, - ggml_view_2d(ctx, lc.comp_kv, head_dim, n_comp_attn, lc.comp_kv->nb[1], 0), - GGML_TYPE_F32); - kv_f32 = ggml_concat(ctx, kv_f32, comp_f32, 1); + if (n_comp_attn > 0 && comp_kv_source) { + ggml_tensor * comp = ggml_view_2d(ctx, comp_kv_source, head_dim, n_comp_attn, comp_kv_source->nb[1], 0); + comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); + kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } - // kv_f32: [head_dim, n_attn] + // kv_attn: [head_dim, n_attn] // Flatten q to [head_dim, n_head*n_tokens] for batched matmul ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, n_head * n_tokens); - // Scores: mul_mat(kv_f32, q_flat) = kv_f32^T[n_attn, head_dim] @ q_flat[head_dim, n_head*n_tokens] + // Scores: mul_mat(kv_attn, q_flat) = kv_attn^T[n_attn, head_dim] @ q_flat[head_dim, n_head*n_tokens] // → [n_attn, n_head*n_tokens] - ggml_tensor * scores = ggml_mul_mat(ctx, kv_f32, q_flat); + ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); scores = ggml_scale(ctx, scores, kq_scale); + if (masked_kv) { + // Broadcast-add the [n_attn,1] additive mask across all query columns. + scores = ggml_add(ctx, scores, cached_inputs->attn_row_mask); + } + (void) n_valid_raw; // Sink-aware softmax: DS4 adds one learned per-head sink logit to the // denominator, but the sink contributes no value vector. @@ -1009,8 +1141,8 @@ static ggml_tensor * build_mla_attention( // probs: [n_attn, n_head*n_tokens] // Context: kv_T^T[head_dim, n_attn] @ probs[n_attn, n_head*n_tokens] → [head_dim, n_head*n_tokens] - // i.e. mul_mat(kv_T, probs) where kv_T = cont(transpose(kv_f32)) = [n_raw, head_dim] - ggml_tensor * kv_T = ggml_cont(ctx, ggml_transpose(ctx, kv_f32)); + // i.e. mul_mat(kv_T, probs) where kv_T = cont(transpose(kv_attn)) = [n_raw, head_dim] + ggml_tensor * kv_T = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); ggml_tensor * context = ggml_mul_mat(ctx, kv_T, probs); // context: [head_dim, n_head*n_tokens] @@ -1109,6 +1241,113 @@ struct DeepSeek4CachedDecodeHcPostGraph { } }; +// Per-step decode scalar inputs shared by all cached per-layer decode graphs. +// Values depend only on (kv_start, ratio), so one tensor per slot serves every +// layer with that ratio. i32 layout per ratio-slot: {rope_pos, neg_pos, +// ape_row, comp_pos, index_ape_row, index_comp_pos}; i64 layout: {raw_kv_row, +// state_row, comp_row, index_state_row, index_comp_row}. +struct Ds4DecodeSharedInputs { + static constexpr int MAX_RATIOS = 4; + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + int ratios[MAX_RATIOS] = {0}; + int n_ratios = 0; + ggml_tensor * i32_bundle = nullptr; // [6 * n_ratios] + ggml_tensor * i64_bundle = nullptr; // [5 * n_ratios] + // Per-slot views handed to the graph builders. + ggml_tensor * v_rope_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_neg_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_ape_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_comp_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_ape[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_cpos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_raw_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_state_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_comp_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_state[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_comp[MAX_RATIOS] = {nullptr}; + + void free() { + if (buf) { ggml_backend_buffer_free(buf); buf = nullptr; } + if (ctx) { ggml_free(ctx); ctx = nullptr; } + owner_ctx = nullptr; backend = nullptr; n_ratios = 0; + i32_bundle = nullptr; i64_bundle = nullptr; + } + + int slot(int ratio) const { + for (int i = 0; i < n_ratios; ++i) if (ratios[i] == ratio) return i; + return -1; + } + + bool ensure(const DeepSeek4Weights & w, ggml_backend_t bk) { + if (ctx && owner_ctx == w.ctx && backend == bk) return true; + free(); + n_ratios = 0; + for (int il = 0; il < w.n_layer; ++il) { + const int r = (int) w.compress_ratios[il]; + if (slot(r) >= 0) continue; + if (n_ratios >= MAX_RATIOS) return false; + ratios[n_ratios++] = r; + } + ggml_init_params p{}; + p.mem_size = ggml_tensor_overhead() * (size_t) (2 + 11 * MAX_RATIOS) + 4096; + p.no_alloc = true; + ctx = ggml_init(p); + if (!ctx) return false; + i32_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 6 * (int64_t) n_ratios); + i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 5 * (int64_t) n_ratios); + for (int s = 0; s < n_ratios; ++s) { + v_rope_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 0) * sizeof(int32_t)); + v_neg_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 1) * sizeof(int32_t)); + v_ape_row[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 2) * sizeof(int32_t)); + v_comp_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 3) * sizeof(int32_t)); + v_index_ape[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 4) * sizeof(int32_t)); + v_index_cpos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 5) * sizeof(int32_t)); + v_raw_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 0) * sizeof(int64_t)); + v_state_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 1) * sizeof(int64_t)); + v_comp_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 2) * sizeof(int64_t)); + v_index_state[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 3) * sizeof(int64_t)); + v_index_comp[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 4) * sizeof(int64_t)); + } + buf = ggml_backend_alloc_ctx_tensors(ctx, bk); + if (!buf) { free(); return false; } + owner_ctx = w.ctx; + backend = bk; + return true; + } + + // Upload all per-step values in two writes. + void set_step(const DeepSeek4Weights & w, int kv_start) { + const int token_pos = kv_start; + int32_t i32v[6 * MAX_RATIOS] = {0}; + int64_t i64v[5 * MAX_RATIOS] = {0}; + for (int s = 0; s < n_ratios; ++s) { + const int ratio = ratios[s]; + i32v[s * 6 + 0] = kv_start; + i32v[s * 6 + 1] = -kv_start; + i64v[s * 5 + 0] = kv_start % w.n_swa; + if (ratio > 0) { + const int pos_mod = token_pos % ratio; + i32v[s * 6 + 2] = pos_mod; + i32v[s * 6 + 3] = token_pos + 1 - ratio; + i64v[s * 5 + 1] = (ratio == 4) ? (int64_t) (ratio + pos_mod) : (int64_t) pos_mod; + i64v[s * 5 + 2] = token_pos / ratio; + } + if (ratio == 4) { + const int pos_mod = token_pos % ratio; + i32v[s * 6 + 4] = pos_mod; + i32v[s * 6 + 5] = token_pos + 1 - ratio; + i64v[s * 5 + 3] = ratio + pos_mod; + i64v[s * 5 + 4] = token_pos / ratio; + } + } + ggml_backend_tensor_set(i32_bundle, i32v, 0, sizeof(int32_t) * 6 * (size_t) n_ratios); + ggml_backend_tensor_set(i64_bundle, i64v, 0, sizeof(int64_t) * 5 * (size_t) n_ratios); + } +}; + static bool build_cached_decode_attn_graph( DeepSeek4CachedDecodeAttnGraph & out, ggml_backend_t backend, @@ -1119,7 +1358,8 @@ static bool build_cached_decode_attn_graph( int kv_start, int raw_attn_count, int comp_attn_count, - int index_comp_count) { + int index_comp_count, + const Ds4DecodeSharedInputs * shared = nullptr) { out.free(); const size_t ctx_size = 48 * 1024 * 1024; @@ -1137,6 +1377,8 @@ static bool build_cached_decode_attn_graph( out.n_raw = raw_attn_count; out.n_comp_attn = comp_attn_count; out.n_index_comp = index_comp_count; + out.attn_flush = ratio > 0 && (((kv_start + 1) % ratio) == 0); + out.index_flush = ratio == 4 && (((kv_start + 1) % ratio) == 0); out.compressed = ratio > 0; out.indexed = ratio == 4; @@ -1144,32 +1386,61 @@ static bool build_cached_decode_attn_graph( ggml_set_input(out.sg.inp_embed); out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 2048, false); - out.inputs.rope_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.neg_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - ggml_set_input(out.inputs.rope_pos); - ggml_set_input(out.inputs.neg_pos); - - out.inputs.raw_kv_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - ggml_set_input(out.inputs.raw_kv_rows); - if (ratio > 0) { - out.inputs.attn_ape_row = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.attn_comp_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.attn_state_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - out.inputs.attn_comp_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - ggml_set_input(out.inputs.attn_ape_row); - ggml_set_input(out.inputs.attn_comp_pos); - ggml_set_input(out.inputs.attn_state_rows); - ggml_set_input(out.inputs.attn_comp_rows); - } - if (ratio == 4) { - out.inputs.index_ape_row = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.index_comp_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); - out.inputs.index_state_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - out.inputs.index_comp_rows = ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); - ggml_set_input(out.inputs.index_ape_row); - ggml_set_input(out.inputs.index_comp_pos); - ggml_set_input(out.inputs.index_state_rows); - ggml_set_input(out.inputs.index_comp_rows); + const int shared_slot = shared ? shared->slot(ratio) : -1; + if (shared_slot >= 0) { + out.inputs.rope_pos = shared->v_rope_pos[shared_slot]; + out.inputs.neg_pos = shared->v_neg_pos[shared_slot]; + out.inputs.raw_kv_rows = shared->v_raw_row[shared_slot]; + if (ratio > 0) { + out.inputs.attn_ape_row = shared->v_ape_row[shared_slot]; + out.inputs.attn_comp_pos = shared->v_comp_pos[shared_slot]; + out.inputs.attn_state_rows = shared->v_state_row[shared_slot]; + out.inputs.attn_comp_rows = shared->v_comp_row[shared_slot]; + } + if (ratio == 4) { + out.inputs.index_ape_row = shared->v_index_ape[shared_slot]; + out.inputs.index_comp_pos = shared->v_index_cpos[shared_slot]; + out.inputs.index_state_rows = shared->v_index_state[shared_slot]; + out.inputs.index_comp_rows = shared->v_index_comp[shared_slot]; + } + out.uses_shared_inputs = true; + } else { + out.inputs.rope_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.neg_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + ggml_set_input(out.inputs.rope_pos); + ggml_set_input(out.inputs.neg_pos); + + out.inputs.raw_kv_rows = + ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + ggml_set_input(out.inputs.raw_kv_rows); + if (ratio > 0) { + out.inputs.attn_ape_row = + ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.attn_comp_pos = + ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.attn_state_rows = + ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + out.inputs.attn_comp_rows = + ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + ggml_set_input(out.inputs.attn_ape_row); + ggml_set_input(out.inputs.attn_comp_pos); + ggml_set_input(out.inputs.attn_state_rows); + ggml_set_input(out.inputs.attn_comp_rows); + } + if (ratio == 4) { + out.inputs.index_ape_row = + ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.index_comp_pos = + ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); + out.inputs.index_state_rows = + ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + out.inputs.index_comp_rows = + ggml_new_tensor_2d(out.sg.ctx, GGML_TYPE_I64, 1, 1); + ggml_set_input(out.inputs.index_ape_row); + ggml_set_input(out.inputs.index_comp_pos); + ggml_set_input(out.inputs.index_state_rows); + ggml_set_input(out.inputs.index_comp_rows); + } } std::vector i32_inputs; @@ -1275,6 +1546,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, bool ffn, ggml_tensor * hc_state, ggml_tensor * fn_tensor, + const void * fn_device_override, ggml_tensor * scale_tensor, ggml_tensor * base_tensor, const float * scale_data, @@ -1284,12 +1556,13 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, int sinkhorn_iters, float hc_eps) { #if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) - if (!working || !post || !comb || !hc_state || !fn_tensor || !scale_data || !base_data || - !working->data || !post->data || !comb->data || !hc_state->data || !fn_tensor->data) { + const void * fn_device = fn_device_override ? fn_device_override : (fn_tensor ? fn_tensor->data : nullptr); + if (!working || !post || !comb || !hc_state || !fn_device || !scale_data || !base_data || + !working->data || !post->data || !comb->data || !hc_state->data) { return false; } const bool can_use_device_params = - ds4_backend_is_cuda(backend) && + ds4_backend_is_gpu(backend) && scale_tensor && base_tensor && scale_tensor->data && base_tensor->data && scale_tensor->buffer && base_tensor->buffer && @@ -1297,7 +1570,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, !ggml_backend_buffer_is_host(base_tensor->buffer); if (can_use_device_params) { return deepseek4_cuda_hc_pre_device(hc_state->data, - fn_tensor->data, + fn_device, scale_tensor->data, base_tensor->data, n_embd, @@ -1309,7 +1582,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, comb->data); } return deepseek4_cuda_hc_pre_device_params(hc_state->data, - fn_tensor->data, + fn_device, scale_data, base_data, n_embd, @@ -1328,6 +1601,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, (void) ffn; (void) hc_state; (void) fn_tensor; + (void) fn_device_override; (void) scale_tensor; (void) base_tensor; (void) scale_data; @@ -1689,10 +1963,11 @@ static Ds4MoeRouting build_moe_routing( selection = ggml_add(ctx, selection, L.ffn_exp_probs_b); } - out.selected = ggml_top_k(ctx, selection, w.n_expert_used); + const int k_used = ds4_effective_expert_count(w); + out.selected = ggml_top_k(ctx, selection, k_used); ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, w.n_expert, n_tokens); out.weights = ggml_get_rows(ctx, probs_3d, out.selected); - out.weights = ggml_reshape_2d(ctx, out.weights, w.n_expert_used, n_tokens); + out.weights = ggml_reshape_2d(ctx, out.weights, k_used, n_tokens); ggml_tensor * w_sum = ggml_sum_rows(ctx, out.weights); w_sum = ggml_clamp(ctx, w_sum, 6.103515625e-5f, INFINITY); @@ -1712,7 +1987,7 @@ static ggml_tensor * build_moe_ffn( int n_tokens) { const int n_embd = w.n_embd; - const int n_used = w.n_expert_used; + int n_used = w.n_expert_used; const int n_ff_exp = w.n_ff_exp; ggml_tensor * shared_out = build_shared_ffn(ctx, cur, w, L); ggml_tensor * routed_out = nullptr; @@ -1721,6 +1996,7 @@ static ggml_tensor * build_moe_ffn( routed_out = ggml_scale(ctx, cur, 0.0f); } else { Ds4MoeRouting routing = build_moe_routing(ctx, cur, w, L, n_tokens); + n_used = (int) routing.selected->ne[0]; 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, routing.selected); ggml_tensor * up_e = ggml_mul_mat_id(ctx, L.ffn_up_exps, cur_3d, routing.selected); @@ -1734,8 +2010,6 @@ static ggml_tensor * build_moe_ffn( ggml_tensor * weights_3d = ggml_reshape_3d(ctx, routing.weights, 1, n_used, n_tokens); routed_out = ggml_mul(ctx, down_e, weights_3d); - // Sum over dim-1 (n_used experts): permute [n_embd,n_used,n_tokens] -> [n_used,n_embd,n_tokens], - // then sum_rows reduces dim-0, yielding [1,n_embd,n_tokens], reshape to [n_embd,n_tokens]. routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); routed_out = ggml_sum_rows(ctx, routed_out); routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); @@ -1804,6 +2078,21 @@ struct HashRoutingTableCpu { bool loaded = false; }; +static const int32_t * hash_routing_row( + const HashRoutingTableCpu & table, + int32_t token_id, + int table_width) { + if (!table.loaded || token_id < 0 || table_width <= 0) { + return nullptr; + } + const size_t offset = (size_t) token_id * (size_t) table_width; + if (offset > table.ids.size() || + table.ids.size() - offset < (size_t) table_width) { + return nullptr; + } + return table.ids.data() + offset; +} + static void cpu_rms_norm(float * out, const float * x, int n, float eps) { float ss = 0.0f; for (int i = 0; i < n; i++) ss += x[i] * x[i]; @@ -1811,7 +2100,7 @@ static void cpu_rms_norm(float * out, const float * x, int n, float eps) { for (int i = 0; i < n; i++) out[i] = x[i] * scale; } -static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { +static float cpu_dot_f16_row_scalar(const uint16_t * row, const float * x, int cols) { float acc = 0.0f; for (int c = 0; c < cols; c++) { acc += ggml_fp16_to_fp32(row[c]) * x[c]; @@ -1819,6 +2108,82 @@ static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { return acc; } +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) +static bool ds4_cpu_has_f16c() { + static int supported = -1; + if (supported < 0) { + __builtin_cpu_init(); + supported = (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("f16c")) ? 1 : 0; + } + return supported == 1; +} + +__attribute__((target("avx2,f16c"))) +static float cpu_dot_f16_row_f16c(const uint16_t * row, const float * x, int cols) { + float acc = 0.0f; + int c = 0; + alignas(32) float prod[8]; + for (; c + 7 < cols; c += 8) { + const __m128i h = _mm_loadu_si128(reinterpret_cast(row + c)); + const __m256 wf = _mm256_cvtph_ps(h); + const __m256 xf = _mm256_loadu_ps(x + c); + _mm256_store_ps(prod, _mm256_mul_ps(wf, xf)); + acc += prod[0]; + acc += prod[1]; + acc += prod[2]; + acc += prod[3]; + acc += prod[4]; + acc += prod[5]; + acc += prod[6]; + acc += prod[7]; + } + for (; c < cols; ++c) { + acc += ggml_fp16_to_fp32(row[c]) * x[c]; + } + return acc; +} +#endif + +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) +__attribute__((target("avx2,f16c"))) +static void cpu_dot_f16_rows3_f16c(const uint16_t * r0, const uint16_t * r1, const uint16_t * r2, + const float * x, int cols, + float * o0, float * o1, float * o2) { + float a0 = 0.0f, a1 = 0.0f, a2 = 0.0f; + int c = 0; + alignas(32) float p0[8], p1[8], p2[8]; + for (; c + 7 < cols; c += 8) { + const __m256 xf = _mm256_loadu_ps(x + c); + _mm256_store_ps(p0, _mm256_mul_ps(_mm256_cvtph_ps(_mm_loadu_si128(reinterpret_cast(r0 + c))), xf)); + _mm256_store_ps(p1, _mm256_mul_ps(_mm256_cvtph_ps(_mm_loadu_si128(reinterpret_cast(r1 + c))), xf)); + _mm256_store_ps(p2, _mm256_mul_ps(_mm256_cvtph_ps(_mm_loadu_si128(reinterpret_cast(r2 + c))), xf)); + a0 += p0[0]; a1 += p1[0]; a2 += p2[0]; + a0 += p0[1]; a1 += p1[1]; a2 += p2[1]; + a0 += p0[2]; a1 += p1[2]; a2 += p2[2]; + a0 += p0[3]; a1 += p1[3]; a2 += p2[3]; + a0 += p0[4]; a1 += p1[4]; a2 += p2[4]; + a0 += p0[5]; a1 += p1[5]; a2 += p2[5]; + a0 += p0[6]; a1 += p1[6]; a2 += p2[6]; + a0 += p0[7]; a1 += p1[7]; a2 += p2[7]; + } + for (; c < cols; ++c) { + a0 += ggml_fp16_to_fp32(r0[c]) * x[c]; + a1 += ggml_fp16_to_fp32(r1[c]) * x[c]; + a2 += ggml_fp16_to_fp32(r2[c]) * x[c]; + } + *o0 = a0; *o1 = a1; *o2 = a2; +} +#endif + +static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) + if (ds4_cpu_has_f16c()) { + return cpu_dot_f16_row_f16c(row, x, cols); + } +#endif + return cpu_dot_f16_row_scalar(row, x, cols); +} + static void cpu_matvec_f16_serial(float * out, const uint16_t * mat, const float * x, int rows, int cols) { // mat: [cols, rows] in row-major F16 (ggml layout: ne[0]=cols, ne[1]=rows) // out[r] = dot(mat_row_r, x) for r in [0, rows) @@ -1839,6 +2204,157 @@ static void cpu_matvec_f16(float * out, const uint16_t * mat, const float * x, i }); } +// Persistent worker pool for the decode-path HC fn matvec. Splitting rows +// across threads leaves each row's accumulation order untouched, so results +// are bit-identical to the serial path; only wall time changes. Decode issues +// ~86 of these 24x16384 matvecs per token, so workers spin briefly to catch +// adjacent jobs, then park on a condition variable while the server is idle. +struct Ds4HcMatvecPool { + struct Job { + const uint16_t * mat; + const float * x; + float * out; + int rows; + int cols; + int active_workers; + }; + std::mutex client_mu; + std::mutex wait_mu; + std::condition_variable wait_cv; + std::atomic seq{0}; + std::atomic remaining{0}; + Job job{}; + std::vector workers; + std::atomic stop{false}; + int nth = 0; + + static void cpu_relax() { +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) + __builtin_ia32_pause(); +#endif + } + + Ds4HcMatvecPool() { + unsigned hw = std::thread::hardware_concurrency(); + nth = hw == 0 ? 4 : (int)(hw < 8 ? hw : 8); + for (int i = 0; i < nth; ++i) { + workers.emplace_back([this, i]() { + uint64_t last = 0; + for (;;) { + uint64_t s = last; + for (int spins = 0; spins < 65536; ++spins) { + s = seq.load(std::memory_order_acquire); + if (s != last || stop.load(std::memory_order_relaxed)) { + break; + } + cpu_relax(); + } + if (s == last && !stop.load(std::memory_order_relaxed)) { + std::unique_lock lk(wait_mu); + wait_cv.wait(lk, [&]() { + return stop.load(std::memory_order_relaxed) || + seq.load(std::memory_order_acquire) != last; + }); + s = seq.load(std::memory_order_acquire); + } + if (stop.load(std::memory_order_relaxed)) return; + last = s; + const Job j = job; + if (i >= j.active_workers) continue; + const int chunk = (j.rows + j.active_workers - 1) / j.active_workers; + const int r0 = i * chunk; + const int r1 = j.rows < r0 + chunk ? j.rows : r0 + chunk; + if (row_fn) { + for (int r = r0; r < r1; ++r) row_fn(r); + } else { + int r = r0; +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) + if (ds4_cpu_has_f16c()) { + for (; r + 2 < r1; r += 3) { + cpu_dot_f16_rows3_f16c( + j.mat + (size_t) (r + 0) * j.cols, + j.mat + (size_t) (r + 1) * j.cols, + j.mat + (size_t) (r + 2) * j.cols, + j.x, j.cols, + &j.out[r + 0], &j.out[r + 1], &j.out[r + 2]); + } + } +#endif + for (; r < r1; ++r) { + j.out[r] = cpu_dot_f16_row(j.mat + (size_t) r * j.cols, j.x, j.cols); + } + } + remaining.fetch_sub(1, std::memory_order_acq_rel); + } + }); + } + } + ~Ds4HcMatvecPool() { + { + // Protect predicate changes with the same mutex used by wait(). + // Otherwise a worker can test the predicate, miss notify_all(), + // and sleep forever between the test and the wait. + std::lock_guard lk(wait_mu); + stop.store(true, std::memory_order_release); + } + wait_cv.notify_all(); + for (auto & t : workers) t.join(); + } + void run(const uint16_t * mat, const float * x, float * out, int rows, int cols) { + if (rows <= 0) return; + std::lock_guard lk(client_mu); + row_fn = nullptr; + const int active_workers = std::min(nth, rows); + job = {mat, x, out, rows, cols, active_workers}; + remaining.store(active_workers, std::memory_order_release); + { + // Publish the new generation while holding wait_mu so a worker + // cannot miss the transition between its predicate check and + // blocking in wait(). + std::lock_guard wake_lk(wait_mu); + seq.fetch_add(1, std::memory_order_release); + } + wait_cv.notify_all(); + int spins = 0; + while (remaining.load(std::memory_order_acquire) != 0) { + if (++spins < 65536) { cpu_relax(); } + else { std::this_thread::yield(); spins = 0; } + } + } + + // Generic variant: invoke fn(row) for each row in [0, rows), rows split + // across workers with the same static chunking as run(). + std::function row_fn; + void run_custom(int rows, std::function fn) { + if (rows <= 0) return; + std::lock_guard lk(client_mu); + row_fn = std::move(fn); + const int active_workers = std::min(nth, rows); + job = {nullptr, nullptr, nullptr, rows, 0, active_workers}; + remaining.store(active_workers, std::memory_order_release); + { + std::lock_guard wake_lk(wait_mu); + seq.fetch_add(1, std::memory_order_release); + } + wait_cv.notify_all(); + int spins = 0; + while (remaining.load(std::memory_order_acquire) != 0) { + if (++spins < 65536) { cpu_relax(); } + else { std::this_thread::yield(); spins = 0; } + } + row_fn = nullptr; + } +}; + +static Ds4HcMatvecPool & ds4_hc_matvec_pool() { + static Ds4HcMatvecPool pool; + return pool; +} + +static void cpu_matvec_f16_pooled(float * out, const uint16_t * mat, const float * x, int rows, int cols) { + ds4_hc_matvec_pool().run(mat, x, out, rows, cols); +} + static void cpu_hc_sinkhorn(float * out, const float * mix, const float * scale, const float * base, int n_hc, int iters, float eps) { const float pre_scale = scale[0]; @@ -1969,7 +2485,7 @@ static void cpu_hc_pre_into(float * working, if (serial_fn) { cpu_matvec_f16_serial(mix, fn_data, flat, mix_dim, hc_dim); } else { - cpu_matvec_f16(mix, fn_data, flat, mix_dim, hc_dim); + cpu_matvec_f16_pooled(mix, fn_data, flat, mix_dim, hc_dim); } finish_hc_pre_from_mix_into(working, post, comb, hc_state, mix, scale_data, base_data, @@ -2034,7 +2550,8 @@ static void hc_pre_auto_into(float * working, int sinkhorn_iters, float hc_eps, float * flat, - float * mix_scratch) { + float * mix_scratch, + bool serial_fn) { #if defined(DFLASH27B_BACKEND_CUDA) if (ds4_hc_cuda_enabled() && fn_tensor && fn_tensor->data) { float mix[24]; @@ -2053,7 +2570,7 @@ static void hc_pre_auto_into(float * working, cpu_hc_pre_into(working, post, comb, hc_state, weights.fn_data.data(), weights.scale_data.data(), weights.base_data.data(), - n_embd, n_hc, sinkhorn_iters, hc_eps, flat, mix_scratch, true); + n_embd, n_hc, sinkhorn_iters, hc_eps, flat, mix_scratch, serial_fn); } static void hc_pre_batch(std::vector & working, @@ -2087,7 +2604,8 @@ static void hc_pre_batch(std::vector & working, sinkhorn_iters, hc_eps, flat.data(), - mix); + mix, + /*serial_fn=*/n_tokens > 1); } }); } @@ -2116,6 +2634,23 @@ static void hc_post_batch(std::vector & out_hc, int n_hc) { const size_t hc_dim = (size_t)n_embd * (size_t)n_hc; out_hc.resize((size_t)n_tokens * hc_dim); + if (n_tokens == 1) { + // Decode: split the n_hc independent destination streams across the + // persistent pool. Per-element accumulation order is unchanged, so + // the result is bit-identical to the serial loop. + struct Ctx { const float * block; const float * res; const float * post; const float * comb; float * out; int n_embd; int n_hc; }; + Ctx c{block_out, residual_hc, post, comb, out_hc.data(), n_embd, n_hc}; + ds4_hc_matvec_pool().run_custom(n_hc, [&c](int h) { + for (int d = 0; d < c.n_embd; ++d) { + float acc = c.block[d] * c.post[h]; + for (int src = 0; src < c.n_hc; ++src) { + acc += c.comb[h + src * c.n_hc] * c.res[(size_t)src * c.n_embd + d]; + } + c.out[(size_t)h * c.n_embd + d] = acc; + } + }); + return; + } ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { for (int t = t0; t < t1; ++t) { cpu_hc_post(out_hc.data() + (size_t)t * hc_dim, @@ -2163,18 +2698,88 @@ static void hc_output_batch(std::vector & final_embd, }); } +static bool load_tensor_to_f32_cpu(std::vector & dst, ggml_tensor * t) { + if (!t) return false; + + const size_t elems = ggml_nelements(t); + dst.resize(elems); + if (elems == 0) return true; + + if (t->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(t, dst.data(), 0, ggml_nbytes(t)); + return true; + } + + const ggml_type_traits * tr = ggml_get_type_traits(t->type); + if (!tr || !tr->to_float || t->ne[0] <= 0) return false; + + std::vector raw(ggml_nbytes(t)); + ggml_backend_tensor_get(t, raw.data(), 0, raw.size()); + + const int64_t cols = t->ne[0]; + const int64_t rows = (int64_t)elems / cols; + const size_t row_bytes = ggml_row_size(t->type, cols); + for (int64_t r = 0; r < rows; ++r) { + tr->to_float(raw.data() + (size_t)r * row_bytes, + dst.data() + (size_t)r * (size_t)cols, + cols); + } + return true; +} + +static bool load_tensor_to_f16_cpu(std::vector & dst, ggml_tensor * t) { + if (!t) return false; + + const size_t elems = ggml_nelements(t); + dst.resize(elems); + if (elems == 0) return true; + + if (t->type == GGML_TYPE_F16) { + ggml_backend_tensor_get(t, dst.data(), 0, ggml_nbytes(t)); + return true; + } + + std::vector f32; + if (!load_tensor_to_f32_cpu(f32, t)) return false; + ggml_fp32_to_fp16_row(f32.data(), reinterpret_cast(dst.data()), (int64_t)elems); + return true; +} + static void load_hc_weights_cpu(HcWeightsCpu & dst, ggml_tensor * fn, ggml_tensor * scale, ggml_tensor * base) { if (!fn || !scale || !base || dst.loaded) return; - dst.fn_data.resize(ggml_nelements(fn)); - dst.scale_data.resize(ggml_nelements(scale)); - dst.base_data.resize(ggml_nelements(base)); - ggml_backend_tensor_get(fn, dst.fn_data.data(), 0, ggml_nbytes(fn)); - ggml_backend_tensor_get(scale, dst.scale_data.data(), 0, ggml_nbytes(scale)); - ggml_backend_tensor_get(base, dst.base_data.data(), 0, ggml_nbytes(base)); + if (!load_tensor_to_f16_cpu(dst.fn_data, fn) || + !load_tensor_to_f32_cpu(dst.scale_data, scale) || + !load_tensor_to_f32_cpu(dst.base_data, base)) { + dst.fn_data.clear(); + dst.scale_data.clear(); + dst.base_data.clear(); + return; + } dst.loaded = true; } +static void reset_hc_weights_cpu(HcWeightsCpu & w) { + w.fn_data.clear(); + w.scale_data.clear(); + w.base_data.clear(); + w.loaded = false; +} + +static void reset_hc_layer_weights_cpu(std::vector & weights) { + for (HcLayerWeightsCpu & layer : weights) { + reset_hc_weights_cpu(layer.attn); + reset_hc_weights_cpu(layer.ffn); + } + weights.clear(); +} + +static const void * hc_fn_device_ptr(const HcWeightsCpu &, ggml_tensor * fn) { + if (!fn) return nullptr; + if (fn->type == GGML_TYPE_F16) return fn->data; + return nullptr; +} + static bool load_hash_routing_cpu(HashRoutingTableCpu & dst, ggml_tensor * table) { if (dst.loaded) return true; if (!table) return false; @@ -2405,33 +3010,34 @@ static bool deepseek4_step_hybrid( return false; } - std::vector selected_host((size_t)w.n_expert_used * (size_t)n_tokens); - std::vector weights_host((size_t)w.n_expert_used * (size_t)n_tokens); + const int n_used = ds4_effective_expert_count(w); + std::vector selected_host((size_t)n_used * (size_t)n_tokens); + std::vector weights_host((size_t)n_used * (size_t)n_tokens); const auto route_select_t0 = Ds4TimingClock::now(); - const auto & hash_table = hash_routing_tables[(size_t)il].ids; for (int ti = 0; ti < n_tokens; ++ti) { const int32_t tok = token_ids[ti]; - if (tok < 0 || tok >= w.n_vocab) { + const int32_t * row = hash_routing_row( + hash_routing_tables[(size_t)il], tok, w.n_expert_used); + if (!row) { std::fprintf(stderr, "[deepseek4] token id %d outside hash table for layer %d\n", tok, il); if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; } - const int32_t * row = hash_table.data() + (size_t)tok * (size_t)w.n_expert_used; float sum = 0.0f; - for (int ei = 0; ei < w.n_expert_used; ++ei) { + for (int ei = 0; ei < n_used; ++ei) { const int32_t expert = row[ei]; - selected_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei] = expert; + selected_host[(size_t)ti * (size_t)n_used + (size_t)ei] = expert; float prob = 0.0f; if (expert >= 0 && expert < w.n_expert) { prob = probs_host[(size_t)ti * (size_t)w.n_expert + (size_t)expert]; } - weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei] = prob; + weights_host[(size_t)ti * (size_t)n_used + (size_t)ei] = prob; sum += prob; } sum = std::max(sum, 6.103515625e-5f); - for (int ei = 0; ei < w.n_expert_used; ++ei) { - float & weight = weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)ei]; + for (int ei = 0; ei < n_used; ++ei) { + float & weight = weights_host[(size_t)ti * (size_t)n_used + (size_t)ei]; weight = weight / sum * w.expert_weight_scale; } } @@ -2439,8 +3045,8 @@ static bool deepseek4_step_hybrid( if (routing_stats) { for (int ti = 0; ti < n_tokens; ++ti) { routing_stats->observe(il, - selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, - w.n_expert_used); + selected_host.data() + (size_t)ti * (size_t)n_used, + n_used); } } @@ -2449,7 +3055,7 @@ static bool deepseek4_step_hybrid( auto & storage = moe_hybrid.layers[(size_t) il]; if (!eval_ds4_hybrid( backend, cpu_backend, hybrid_cfg, desc, &moe_hybrid, storage, stream_engine, - il, n_embd, w.n_expert_used, + il, n_embd, n_used, ffn_normed_host.data(), selected_host.data(), weights_host.data(), n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, telemetry)) { if (hot_alloc) ggml_gallocr_free(hot_alloc); @@ -2498,8 +3104,9 @@ static bool deepseek4_step_hybrid( std::vector ffn_normed_host((size_t)n_embd * (size_t)n_tokens); std::vector probs_host((size_t)w.n_expert * (size_t)n_tokens); - std::vector selected_host((size_t)w.n_expert_used * (size_t)n_tokens); - std::vector weights_host((size_t)w.n_expert_used * (size_t)n_tokens); + const int n_used = ds4_effective_expert_count(w); + std::vector selected_host((size_t)n_used * (size_t)n_tokens); + std::vector weights_host((size_t)n_used * (size_t)n_tokens); const auto route_read_t0 = Ds4TimingClock::now(); ggml_backend_tensor_get(ffn_normed, ffn_normed_host.data(), 0, sizeof(float) * ffn_normed_host.size()); ggml_backend_tensor_get(router_probs, probs_host.data(), 0, sizeof(float) * probs_host.size()); @@ -2516,18 +3123,18 @@ static bool deepseek4_step_hybrid( } for (int ti = 0; ti < n_tokens; ++ti) { const float * probs = probs_host.data() + (size_t)ti * (size_t)w.n_expert; - std::vector top((size_t)w.n_expert_used, -1); + std::vector top((size_t)n_used, -1); for (int expert = 0; expert < w.n_expert; ++expert) { const float score = probs[expert] + (!bias_host.empty() ? bias_host[(size_t)expert] : 0.0f); - for (int slot = 0; slot < w.n_expert_used; ++slot) { + for (int slot = 0; slot < n_used; ++slot) { const int32_t cur_expert = top[(size_t)slot]; const float cur_score = cur_expert >= 0 ? probs[cur_expert] + (!bias_host.empty() ? bias_host[(size_t)cur_expert] : 0.0f) : -INFINITY; if (cur_expert < 0 || score > cur_score) { - for (int m = w.n_expert_used - 1; m > slot; --m) { + for (int m = n_used - 1; m > slot; --m) { top[(size_t)m] = top[(size_t)m - 1]; } top[(size_t)slot] = expert; @@ -2536,16 +3143,16 @@ static bool deepseek4_step_hybrid( } } float sum = 0.0f; - for (int slot = 0; slot < w.n_expert_used; ++slot) { + for (int slot = 0; slot < n_used; ++slot) { const int32_t expert = top[(size_t)slot]; - selected_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot] = expert; + selected_host[(size_t)ti * (size_t)n_used + (size_t)slot] = expert; const float weight = expert >= 0 ? probs[expert] : 0.0f; - weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot] = weight; + weights_host[(size_t)ti * (size_t)n_used + (size_t)slot] = weight; sum += weight; } sum = std::max(sum, 6.103515625e-5f); - for (int slot = 0; slot < w.n_expert_used; ++slot) { - float & weight = weights_host[(size_t)ti * (size_t)w.n_expert_used + (size_t)slot]; + for (int slot = 0; slot < n_used; ++slot) { + float & weight = weights_host[(size_t)ti * (size_t)n_used + (size_t)slot]; weight = weight / sum * w.expert_weight_scale; } } @@ -2553,8 +3160,8 @@ static bool deepseek4_step_hybrid( if (routing_stats) { for (int ti = 0; ti < n_tokens; ++ti) { routing_stats->observe(il, - selected_host.data() + (size_t)ti * (size_t)w.n_expert_used, - w.n_expert_used); + selected_host.data() + (size_t)ti * (size_t)n_used, + n_used); } } @@ -2563,7 +3170,7 @@ static bool deepseek4_step_hybrid( auto & storage = moe_hybrid.layers[(size_t) il]; if (!eval_ds4_hybrid( backend, cpu_backend, hybrid_cfg, desc, &moe_hybrid, storage, stream_engine, - il, n_embd, w.n_expert_used, + il, n_embd, n_used, ffn_normed_host.data(), selected_host.data(), weights_host.data(), n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, telemetry)) { if (hot_alloc) ggml_gallocr_free(hot_alloc); @@ -2669,141 +3276,605 @@ bool deepseek4_step( MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats) { - if (w.moe_hybrid && moe_hybrid != nullptr) { return deepseek4_step_hybrid(backend, w, cache, *moe_hybrid, embed, n_tokens, kv_start, out_logits, token_ids, stream_engine, telemetry, routing_stats); } - const int n_embd = w.n_embd; - const int n_layer = w.n_layer; + std::vector hc_state; + return deepseek4_step_layer_range( + backend, w, cache, hc_state, embed, n_tokens, kv_start, + 0, w.n_layer, &out_logits, token_ids, telemetry); +} - // Create compute graph context — need large budget for MoE layers - const size_t ctx_size = ggml_tensor_overhead() * 65536 + 16 * 1024 * 1024; - ggml_init_params params{}; - params.mem_size = ctx_size; - params.mem_buffer = nullptr; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - if (!ctx) return false; +// ─── Fused single-graph decode (n_tokens == 1) ────────────────────────── +// Chains all layers (HC pre → attention → HC post → HC pre → FFN → HC post) +// plus the output HC merge and lm_head into ONE cached ggml graph, so a +// decode step is a single ggml_backend_graph_compute with one logits +// readback instead of ~90 per-layer graph launches with host round-trips. +// HC Sinkhorn mixing runs in the fused GGML_OP_DS4_HC op (one kernel per +// sublayer instead of ~170 tiny ops for 20 Sinkhorn iterations). +// +// Compressed-KV reads are padded to a fixed stride with an additive +// score mask, and each structural variant (flush pattern) lives in its own +// slot with a private metadata arena. Tensor addresses therefore stay stable +// while a variant recurs, which is what the ggml-cuda/HIP graph cache keys +// on, enabling graph replay for the bulk of decode steps. + +struct DeepSeek4FusedDecodeGraph { + std::vector shape_key; + uint64_t last_use = 0; + StepGraph sg; + ggml_tensor * inp_embed = nullptr; + ggml_tensor * i32_bundle = nullptr; + ggml_tensor * i64_bundle = nullptr; + ggml_tensor * mask_bundle = nullptr; // additive score mask (0 / -1e30), may be null + std::vector hash_ids; + ggml_tensor * logits = nullptr; + + void reset_nodes() { + inp_embed = nullptr; + i32_bundle = nullptr; + i64_bundle = nullptr; + mask_bundle = nullptr; + logits = nullptr; + hash_ids.clear(); + shape_key.clear(); + last_use = 0; + } - // Input embeddings - ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); - ggml_set_name(inp, "inp_embed"); - ggml_set_input(inp); + bool built() const { + return sg.ctx && sg.gf && logits; + } - ggml_tensor * cur = inp; - ggml_cgraph * gf = ggml_new_graph_custom(ctx, 32768, false); - std::vector i32_inputs; - std::vector i32_array_inputs; - std::vector i64_array_inputs; + void destroy() { + step_graph_destroy(sg); + reset_nodes(); + } +}; - // Layer loop - for (int il = 0; il < n_layer; il++) { - const DeepSeek4Layer & L = w.layers[il]; - DeepSeek4LayerCache & lc = cache.layers[il]; +struct DeepSeek4FusedDecodeCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + bool disabled = false; + uint64_t counter = 0; + std::array slots; + + // Persistent F16 mirrors of the (quantized) HC fn projection weights so + // the fused graph matches the numerics of the reference HC paths, which + // always dequantize fn to F16 before the mix matvec. + ggml_context * fn_ctx = nullptr; + ggml_backend_buffer_t fn_buf = nullptr; + std::vector fn_attn_f16; + std::vector fn_ffn_f16; + ggml_tensor * fn_out_f16 = nullptr; + + void destroy() { + for (auto & s : slots) s.destroy(); + if (fn_buf) { ggml_backend_buffer_free(fn_buf); fn_buf = nullptr; } + if (fn_ctx) { ggml_free(fn_ctx); fn_ctx = nullptr; } + fn_attn_f16.clear(); + fn_ffn_f16.clear(); + fn_out_f16 = nullptr; + owner_ctx = nullptr; + backend = nullptr; + disabled = false; + counter = 0; + } +}; - // ── HC pre (attention) ────────────────────────────────────── - // TODO: Full HC implementation. For now, pass cur through directly. - ggml_tensor * attn_in = cur; +static ggml_tensor * ds4_fused_hc_base_f32(ggml_context * ctx, ggml_tensor * base) { + if (!base) return nullptr; + ggml_tensor * b = base; + if (b->type != GGML_TYPE_F32) { + b = ggml_cast(ctx, b, GGML_TYPE_F32); + } + return ggml_reshape_1d(ctx, b, ggml_nelements(b)); +} - // ── Attention norm ────────────────────────────────────────── - ggml_tensor * normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); +static ggml_tensor * ds4_build_fused_hc_pre( + ggml_context * ctx, + const DeepSeek4Weights & w, + ggml_tensor * hc_flat, // [n_embd*n_hc] contiguous f32 + ggml_tensor * fn, + ggml_tensor * base, + const HcWeightsCpu & cw, + ggml_tensor ** out_split) { + if (!fn || !base || !cw.loaded || cw.scale_data.size() < 3) return nullptr; + const int mix_dim = 2 * w.n_hc + w.n_hc * w.n_hc; + ggml_tensor * normed = ggml_rms_norm(ctx, hc_flat, w.hc_eps); + ggml_tensor * mix = ggml_mul_mat(ctx, fn, normed); + mix = ggml_reshape_1d(ctx, mix, mix_dim); + ggml_tensor * base_f32 = ds4_fused_hc_base_f32(ctx, base); + ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base_f32, hc_flat, + w.n_hc, w.n_hc_sinkhorn_iter, + cw.scale_data[0], cw.scale_data[1], cw.scale_data[2]); + *out_split = ggml_view_1d(ctx, pre, mix_dim, (size_t) w.n_embd * sizeof(float)); + return ggml_view_1d(ctx, pre, w.n_embd, 0); +} - // ── MLA attention ─────────────────────────────────────────── - ggml_tensor * attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, - il, kv_start, n_tokens, - nullptr, i32_inputs, i32_array_inputs, - i64_array_inputs); +static ggml_tensor * ds4_build_hash_routed_ffn( + ggml_context * ctx, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + ggml_tensor * ffn_normed, + ggml_tensor * hash_ids) { + ggml_tensor * shared_out = build_shared_ffn(ctx, ffn_normed, w, L); + ggml_tensor * logits = ggml_mul_mat(ctx, L.ffn_gate_inp, ffn_normed); + ggml_tensor * probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); - // ── Residual ──────────────────────────────────────────────── - cur = ggml_add(ctx, cur, attn_out); + const int n_used = (int) hash_ids->ne[0]; + const int n_ff_exp = w.n_ff_exp; + ggml_tensor * cur_3d = ggml_reshape_3d(ctx, ffn_normed, w.n_embd, 1, 1); + ggml_tensor * gate_e = ggml_mul_mat_id(ctx, L.ffn_gate_exps, cur_3d, hash_ids); + ggml_tensor * up_e = ggml_mul_mat_id(ctx, L.ffn_up_exps, cur_3d, hash_ids); + gate_e = ggml_reshape_3d(ctx, gate_e, n_ff_exp, n_used, 1); + up_e = ggml_reshape_3d(ctx, up_e, n_ff_exp, n_used, 1); + 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, hash_ids); + down_e = ggml_reshape_3d(ctx, down_e, w.n_embd, n_used, 1); + + ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, w.n_expert, 1); + ggml_tensor * weights = ggml_get_rows(ctx, probs_3d, hash_ids); + weights = ggml_reshape_2d(ctx, weights, n_used, 1); + ggml_tensor * w_sum = ggml_sum_rows(ctx, weights); + w_sum = ggml_clamp(ctx, w_sum, 6.103515625e-5f, INFINITY); + weights = ggml_div(ctx, weights, w_sum); + if (w.expert_weight_scale != 1.0f) { + weights = ggml_scale(ctx, weights, w.expert_weight_scale); + } - // ── HC pre (FFN) ──────────────────────────────────────────── - ggml_tensor * ffn_in = cur; + ggml_tensor * weights_3d = ggml_reshape_3d(ctx, weights, 1, n_used, 1); + ggml_tensor * routed_out = ggml_mul(ctx, down_e, weights_3d); + routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(ctx, routed_out); + routed_out = ggml_reshape_2d(ctx, routed_out, w.n_embd, 1); + return ggml_add(ctx, shared_out, routed_out); +} - // ── FFN norm ──────────────────────────────────────────────── - ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); +static bool ds4_fused_ensure_fn_mirrors( + DeepSeek4FusedDecodeCache & fc, + ggml_backend_t backend, + const DeepSeek4Weights & w, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights) { + if (fc.fn_ctx && fc.fn_buf && fc.fn_attn_f16.size() == (size_t) w.n_layer && fc.fn_out_f16) { + return true; + } + const int64_t hc_dim = (int64_t) w.n_embd * w.n_hc; + const int64_t mix_dim = 2 * (int64_t) w.n_hc + (int64_t) w.n_hc * w.n_hc; + if (hc_weights.size() != (size_t) w.n_layer) { + return false; + } + for (int il = 0; il < w.n_layer; ++il) { + const auto & attn = hc_weights[(size_t) il].attn.fn_data; + const auto & ffn = hc_weights[(size_t) il].ffn.fn_data; + if ((int64_t) attn.size() != hc_dim * mix_dim || + (int64_t) ffn.size() != hc_dim * mix_dim) { + return false; + } + } + if ((int64_t) hc_out_weights.fn_data.size() != hc_dim * w.n_hc) { + return false; + } - // ── MoE FFN ───────────────────────────────────────────────── - ggml_tensor * ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, n_tokens); + if (fc.fn_buf) { ggml_backend_buffer_free(fc.fn_buf); fc.fn_buf = nullptr; } + if (fc.fn_ctx) { ggml_free(fc.fn_ctx); fc.fn_ctx = nullptr; } + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * (size_t) (2 * w.n_layer + 4) + 4096; + params.mem_buffer = nullptr; + params.no_alloc = true; + fc.fn_ctx = ggml_init(params); + if (!fc.fn_ctx) return false; + fc.fn_attn_f16.assign((size_t) w.n_layer, nullptr); + fc.fn_ffn_f16.assign((size_t) w.n_layer, nullptr); + for (int il = 0; il < w.n_layer; ++il) { + fc.fn_attn_f16[(size_t) il] = ggml_new_tensor_2d(fc.fn_ctx, GGML_TYPE_F16, hc_dim, mix_dim); + fc.fn_ffn_f16[(size_t) il] = ggml_new_tensor_2d(fc.fn_ctx, GGML_TYPE_F16, hc_dim, mix_dim); + } + fc.fn_out_f16 = ggml_new_tensor_2d(fc.fn_ctx, GGML_TYPE_F16, hc_dim, w.n_hc); + fc.fn_buf = ggml_backend_alloc_ctx_tensors(fc.fn_ctx, backend); + if (!fc.fn_buf) { + ggml_free(fc.fn_ctx); + fc.fn_ctx = nullptr; + return false; + } + for (int il = 0; il < w.n_layer; ++il) { + const auto & a = hc_weights[(size_t) il].attn.fn_data; + const auto & f = hc_weights[(size_t) il].ffn.fn_data; + ggml_backend_tensor_set(fc.fn_attn_f16[(size_t) il], a.data(), 0, a.size() * sizeof(uint16_t)); + ggml_backend_tensor_set(fc.fn_ffn_f16[(size_t) il], f.data(), 0, f.size() * sizeof(uint16_t)); + } + const auto & o = hc_out_weights.fn_data; + ggml_backend_tensor_set(fc.fn_out_f16, o.data(), 0, o.size() * sizeof(uint16_t)); + return true; +} - // ── Residual ──────────────────────────────────────────────── - cur = ggml_add(ctx, cur, ffn_out); +static bool ds4_build_fused_decode_graph( + DeepSeek4FusedDecodeCache & fc, + DeepSeek4FusedDecodeGraph & fg, + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, + int kv_start, + bool have_token_ids, + std::vector && shape_key) { + step_graph_free(fg.sg); + fg.reset_nodes(); + fg.hash_ids.assign((size_t) w.n_layer, nullptr); + + const int n_embd = w.n_embd; + const int n_hc = w.n_hc; + const int token_pos = kv_start; + const size_t arena_size = 192u * 1024 * 1024; + if (fg.sg.meta_arena.size() < arena_size) { + fg.sg.meta_arena.resize(arena_size); } + ggml_init_params params{}; + params.mem_size = fg.sg.meta_arena.size(); + params.mem_buffer = fg.sg.meta_arena.data(); + params.no_alloc = true; + fg.sg.ctx = ggml_init(params); + if (!fg.sg.ctx) return false; + ggml_context * ctx = fg.sg.ctx; + fg.sg.gf = ggml_new_graph_custom(ctx, 32768, false); + ggml_cgraph * gf = fg.sg.gf; + + fg.inp_embed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, 1); + ggml_set_input(fg.inp_embed); + fg.i32_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 6 * (int64_t) w.n_layer); + ggml_set_input(fg.i32_bundle); + fg.i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 5 * (int64_t) w.n_layer); + ggml_set_input(fg.i64_bundle); + + // One additive score-mask bundle covering EVERY layer: [n_swa raw rows ++ + // padded comp rows]. All layers take the masked full-ring attention branch + // so the graph topology never depends on the live raw-row count. + int64_t mask_total = 0; + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + int padded = 0; + if (ratio > 0 && cache.layers[(size_t) il].comp_kv) { + const int n_comp = ds4_comp_rows_used(cache.layers[(size_t) il].comp_kv, + cache.layers[(size_t) il].n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) cache.layers[(size_t) il].comp_kv->ne[1]); + } + mask_total += w.n_swa + padded; + } + fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mask_total); + ggml_set_input(fg.mask_bundle); + int64_t mask_off = 0; - // ── Output head ───────────────────────────────────────────────────── - // TODO: HC output pre (merge residual streams for final projection) + // hc state starts as the token embedding replicated into every stream + ggml_tensor * hc_cur = ggml_repeat_4d(ctx, fg.inp_embed, n_embd, n_hc, 1, 1); - // Final RMSNorm - cur = build_rms_norm(ctx, cur, w.out_norm, w.rms_eps); + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t) il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + const int ratio = (int) w.compress_ratios[il]; + + // ── HC pre (attention) ───────────────────────────────────── + ggml_tensor * hc_flat = ggml_reshape_1d(ctx, hc_cur, (int64_t) n_embd * n_hc); + ggml_tensor * split_attn = nullptr; + ggml_tensor * working = ds4_build_fused_hc_pre(ctx, w, hc_flat, + fc.fn_attn_f16[(size_t) il], L.hc_attn_base, + hlw.attn, &split_attn); + if (!working) return false; + ggml_tensor * attn_in = ggml_reshape_2d(ctx, working, n_embd, 1); + + // ── Attention (inputs are views into the shared bundles) ── + DeepSeek4AttentionGraphInputs ain{}; + ain.rope_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 0) * sizeof(int32_t)); + ain.neg_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 1) * sizeof(int32_t)); + ain.raw_kv_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 0) * sizeof(int64_t)); + if (ratio > 0) { + ain.attn_ape_row = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 2) * sizeof(int32_t)); + ain.attn_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 3) * sizeof(int32_t)); + ain.attn_state_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 1) * sizeof(int64_t)); + ain.attn_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 2) * sizeof(int64_t)); + } + if (ratio == 4) { + ain.index_ape_row = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 4) * sizeof(int32_t)); + ain.index_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 5) * sizeof(int32_t)); + ain.index_state_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 3) * sizeof(int64_t)); + ain.index_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 4) * sizeof(int64_t)); + } + { + int padded = 0; + if (ratio > 0 && lc.comp_kv) { + const int n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + const int64_t n_attn = (int64_t) w.n_swa + padded; + ain.attn_row_mask = ggml_view_2d(ctx, fg.mask_bundle, n_attn, 1, + n_attn * sizeof(float), + (size_t) mask_off * sizeof(float)); + ain.padded_comp = padded; + mask_off += n_attn; + } - // lm_head projection - ggml_tensor * logits = ggml_mul_mat(ctx, w.output, cur); - ggml_set_name(logits, "logits"); - ggml_set_output(logits); + std::vector i32b; + std::vector i32ab; + std::vector i64ab; + ggml_tensor * normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); + ggml_tensor * attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, + kv_start, 1, &ain, + i32b, i32ab, i64ab); + if (!attn_out) return false; + if (!i32b.empty() || !i32ab.empty() || !i64ab.empty()) { + std::fprintf(stderr, + "[deepseek4] fused decode: layer %d created %zu/%zu/%zu dynamic bindings; cannot fuse\n", + il, i32b.size(), i32ab.size(), i64ab.size()); + return false; + } - // ── Build and run graph ───────────────────────────────────────────── - ggml_build_forward_expand(gf, logits); + // ── HC post (attention) ──────────────────────────────────── + ggml_tensor * attn_out_flat = ggml_reshape_1d(ctx, attn_out, n_embd); + ggml_tensor * hc_next = ggml_ds4_hc_post(ctx, hc_flat, attn_out_flat, split_attn, n_hc); + hc_cur = ggml_reshape_2d(ctx, hc_next, n_embd, n_hc); + + // ── HC pre (FFN) ─────────────────────────────────────────── + hc_flat = ggml_reshape_1d(ctx, hc_cur, (int64_t) n_embd * n_hc); + ggml_tensor * split_ffn = nullptr; + ggml_tensor * fworking = ds4_build_fused_hc_pre(ctx, w, hc_flat, + fc.fn_ffn_f16[(size_t) il], L.hc_ffn_base, + hlw.ffn, &split_ffn); + if (!fworking) return false; + ggml_tensor * ffn_in = ggml_reshape_2d(ctx, fworking, n_embd, 1); + + // ── FFN ──────────────────────────────────────────────────── + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); + ggml_tensor * ffn_out = nullptr; + const bool hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && + have_token_ids && hash_tables[(size_t) il].loaded; + if (hash_routed) { + ggml_tensor * hids = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, ds4_effective_expert_count(w), 1); + ggml_set_input(hids); + fg.hash_ids[(size_t) il] = hids; + ffn_out = ds4_build_hash_routed_ffn(ctx, w, L, ffn_normed, hids); + } else { + ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, 1); + } + if (!ffn_out) return false; - // Allocate - ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); - if (!ggml_gallocr_alloc_graph(alloc, gf)) { - std::fprintf(stderr, "[deepseek4] graph allocation failed\n"); - ggml_gallocr_free(alloc); - ggml_free(ctx); + // ── HC post (FFN) ────────────────────────────────────────── + ggml_tensor * ffn_out_flat = ggml_reshape_1d(ctx, ffn_out, n_embd); + hc_next = ggml_ds4_hc_post(ctx, hc_flat, ffn_out_flat, split_ffn, n_hc); + hc_cur = ggml_reshape_2d(ctx, hc_next, n_embd, n_hc); + } + + // ── Output: HC merge → norm → lm_head ────────────────────────── + ggml_tensor * hc_flat = ggml_reshape_1d(ctx, hc_cur, (int64_t) n_embd * n_hc); + ggml_tensor * onorm = ggml_rms_norm(ctx, hc_flat, w.hc_eps); + ggml_tensor * omix = ggml_mul_mat(ctx, fc.fn_out_f16, onorm); + omix = ggml_reshape_1d(ctx, omix, ggml_nelements(omix)); + ggml_tensor * obase = ds4_fused_hc_base_f32(ctx, w.output_hc_base); + if (!obase || hc_out_weights.scale_data.empty()) return false; + ggml_tensor * final_embd = ggml_ds4_hc_out(ctx, omix, obase, hc_flat, n_hc, + hc_out_weights.scale_data[0]); + ggml_tensor * final_2d = ggml_reshape_2d(ctx, final_embd, n_embd, 1); + ggml_tensor * out_normed = build_rms_norm(ctx, final_2d, w.out_norm, w.rms_eps); + fg.logits = ggml_mul_mat(ctx, w.output, out_normed); + ggml_set_output(fg.logits); + ggml_build_forward_expand(gf, fg.logits); + + if (!fg.sg.alloc) { + fg.sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + if (!fg.sg.alloc || !ggml_gallocr_alloc_graph(fg.sg.alloc, fg.sg.gf)) { + std::fprintf(stderr, "[deepseek4] fused decode graph alloc failed\n"); return false; } - // Set input data - ggml_backend_tensor_set(inp, embed, 0, n_embd * n_tokens * sizeof(float)); - for (const DeepSeek4I32InputBinding & binding : i32_inputs) { - ggml_backend_tensor_set(binding.tensor, &binding.value, 0, sizeof(binding.value)); + fg.shape_key = std::move(shape_key); + return true; +} + +// Returns 1 on success (out_logits filled), 0 to fall back to the per-layer +// path, -1 on a hard failure after cache state may have been touched. +static int ds4_try_fused_decode_step( + DeepSeek4FusedDecodeCache & fc, + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, + std::vector & hash_scratch, + const float * embed, + int kv_start, + std::vector & out_logits, + const int32_t * token_ids, + DeepSeek4StepTelemetry * telemetry) { + if (fc.disabled) return 0; + if (!hc_out_weights.loaded || hc_out_weights.scale_data.empty() || + !w.output_hc_fn || !w.output_hc_base) { + fc.disabled = true; + return 0; } - for (const DeepSeek4I32ArrayBinding & binding : i32_array_inputs) { - ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, - sizeof(int32_t) * binding.values.size()); + for (int il = 0; il < w.n_layer; ++il) { + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + const DeepSeek4Layer & L = w.layers[(size_t) il]; + if (!hlw.attn.loaded || hlw.attn.scale_data.size() < 3 || + !hlw.ffn.loaded || hlw.ffn.scale_data.size() < 3 || + !L.hc_attn_fn || !L.hc_ffn_fn || !L.hc_attn_base || !L.hc_ffn_base) { + fc.disabled = true; + return 0; + } + } + + if (fc.owner_ctx != w.ctx || fc.backend != backend) { + fc.destroy(); + fc.owner_ctx = w.ctx; + fc.backend = backend; } - for (const DeepSeek4I64ArrayBinding & binding : i64_array_inputs) { - ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, - sizeof(int64_t) * binding.values.size()); + if (!ds4_fused_ensure_fn_mirrors(fc, backend, w, hc_weights, hc_out_weights)) { + std::fprintf(stderr, "[deepseek4] fused decode: HC fn mirror upload failed; using per-layer path\n"); + fc.disabled = true; + return 0; } - // Compute - if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "[deepseek4] graph compute failed\n"); - ggml_gallocr_free(alloc); - ggml_free(ctx); - return false; + const int token_pos = kv_start; + std::vector key; + key.reserve((size_t) w.n_layer + 2); + key.push_back(w.n_swa); + key.push_back(token_ids ? 1 : 0); + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + int padded = 0; + if (ratio > 0 && lc.comp_kv) { + const int n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + const bool flush = ratio > 0 && (((token_pos + 1) % ratio) == 0); + key.push_back(((int64_t) padded << 1) | (flush ? 1 : 0)); } - // Read logits (only last token for generation) - out_logits.resize(w.n_vocab); - const size_t logits_offset = (size_t)(n_tokens - 1) * w.n_vocab * sizeof(float); - ggml_backend_tensor_get(logits, out_logits.data(), logits_offset, - w.n_vocab * sizeof(float)); + // Pick the slot whose shape key matches; otherwise rebuild the LRU slot. + fc.counter++; + DeepSeek4FusedDecodeGraph * fg = nullptr; + for (auto & s : fc.slots) { + if (s.built() && s.shape_key == key) { + fg = &s; + break; + } + } + if (!fg) { + for (auto & s : fc.slots) { + if (!s.built()) { fg = &s; break; } + } + if (!fg) { + fg = &fc.slots[0]; + for (auto & s : fc.slots) { + if (s.last_use < fg->last_use) fg = &s; + } + } + const auto build_t0 = Ds4TimingClock::now(); + if (!ds4_build_fused_decode_graph(fc, *fg, backend, w, cache, + hc_weights, hc_out_weights, hash_tables, + kv_start, token_ids != nullptr, std::move(key))) { + std::fprintf(stderr, + "[deepseek4] fused decode graph build failed; using per-layer path\n"); + step_graph_free(fg->sg); + fg->reset_nodes(); + fc.disabled = true; + return 0; + } + if (telemetry) telemetry->full_graph_build_us += ds4_elapsed_us(build_t0, Ds4TimingClock::now()); + } + fg->last_use = fc.counter; - ggml_gallocr_free(alloc); - ggml_free(ctx); + // ── Fill inputs ───────────────────────────────────────────────── + const auto set_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_set(fg->inp_embed, embed, 0, sizeof(float) * (size_t) w.n_embd); - const int next_pos = kv_start + n_tokens; - for (int il = 0; il < n_layer; ++il) { - const uint32_t ratio = w.compress_ratios[il]; - if (ratio <= 0 || (next_pos % (int) ratio) != 0) { - continue; + std::vector i32v((size_t) w.n_layer * 6, 0); + std::vector i64v((size_t) w.n_layer * 5, 0); + const int64_t raw_row = kv_start % w.n_swa; + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + i32v[(size_t) il * 6 + 0] = kv_start; + i32v[(size_t) il * 6 + 1] = -kv_start; + i64v[(size_t) il * 5 + 0] = raw_row; + if (ratio > 0) { + const int pos_mod = token_pos % ratio; + i32v[(size_t) il * 6 + 2] = pos_mod; + i32v[(size_t) il * 6 + 3] = token_pos + 1 - ratio; + i64v[(size_t) il * 5 + 1] = (ratio == 4) ? (int64_t) (ratio + pos_mod) : (int64_t) pos_mod; + i64v[(size_t) il * 5 + 2] = token_pos / ratio; } - cache.layers[il].n_comp = std::max(cache.layers[il].n_comp, next_pos / (int) ratio); if (ratio == 4) { - cache.layers[il].n_index_comp = std::max(cache.layers[il].n_index_comp, - next_pos / (int) ratio); + const int pos_mod = token_pos % ratio; + i32v[(size_t) il * 6 + 4] = pos_mod; + i32v[(size_t) il * 6 + 5] = token_pos + 1 - ratio; + i64v[(size_t) il * 5 + 3] = ratio + pos_mod; + i64v[(size_t) il * 5 + 4] = token_pos / ratio; + } + } + ggml_backend_tensor_set(fg->i32_bundle, i32v.data(), 0, sizeof(int32_t) * i32v.size()); + ggml_backend_tensor_set(fg->i64_bundle, i64v.data(), 0, sizeof(int64_t) * i64v.size()); + + if (fg->mask_bundle) { + std::vector maskv((size_t) ggml_nelements(fg->mask_bundle), 0.0f); + size_t off = 0; + const int n_valid_raw = std::min(kv_start + 1, w.n_swa); + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + int n_comp = 0, padded = 0; + if (ratio > 0 && lc.comp_kv) { + n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + for (int j = n_valid_raw; j < w.n_swa; ++j) { + maskv[off + (size_t) j] = -1.0e30f; + } + for (int j = n_comp; j < padded; ++j) { + maskv[off + (size_t) w.n_swa + (size_t) j] = -1.0e30f; + } + off += (size_t) w.n_swa + (size_t) padded; + } + if (off != maskv.size()) { + std::fprintf(stderr, "[deepseek4] fused decode: mask layout mismatch (%zu vs %zu)\n", + off, maskv.size()); + return -1; } + ggml_backend_tensor_set(fg->mask_bundle, maskv.data(), 0, sizeof(float) * maskv.size()); } - cache.cur_pos = next_pos; - return true; + if (token_ids) { + for (int il = 0; il < w.n_layer; ++il) { + ggml_tensor * hids = fg->hash_ids[(size_t) il]; + if (!hids) continue; + const int n_used = ds4_effective_expert_count(w); + hash_scratch.resize((size_t) n_used); + const int32_t tok = token_ids[0]; + const int32_t * row = hash_routing_row( + hash_tables[(size_t) il], tok, w.n_expert_used); + if (!row) { + std::fprintf(stderr, + "[deepseek4] token id %d outside hash table for layer %d\n", + tok, il); + return -1; + } + std::memcpy(hash_scratch.data(), row, + (size_t) n_used * sizeof(int32_t)); + ggml_backend_tensor_set(hids, hash_scratch.data(), 0, + sizeof(int32_t) * (size_t) n_used); + } + } + if (telemetry) telemetry->full_graph_set_us += ds4_elapsed_us(set_t0, Ds4TimingClock::now()); + + // ── Compute ───────────────────────────────────────────────────── + const auto compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, fg->sg.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[deepseek4] fused decode graph compute failed\n"); + return -1; + } + if (telemetry) telemetry->full_graph_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); + + // ── Read logits ───────────────────────────────────────────────── + const auto read_t0 = Ds4TimingClock::now(); + out_logits.resize((size_t) w.n_vocab); + ggml_backend_tensor_get(fg->logits, out_logits.data(), 0, + sizeof(float) * (size_t) w.n_vocab); + if (telemetry) telemetry->full_graph_read_us += ds4_elapsed_us(read_t0, Ds4TimingClock::now()); + return 1; } + bool deepseek4_step_layer_range( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -2819,9 +3890,6 @@ bool deepseek4_step_layer_range( DeepSeek4StepTelemetry * telemetry) { const auto step_t0 = Ds4TimingClock::now(); - // NOTE: The old deepseek4_step() lacks HC implementation. - // Always use the HC-enabled layer_range path below. - // ── Partial layer-range forward with HC ───────────────────────────── const int n_embd = w.n_embd; const int n_hc = w.n_hc; @@ -2860,10 +3928,15 @@ bool deepseek4_step_layer_range( static std::vector cached_decode_ffn_graphs; static DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; static DeepSeek4CachedLayerAlloc cached_dynamic_output_alloc; + static DeepSeek4FusedDecodeCache fused_decode_graph_cache; + static Ds4DecodeSharedInputs decode_shared_inputs; static int hc_loaded_n_layer = 0; - if (hc_loaded_n_layer != w.n_layer) { + static const ggml_context * hc_loaded_ctx = nullptr; + if (hc_loaded_n_layer != w.n_layer || hc_loaded_ctx != w.ctx) { + reset_hc_layer_weights_cpu(hc_layer_weights_range); + reset_hc_weights_cpu(hc_output_weights_range); hc_layer_weights_range.resize((size_t)w.n_layer); - hash_routing_tables_range.resize((size_t)w.n_layer); + hash_routing_tables_range.assign((size_t)w.n_layer, {}); for (auto & alloc : cached_attn_allocs) { alloc.free(); } @@ -2889,6 +3962,8 @@ bool deepseek4_step_layer_range( cached_decode_ffn_graphs.assign((size_t)w.n_layer, {}); cached_decode_output_graph.free(); cached_dynamic_output_alloc.free(); + fused_decode_graph_cache.destroy(); + decode_shared_inputs.free(); for (int il = 0; il < w.n_layer; il++) { const DeepSeek4Layer & L = w.layers[(size_t)il]; load_hc_weights_cpu(hc_layer_weights_range[il].attn, L.hc_attn_fn, L.hc_attn_scale, L.hc_attn_base); @@ -2899,11 +3974,34 @@ bool deepseek4_step_layer_range( } load_hc_weights_cpu(hc_output_weights_range, w.output_hc_fn, w.output_hc_scale, w.output_hc_base); hc_loaded_n_layer = w.n_layer; + hc_loaded_ctx = w.ctx; } // Per-layer execution with CPU-side HC static thread_local DeepSeek4LayerRangeScratch scratch; - scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, w.n_expert_used); + const int n_expert_used = ds4_effective_expert_count(w); + scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, n_expert_used); + + if (n_tokens == 1 && layer_begin == 0 && is_last_shard && + out_logits && ds4_backend_is_gpu(backend) && w.fused_decode) { + const int rc = ds4_try_fused_decode_step( + fused_decode_graph_cache, backend, w, cache, hc_layer_weights_range, + hc_output_weights_range, hash_routing_tables_range, scratch.hash_expert_ids, + embed, kv_start, *out_logits, token_ids, telemetry); + if (rc < 0) return false; + if (rc > 0) { + const int np = kv_start + 1; + for (int il = layer_begin; il < layer_end; ++il) { + const uint32_t ratio = w.compress_ratios[il]; + if (ratio <= 0 || (np % (int) ratio) != 0) continue; + cache.layers[il].n_comp = std::max(cache.layers[il].n_comp, np / (int) ratio); + if (ratio == 4) cache.layers[il].n_index_comp = std::max(cache.layers[il].n_index_comp, np / (int) ratio); + } + cache.cur_pos = np; + if (telemetry) telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + return true; + } + } std::vector & cur = scratch.cur; std::vector & ffn_working = scratch.ffn_working; std::vector & hc_post = scratch.hc_post; @@ -2912,12 +4010,32 @@ bool deepseek4_step_layer_range( std::vector & attn_out_host = scratch.attn_out_host; std::vector & ffn_out_host = scratch.ffn_out_host; std::vector & hash_expert_ids_host = scratch.hash_expert_ids; - const bool reuse_decode_graphs = (n_tokens == 1); - const bool use_backend_decode_hc = reuse_decode_graphs && ds4_backend_is_gpu(backend); + const bool reuse_decode_graphs = n_tokens == 1; + Ds4DecodeSharedInputs * shared_inputs = nullptr; + if (reuse_decode_graphs && ds4_backend_is_gpu(backend) && + decode_shared_inputs.ensure(w, backend)) { + decode_shared_inputs.set_step(w, kv_start); + shared_inputs = &decode_shared_inputs; + } + + bool backend_decode_hc_supported = true; + for (int il = layer_begin; il < layer_end; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t)il]; + const HcLayerWeightsCpu & hc_lw = hc_layer_weights_range[(size_t)il]; + if (!hc_fn_device_ptr(hc_lw.attn, L.hc_attn_fn) || + !hc_fn_device_ptr(hc_lw.ffn, L.hc_ffn_fn)) { + backend_decode_hc_supported = false; + break; + } + } + const bool use_backend_decode_hc = + reuse_decode_graphs && + ds4_backend_is_gpu(backend) && + backend_decode_hc_supported; const bool use_backend_decode_hc_direct = use_backend_decode_hc && ds4_backend_is_hip(backend); const bool use_backend_decode_hc_graph = use_backend_decode_hc && !use_backend_decode_hc_direct; - const ggml_tensor * hc_state_backend = nullptr; + ggml_tensor * hc_state_backend = nullptr; if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { if (!cached_decode_hc_post_graph.valid() || cached_decode_hc_post_graph.owner_ctx != w.ctx || @@ -2946,10 +4064,6 @@ bool deepseek4_step_layer_range( // ── HC pre (attention) ────────────────────────────────────── const auto hc_pre_attn_t0 = Ds4TimingClock::now(); if (use_backend_decode_hc_direct) { - const auto hc_pre_attn_input_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_set(cached_decode_hc_post_graph.residual_hc, - hc_state.data(), 0, sizeof(float) * hc_state.size()); - if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_attn_input_t0, Ds4TimingClock::now()); auto & cached = cached_decode_attn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || cached.owner_ctx != w.ctx || @@ -2970,8 +4084,9 @@ bool deepseek4_step_layer_range( backend, il, false, - cached_decode_hc_post_graph.residual_hc, + hc_state_backend, L.hc_attn_fn, + hc_fn_device_ptr(hc_lw.attn, L.hc_attn_fn), L.hc_attn_scale, L.hc_attn_base, hc_lw.attn.scale_data.data(), @@ -3023,7 +4138,7 @@ bool deepseek4_step_layer_range( // ── Build & run attention graph ───────────────────────────── { const int token_pos = kv_start + n_tokens - 1; - const bool reuse_decode_attn = n_tokens == 1; + const bool reuse_decode_attn = reuse_decode_graphs; ggml_tensor * attn_out = nullptr; ggml_cgraph * gf = nullptr; ggml_context * ctx = nullptr; @@ -3033,6 +4148,8 @@ bool deepseek4_step_layer_range( const int n_raw = std::min(kv_start + 1, w.n_swa); const int n_comp_attn = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; const int n_index_comp = (ratio == 4) ? ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos) : 0; + const bool attn_flush = ratio > 0 && (((token_pos + 1) % ratio) == 0); + const bool index_flush = ratio == 4 && (((token_pos + 1) % ratio) == 0); auto & per_layer = cached_decode_attn_graphs[(size_t)il]; auto it = std::find_if(per_layer.begin(), per_layer.end(), [&](const DeepSeek4CachedDecodeAttnGraph & candidate) { @@ -3042,7 +4159,9 @@ bool deepseek4_step_layer_range( candidate.layer_idx == il && candidate.n_raw == n_raw && candidate.n_comp_attn == n_comp_attn && - candidate.n_index_comp == n_index_comp; + candidate.n_index_comp == n_index_comp && + candidate.attn_flush == attn_flush && + candidate.index_flush == index_flush; }); if (it == per_layer.end()) { if (per_layer.size() >= 20) { @@ -3053,7 +4172,8 @@ bool deepseek4_step_layer_range( auto & candidate = per_layer.back(); const auto attn_build_t0 = Ds4TimingClock::now(); if (!build_cached_decode_attn_graph(candidate, backend, w, L, lc, il, kv_start, - n_raw, n_comp_attn, n_index_comp)) { + n_raw, n_comp_attn, n_index_comp, + shared_inputs)) { std::fprintf(stderr, "[deepseek4] cached attn graph alloc failed layer %d\n", il); return false; } @@ -3072,13 +4192,14 @@ bool deepseek4_step_layer_range( } else { ggml_backend_tensor_set(cached_attn->sg.inp_embed, cur.data(), 0, sizeof(float) * cur.size()); } + if (!cached_attn->uses_shared_inputs) { ggml_backend_tensor_set(cached_attn->inputs.rope_pos, &rope_pos, 0, sizeof(rope_pos)); ggml_backend_tensor_set(cached_attn->inputs.neg_pos, &neg_pos, 0, sizeof(neg_pos)); ggml_backend_tensor_set(cached_attn->inputs.raw_kv_rows, &raw_row, 0, sizeof(raw_row)); if (ratio > 0) { const int pos_mod = token_pos % ratio; const int32_t ape_row = pos_mod; - const int64_t state_row = ratio + pos_mod; + const int64_t state_row = (ratio == 4) ? (ratio + pos_mod) : pos_mod; const int64_t comp_row = token_pos / ratio; const int32_t comp_pos = token_pos + 1 - ratio; const bool flush_boundary = ((token_pos + 1) % ratio) == 0; @@ -3103,9 +4224,10 @@ bool deepseek4_step_layer_range( ggml_backend_tensor_set(cached_attn->inputs.index_comp_pos, &comp_pos, 0, sizeof(comp_pos)); } } + } } else { const auto attn_build_t0 = Ds4TimingClock::now(); - const size_t ctx_size = 48 * 1024 * 1024; + const size_t ctx_size = ds4_attn_step_meta_size(n_tokens); ggml_init_params params{}; params.mem_size = ctx_size; params.mem_buffer = nullptr; @@ -3118,7 +4240,7 @@ bool deepseek4_step_layer_range( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; - const size_t graph_size = n_tokens > 1 ? 32768 : 2048; + const size_t graph_size = ds4_attn_step_graph_size(n_tokens); gf = ggml_new_graph_custom(ctx, graph_size, false); ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps); @@ -3162,7 +4284,7 @@ bool deepseek4_step_layer_range( return false; } if (telemetry) telemetry->attn_compute_us += ds4_elapsed_us(attn_compute_t0, Ds4TimingClock::now()); - if (use_backend_decode_hc_graph) { + if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { if (hc_state_backend != cached_decode_hc_post_graph.residual_hc) { ggml_backend_tensor_copy(hc_state_backend, cached_decode_hc_post_graph.residual_hc); } @@ -3180,16 +4302,12 @@ bool deepseek4_step_layer_range( } else { const auto attn_read_t0 = Ds4TimingClock::now(); ggml_backend_tensor_get(attn_out, attn_out_host.data(), 0, sizeof(float) * attn_out_host.size()); - if (use_backend_decode_hc_direct) { - ggml_backend_tensor_get(attn_post_backend, hc_post.data(), 0, sizeof(float) * hc_post.size()); - ggml_backend_tensor_get(attn_comb_backend, hc_comb.data(), 0, sizeof(float) * hc_comb.size()); - } if (telemetry) telemetry->attn_read_us += ds4_elapsed_us(attn_read_t0, Ds4TimingClock::now()); } if (ctx) ggml_free(ctx); // ── HC post (attention) ───────────────────────────────── - if (!use_backend_decode_hc_graph) { + if (!(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { const auto hc_post_attn_t0 = Ds4TimingClock::now(); hc_post_batch(next_hc, attn_out_host.data(), @@ -3207,10 +4325,6 @@ bool deepseek4_step_layer_range( // ── HC pre (FFN) ──────────────────────────────────────────── const auto hc_pre_ffn_t0 = Ds4TimingClock::now(); if (use_backend_decode_hc_direct) { - const auto hc_pre_ffn_input_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_set(cached_decode_hc_post_graph.residual_hc, - hc_state.data(), 0, sizeof(float) * hc_state.size()); - if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_ffn_input_t0, Ds4TimingClock::now()); auto & cached = cached_decode_ffn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || cached.owner_ctx != w.ctx || @@ -3231,8 +4345,9 @@ bool deepseek4_step_layer_range( backend, il, true, - cached_decode_hc_post_graph.residual_hc, + hc_state_backend, L.hc_ffn_fn, + hc_fn_device_ptr(hc_lw.ffn, L.hc_ffn_fn), L.hc_ffn_scale, L.hc_ffn_base, hc_lw.ffn.scale_data.data(), @@ -3289,12 +4404,18 @@ bool deepseek4_step_layer_range( il < w.n_hash_layer && L.ffn_gate_tid2eid && token_ids && hash_routing_tables_range[(size_t)il].loaded; if (hash_routed) { - const auto & ht = hash_routing_tables_range[(size_t)il].ids; - const int n_used = w.n_expert_used; + const int n_used = n_expert_used; hash_expert_ids_host.resize((size_t)n_used * (size_t)n_tokens); for (int ti = 0; ti < n_tokens; ti++) { const int32_t tok = token_ids[ti]; - const int32_t * row = ht.data() + (size_t)tok * (size_t)n_used; + const int32_t * row = hash_routing_row( + hash_routing_tables_range[(size_t)il], tok, w.n_expert_used); + if (!row) { + std::fprintf(stderr, + "[deepseek4] token id %d outside hash table for layer %d\n", + tok, il); + return false; + } memcpy(hash_expert_ids_host.data() + (size_t)ti * n_used, row, (size_t)n_used * sizeof(int32_t)); } @@ -3306,6 +4427,7 @@ bool deepseek4_step_layer_range( cached.backend != backend || cached.layer_idx != il || cached.n_tokens != n_tokens || + cached.n_expert_used != n_expert_used || cached.hash_routed != hash_routed) { const auto ffn_build_t0 = Ds4TimingClock::now(); if (!build_cached_decode_ffn_graph(cached, backend, w, L, il, n_tokens, hash_routed)) { @@ -3335,7 +4457,7 @@ bool deepseek4_step_layer_range( return false; } - if (use_backend_decode_hc_graph) { + if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { if (hc_state_backend != cached_decode_hc_post_graph.residual_hc) { ggml_backend_tensor_copy(hc_state_backend, cached_decode_hc_post_graph.residual_hc); } @@ -3352,15 +4474,11 @@ bool deepseek4_step_layer_range( } else { const auto ffn_read_t0 = Ds4TimingClock::now(); ggml_backend_tensor_get(ffn_out, ffn_out_host.data(), 0, sizeof(float) * ffn_out_host.size()); - if (use_backend_decode_hc_direct) { - ggml_backend_tensor_get(ffn_post_backend, hc_post.data(), 0, sizeof(float) * hc_post.size()); - ggml_backend_tensor_get(ffn_comb_backend, hc_comb.data(), 0, sizeof(float) * hc_comb.size()); - } if (telemetry) telemetry->ffn_read_us += ds4_elapsed_us(ffn_read_t0, Ds4TimingClock::now()); } // ── HC post (FFN) ─────────────────────────────────────── - if (!use_backend_decode_hc_graph) { + if (!(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { const auto hc_post_ffn_t0 = Ds4TimingClock::now(); hc_post_batch(next_hc, ffn_out_host.data(), @@ -3376,7 +4494,7 @@ bool deepseek4_step_layer_range( } } - if (use_backend_decode_hc_graph && hc_state_backend) { + if ((use_backend_decode_hc_graph || use_backend_decode_hc_direct) && hc_state_backend) { ggml_backend_tensor_get(hc_state_backend, hc_state.data(), 0, sizeof(float) * hc_state.size()); } @@ -3564,6 +4682,24 @@ bool create_deepseek4_cache(ggml_backend_t backend, return true; } +void free_deepseek4_cache(DeepSeek4Cache & c) { + if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } + if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } + c.layers.clear(); + c.hc_state = nullptr; +} + +void reset_deepseek4_cache(DeepSeek4Cache & c) { + c.cur_pos = 0; + for (DeepSeek4LayerCache & lc : c.layers) { + lc.n_comp = 0; + lc.n_index_comp = 0; + } + if (c.buf) { + ggml_backend_buffer_clear(c.buf, 0); + } +} + namespace { ggml_tensor * clone_snapshot_tensor(ggml_context * ctx, @@ -3744,16 +4880,10 @@ bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, return true; } -void free_deepseek4_cache(DeepSeek4Cache & c) { - if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } - if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } - c.layers.clear(); - c.hc_state = nullptr; -} void free_deepseek4_snapshot(DeepSeek4Snapshot & s) { - if (s.buf) { ggml_backend_buffer_free(s.buf); s.buf = nullptr; } if (s.ctx) { ggml_free(s.ctx); s.ctx = nullptr; } + if (s.buf) { ggml_backend_buffer_free(s.buf); s.buf = nullptr; } s.layers.clear(); s.cur_pos = 0; s.hc_state_snap = nullptr; diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 31c086394..27fb5a1f1 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -64,6 +64,10 @@ struct DeepSeek4StepTelemetry { uint64_t output_us = 0; uint64_t sample_us = 0; uint64_t emit_us = 0; + uint64_t full_graph_build_us = 0; + uint64_t full_graph_set_us = 0; + uint64_t full_graph_compute_us = 0; + uint64_t full_graph_read_us = 0; int hot_selected = 0; int cold_selected = 0; }; @@ -216,6 +220,11 @@ struct DeepSeek4Weights { // MoE hybrid placement (deprecated — layer split replaces expert split) bool moe_hybrid = false; + + // Runtime serving policy. These values are set by the backend after the + // GGUF is loaded; they are not model metadata. + int routed_expert_top_k = 0; // 0 = model default (n_expert_used) + bool fused_decode = false; }; inline bool deepseek4_is_eos_tok(int tok, const DeepSeek4Weights & w) { @@ -265,6 +274,11 @@ struct DeepSeek4Cache { struct DeepSeek4Snapshot; +struct DeepSeek4RawRingSpan { + int row = 0; + int count = 0; +}; + // ─── Configuration ────────────────────────────────────────────────────── struct DeepSeek4BackendConfig { @@ -273,6 +287,8 @@ struct DeepSeek4BackendConfig { int stream_fd = -1; int chunk = 512; // prefill chunk size int max_ctx = 0; // 0 = auto from SWA + compression capacity + int expert_top_k = 0; // 0 = use all model-routed experts + bool fused_decode = false; // single-graph GPU decode }; // ─── Function declarations ────────────────────────────────────────────── @@ -294,6 +310,11 @@ bool create_deepseek4_cache(ggml_backend_t backend, DeepSeek4Cache & out); void free_deepseek4_cache(DeepSeek4Cache & c); +void reset_deepseek4_cache(DeepSeek4Cache & c); +int deepseek4_previous_raw_ring_spans( + int kv_start, + int n_swa, + DeepSeek4RawRingSpan spans[2]); bool deepseek4_snapshot_save(const DeepSeek4Cache & cache, ggml_backend_t snapshot_backend, DeepSeek4Snapshot & out); diff --git a/server/src/rms_norm_hip.cu b/server/src/rms_norm_hip.cu index 51091619e..3d612eb58 100644 --- a/server/src/rms_norm_hip.cu +++ b/server/src/rms_norm_hip.cu @@ -4,31 +4,16 @@ #include -// Wavefront width of the target arch: 32 on RDNA (gfx10/gfx11), 64 on GCN/CDNA. -// The reduction below must span exactly one wavefront, so hardcoding 32 leaves -// lanes 32..63 out of the reduce on wave64 archs. __AMDGCN_WAVEFRONT_SIZE(__) is -// a compile-time builtin defined only in the device pass. The 64 fallback is -// host-only (WAVE is never used off-device); a device pass without the builtin -// is a hard error rather than a silent wrong-width guess. -#ifndef DFLASH_WAVE_SIZE -# if defined(__AMDGCN_WAVEFRONT_SIZE__) -# define DFLASH_WAVE_SIZE __AMDGCN_WAVEFRONT_SIZE__ -# elif defined(__AMDGCN_WAVEFRONT_SIZE) -# define DFLASH_WAVE_SIZE __AMDGCN_WAVEFRONT_SIZE -# elif !defined(__HIP_DEVICE_COMPILE__) -# define DFLASH_WAVE_SIZE 64 // host pass only; WAVE is unused off-device -# else -# error "rms_norm_hip: missing wavefront-size builtin in device pass; set DFLASH_WAVE_SIZE" -# endif -#endif - __global__ void rms_norm_mul_w_f32_kernel( const float * __restrict__ src, const float * __restrict__ w, float * __restrict__ dst, int hidden, float eps) { - constexpr int WAVE = DFLASH_WAVE_SIZE; + // HIP exposes the target-specific width through a device builtin. Unlike + // preprocessor spellings of the AMDGCN builtin, this works across ROCm + // releases and remains correct in multi-architecture builds. + const int wave = warpSize; const int tok = blockIdx.x; const float * row = src + (size_t)tok * hidden; @@ -43,21 +28,21 @@ __global__ void rms_norm_mul_w_f32_kernel( } #pragma unroll - for (int off = WAVE / 2; off > 0; off >>= 1) + for (int off = wave / 2; off > 0; off >>= 1) sumsq += __shfl_xor(sumsq, off); - if ((threadIdx.x & (WAVE - 1)) == 0) - smem[threadIdx.x / WAVE] = sumsq; + if ((threadIdx.x & (wave - 1)) == 0) + smem[threadIdx.x / wave] = sumsq; __syncthreads(); // Final reduce across per-wavefront partials in a single wavefront. Valid - // while n_warps <= WAVE, which holds for the fixed block=256 launch below + // while n_warps <= wave, which holds for the fixed block=256 launch below // (8 warps on wave32, 4 on wave64). - const int n_warps = blockDim.x / WAVE; - if (threadIdx.x < WAVE) { + const int n_warps = blockDim.x / wave; + if (threadIdx.x < wave) { sumsq = (threadIdx.x < n_warps) ? smem[threadIdx.x] : 0.0f; #pragma unroll - for (int off = WAVE / 2; off > 0; off >>= 1) + for (int off = wave / 2; off > 0; off >>= 1) sumsq += __shfl_xor(sumsq, off); if (threadIdx.x == 0) smem[0] = sumsq; diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index fe0ba31f5..9c9c30f76 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -217,6 +217,10 @@ static void print_usage(const char * prog) { " --target-layer-split Reserved layer-split weights\n" " --peer-access Enable peer access for multi-GPU placement\n" " --chunk Chunked-prefill chunk size (default: 512)\n" + " --ds4-fused-decode Enable DeepSeek4 single-graph GPU decode\n" + " --ds4-expert-top-k \n" + " Keep and renormalize the highest-ranked N routed experts\n" + " (0=model default; single-device DeepSeek4 only)\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" @@ -423,6 +427,14 @@ int main(int argc, char ** argv) { bargs.device.peer_access = true; } else if (std::strcmp(argv[i], "--chunk") == 0 && i + 1 < argc) { bargs.chunk = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--ds4-fused-decode") == 0) { + bargs.ds4_fused_decode = true; + } else if (std::strcmp(argv[i], "--ds4-expert-top-k") == 0 && i + 1 < argc) { + bargs.ds4_expert_top_k = std::atoi(argv[++i]); + if (bargs.ds4_expert_top_k < 0) { + std::fprintf(stderr, "[server] --ds4-expert-top-k must be non-negative\n"); + return 2; + } } else if (std::strcmp(argv[i], "--fa-window") == 0 && i + 1 < argc) { bargs.fa_window = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--model-name") == 0 && i + 1 < argc) { @@ -786,6 +798,22 @@ int main(int argc, char ** argv) { g_peer_access_opt_in = bargs.device.peer_access; std::fprintf(stderr, "[server] creating backend...\n"); const std::string arch = detect_arch(bargs.model_path); + const PlacementBackend target_backend = + bargs.device.backend == PlacementBackend::Auto + ? compiled_placement_backend() + : bargs.device.backend; + const bool monolithic_ds4 = + arch == "deepseek4" && + target_backend == PlacementBackend::Hip && + !bargs.device.is_layer_split() && + !bargs.remote_target_shard.enabled(); + if ((bargs.ds4_fused_decode || bargs.ds4_expert_top_k != 0) && + !monolithic_ds4) { + std::fprintf(stderr, + "[server] --ds4-fused-decode and --ds4-expert-top-k " + "currently require single-device HIP DeepSeek4\n"); + return 2; + } if (spark_autotune) { // Self-tuning hot/cold MoE residency: enable the bounded expert cache // (auto-tunes the working set at serve time), auto-load a learned @@ -1072,6 +1100,16 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[server] │ peer_access = %s\n", bargs.device.peer_access ? "ON" : "off"); std::fprintf(stderr, "[server] │ chunk = %d\n", bargs.chunk); + if (arch == "deepseek4") { + std::fprintf(stderr, "[server] │ ds4_fused = %s\n", + bargs.ds4_fused_decode ? "ON" : "off"); + if (bargs.ds4_expert_top_k > 0) { + std::fprintf(stderr, "[server] │ ds4_expert_topk= %d\n", + bargs.ds4_expert_top_k); + } else { + std::fprintf(stderr, "[server] │ ds4_expert_topk= model default\n"); + } + } std::fprintf(stderr, "[server] │ fa_window = %d\n", bargs.fa_window); if (bargs.fa_window > 0) { std::fprintf(stderr, "[server] │ ⚠ fa_window > 0 drops system prompt / " diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index f36b07d33..7f4bd6084 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -545,6 +545,31 @@ static void test_hash_routing_lookup() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_raw_ring_spans_after_wrap() { + std::fprintf(stderr, " test_raw_ring_spans_after_wrap ..."); + + DeepSeek4RawRingSpan spans[2]; + int count = deepseek4_previous_raw_ring_spans(3, 8, spans); + TEST_ASSERT(count == 1); + TEST_ASSERT(spans[0].row == 0); + TEST_ASSERT(spans[0].count == 3); + + count = deepseek4_previous_raw_ring_spans(8, 8, spans); + TEST_ASSERT(count == 1); + TEST_ASSERT(spans[0].row == 1); + TEST_ASSERT(spans[0].count == 7); + + count = deepseek4_previous_raw_ring_spans(10, 8, spans); + TEST_ASSERT(count == 2); + TEST_ASSERT(spans[0].row == 3); + TEST_ASSERT(spans[0].count == 5); + TEST_ASSERT(spans[1].row == 0); + TEST_ASSERT(spans[1].count == 2); + TEST_ASSERT(spans[0].count + spans[1].count == 7); + + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + struct ScopedEnvVar { explicit ScopedEnvVar(const char * name) : name(name ? name : ""), @@ -908,6 +933,43 @@ static void test_reset_request_state() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_reset_deepseek4_cache(ggml_backend_t backend) { + std::fprintf(stderr, " test_reset_deepseek4_cache ..."); + + DeepSeek4Weights weights; + weights.n_layer = 2; + weights.n_swa = 8; + weights.head_dim = 4; + weights.n_hc = 4; + weights.n_embd = 8; + weights.compress_ratios = {4, 0}; + + DeepSeek4Cache cache; + TEST_ASSERT(create_deepseek4_cache(backend, weights, 32, cache)); + if (cache.buf) { + ggml_backend_buffer_clear(cache.buf, 0x7f); + cache.cur_pos = 17; + cache.layers[0].n_comp = 4; + cache.layers[0].n_index_comp = 4; + cache.layers[1].n_comp = 3; + cache.layers[1].n_index_comp = 2; + + reset_deepseek4_cache(cache); + + TEST_ASSERT(cache.cur_pos == 0); + for (const auto & layer : cache.layers) { + TEST_ASSERT(layer.n_comp == 0); + TEST_ASSERT(layer.n_index_comp == 0); + const std::vector bytes = read_tensor_bytes(layer.raw_kv); + TEST_ASSERT(std::all_of(bytes.begin(), bytes.end(), + [](uint8_t value) { return value == 0; })); + } + } + free_deepseek4_cache(cache); + + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_adapter_guard_paths() { std::fprintf(stderr, " test_adapter_guard_paths ..."); @@ -1609,6 +1671,7 @@ int main() { test_rmsnorm_correctness(backend); test_grouped_output_projection_shape(); test_hash_routing_lookup(); + test_raw_ring_spans_after_wrap(); test_auto_split_computation(); test_layer_range_validation(); test_hc_state_dimensions(); @@ -1619,6 +1682,7 @@ int main() { test_loader_rejects_truncated_tensor_data(backend); test_snapshot_save_restore(); test_reset_request_state(); + test_reset_deepseek4_cache(backend); test_adapter_guard_paths(); test_ipc_mode_registration(); test_target_shard_daemon_validation();