diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 78215b4d3..3163196b1 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -9,27 +9,69 @@ #if defined(GGML_USE_HIP) +__device__ __forceinline__ float ds4_warp_reduce_max(float v) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { +#if defined(__HIP_PLATFORM_AMD__) + v = fmaxf(v, __shfl_xor(v, mask, 32)); +#else + v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, mask, 32)); +#endif + } + return v; +} + +__device__ __forceinline__ float ds4_warp_reduce_sum(float v) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { +#if defined(__HIP_PLATFORM_AMD__) + v += __shfl_xor(v, mask, 32); +#else + v += __shfl_xor_sync(0xffffffffu, v, mask, 32); +#endif + } + return v; +} + __device__ static float ds4_fa_block_sum(float v) { - __shared__ float smem[256]; + v = ds4_warp_reduce_sum(v); + __shared__ float smem[8]; const int tid = threadIdx.x; - smem[tid] = v; + const int lane = tid & 31; + const int warp_id = tid >> 5; + if (lane == 0) { + smem[warp_id] = v; + } __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) smem[tid] += smem[tid + stride]; - __syncthreads(); + if (warp_id == 0) { + float val = (lane < 8) ? smem[lane] : 0.0f; + val = ds4_warp_reduce_sum(val); + if (lane == 0) { + smem[0] = val; + } } + __syncthreads(); return smem[0]; } __device__ static float ds4_fa_block_max(float v) { - __shared__ float smem[256]; + v = ds4_warp_reduce_max(v); + __shared__ float smem[8]; const int tid = threadIdx.x; - smem[tid] = v; + const int lane = tid & 31; + const int warp_id = tid >> 5; + if (lane == 0) { + smem[warp_id] = v; + } __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) smem[tid] = fmaxf(smem[tid], smem[tid + stride]); - __syncthreads(); + if (warp_id == 0) { + float val = (lane < 8) ? smem[lane] : -3.402823466e38f; + val = ds4_warp_reduce_max(val); + if (lane == 0) { + smem[0] = val; + } } + __syncthreads(); return smem[0]; } @@ -642,8 +684,12 @@ __global__ static void ds4_flash_attn_d512_shared_kv_kernel( if (mask_v > -1.0e20f) { dot = 0.0f; #pragma unroll - for (int d = 0; d < D; ++d) { - const float qv = inverse_rope.forward_q_enabled && d >= D - 64 + for (int d = 0; d < D - 64; ++d) { + dot += qh[d] * kb[d]; + } +#pragma unroll + for (int d = D - 64; d < D; ++d) { + const float qv = inverse_rope.forward_q_enabled ? q_rope_tail[d - (D - 64)] : qh[d]; dot += qv * kb[d]; } @@ -687,8 +733,12 @@ __global__ static void ds4_flash_attn_d512_shared_kv_kernel( const KV * kr = k + (size_t) r * D; float dot = 0.0f; #pragma unroll - for (int d = 0; d < D; ++d) { - const float qv = inverse_rope.forward_q_enabled && d >= D - 64 + for (int d = 0; d < D - 64; ++d) { + dot += qh[d] * ds4_fa_load(kr + d); + } +#pragma unroll + for (int d = D - 64; d < D; ++d) { + const float qv = inverse_rope.forward_q_enabled ? q_rope_tail[d - (D - 64)] : qh[d]; dot += qv * ds4_fa_load(kr + d); } @@ -884,14 +934,21 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_kernel( if (visible) { const KV * kr = k + (size_t) r * D; #pragma unroll - for (int d = 0; d < D; ++d) { + for (int d = 0; d < D - 64; ++d) { + const float kv = ds4_fa_load(kr + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + dot[j] += qh[j][d] * kv; + } + } +#pragma unroll + for (int d = D - 64; d < D; ++d) { const float kv = ds4_fa_load(kr + d); #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - const float qv = - inverse_rope.forward_q_enabled && d >= D - 64 - ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] - : qh[j][d]; + const float qv = inverse_rope.forward_q_enabled + ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] + : qh[j][d]; dot[j] += qv * kv; } } @@ -906,29 +963,35 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_kernel( } } - // Match ds4_fa_block_max independently for every grouped head. + // Match ds4_fa_block_max independently for every grouped head via warp shuffles. + const int lane = tid & 31; + const int warp_id = tid >> 5; #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - reduction[(size_t) j * N_THREADS + tid] = local_max[j]; + const float v = ds4_warp_reduce_max(local_max[j]); + if (lane == 0) { + reduction[(size_t) j * 8 + warp_id] = v; + } } __syncthreads(); - for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { - if (tid < stride) { + float max_score[HEADS_PER_BLOCK]; + if (warp_id == 0) { #pragma unroll - for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - float * row = reduction + (size_t) j * N_THREADS; - row[tid] = fmaxf(row[tid], row[tid + stride]); + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float v = (lane < 8) ? reduction[(size_t) j * 8 + lane] : -3.402823466e38f; + v = ds4_warp_reduce_max(v); + if (lane == 0) { + reduction[j] = v; } } - __syncthreads(); } - - float max_score[HEADS_PER_BLOCK]; - float local_sum[HEADS_PER_BLOCK] = {}; + __syncthreads(); #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - max_score[j] = reduction[(size_t) j * N_THREADS]; + max_score[j] = reduction[j]; } + float local_sum[HEADS_PER_BLOCK] = {}; + for (int r = tid; r < n_kv; r += blockDim.x) { #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { @@ -945,27 +1008,30 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_kernel( } } - // Match ds4_fa_block_sum independently for every grouped head. + // Match ds4_fa_block_sum independently for every grouped head via warp shuffles. #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - reduction[(size_t) j * N_THREADS + tid] = local_sum[j]; + const float v = ds4_warp_reduce_sum(local_sum[j]); + if (lane == 0) { + reduction[(size_t) j * 8 + warp_id] = v; + } } __syncthreads(); - for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { - if (tid < stride) { + float inv_denom[HEADS_PER_BLOCK]; + if (warp_id == 0) { #pragma unroll - for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - float * row = reduction + (size_t) j * N_THREADS; - row[tid] += row[tid + stride]; + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float v = (lane < 8) ? reduction[(size_t) j * 8 + lane] : 0.0f; + v = ds4_warp_reduce_sum(v); + if (lane == 0) { + reduction[j] = 1.0f / v; } } - __syncthreads(); } - - float inv_denom[HEADS_PER_BLOCK]; + __syncthreads(); #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - inv_denom[j] = 1.0f / reduction[(size_t) j * N_THREADS]; + inv_denom[j] = reduction[j]; } // Underflow can make the non-zero envelope differ by head, so retain one @@ -1079,8 +1145,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_kernel( const int h = h_begin + j; float * out = dst + ((size_t) t * (size_t) n_heads + (size_t) h) * D + D - 64; - out[2 * pair + 0] = y0; - out[2 * pair + 1] = y1; + *reinterpret_cast(out + 2 * pair) = make_float2(y0, y1); } } } @@ -1251,14 +1316,21 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( if (visible) { const KV * kr = k + (size_t) r * D; #pragma unroll - for (int d = 0; d < D; ++d) { + for (int d = 0; d < D - 64; ++d) { + const float kv = ds4_fa_load(kr + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + dot[j] += qh[j][d] * kv; + } + } +#pragma unroll + for (int d = D - 64; d < D; ++d) { const float kv = ds4_fa_load(kr + d); #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - const float qv = - inverse_rope.forward_q_enabled && d >= D - 64 - ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] - : qh[j][d]; + const float qv = inverse_rope.forward_q_enabled + ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] + : qh[j][d]; dot[j] += qv * kv; } } @@ -1273,28 +1345,35 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( } } + // Match ds4_fa_block_max independently for every grouped head via warp shuffles. + const int lane = tid & 31; + const int warp_id = tid >> 5; #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - reduction[(size_t) j * N_THREADS + tid] = local_max[j]; + const float v = ds4_warp_reduce_max(local_max[j]); + if (lane == 0) { + reduction[(size_t) j * 8 + warp_id] = v; + } } __syncthreads(); - for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { - if (tid < stride) { + float max_score[HEADS_PER_BLOCK]; + if (warp_id == 0) { #pragma unroll - for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - float * row = reduction + (size_t) j * N_THREADS; - row[tid] = fmaxf(row[tid], row[tid + stride]); + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float v = (lane < 8) ? reduction[(size_t) j * 8 + lane] : -3.402823466e38f; + v = ds4_warp_reduce_max(v); + if (lane == 0) { + reduction[j] = v; } } - __syncthreads(); } - - float max_score[HEADS_PER_BLOCK]; - float local_sum[HEADS_PER_BLOCK] = {}; + __syncthreads(); #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - max_score[j] = reduction[(size_t) j * N_THREADS]; + max_score[j] = reduction[j]; } + float local_sum[HEADS_PER_BLOCK] = {}; + for (int iteration = 0; iteration < score_iteration_count; ++iteration) { int score_index; int bound_value; @@ -1342,26 +1421,30 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( } } + // Match ds4_fa_block_sum independently for every grouped head via warp shuffles. #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - reduction[(size_t) j * N_THREADS + tid] = local_sum[j]; + const float v = ds4_warp_reduce_sum(local_sum[j]); + if (lane == 0) { + reduction[(size_t) j * 8 + warp_id] = v; + } } __syncthreads(); - for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { - if (tid < stride) { + float inv_denom[HEADS_PER_BLOCK]; + if (warp_id == 0) { #pragma unroll - for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - float * row = reduction + (size_t) j * N_THREADS; - row[tid] += row[tid + stride]; + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float v = (lane < 8) ? reduction[(size_t) j * 8 + lane] : 0.0f; + v = ds4_warp_reduce_sum(v); + if (lane == 0) { + reduction[j] = 1.0f / v; } } - __syncthreads(); } - - float inv_denom[HEADS_PER_BLOCK]; + __syncthreads(); #pragma unroll for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - inv_denom[j] = 1.0f / reduction[(size_t) j * N_THREADS]; + inv_denom[j] = reduction[j]; } // Finish the shared envelope before the value phase reads it. @@ -1519,11 +1602,11 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( } else { float * out = dst + ((size_t) t * (size_t) n_heads + (size_t) h) * D + d0; - out[0] = value0; - out[1] = value1; if constexpr (VALUES_PER_THREAD == 4) { - out[2] = acc2[j] * inv_denom[j]; - out[3] = acc3[j] * inv_denom[j]; + *reinterpret_cast(out) = make_float4( + value0, value1, acc2[j] * inv_denom[j], acc3[j] * inv_denom[j]); + } else { + *reinterpret_cast(out) = make_float2(value0, value1); } } } @@ -1549,8 +1632,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( const int h = h_begin + j; float * out = dst + ((size_t) t * (size_t) n_heads + (size_t) h) * D + D - 64; - out[2 * pair + 0] = y0; - out[2 * pair + 1] = y1; + *reinterpret_cast(out + 2 * pair) = make_float2(y0, y1); } } } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index dd838f237..73ad8afdf 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -113,8 +113,8 @@ struct tile_x_sizes { #ifndef LUCEBOX_RDNA_MMQ_TILE_OVERRIDE #define LUCEBOX_RDNA_MMQ_TILE_OVERRIDE 1 #endif -#define LUCEBOX_RDNA_TILE_HOST(cc) (LUCEBOX_RDNA_MMQ_TILE_OVERRIDE && (GGML_CUDA_CC_IS_RDNA3(cc) || GGML_CUDA_CC_IS_RDNA4(cc))) -#if LUCEBOX_RDNA_MMQ_TILE_OVERRIDE && (defined(RDNA3) || defined(RDNA4)) +#define LUCEBOX_RDNA_TILE_HOST(cc) (LUCEBOX_RDNA_MMQ_TILE_OVERRIDE && (GGML_CUDA_CC_IS_RDNA3_5(cc) || GGML_CUDA_CC_IS_RDNA4(cc))) +#if LUCEBOX_RDNA_MMQ_TILE_OVERRIDE && (defined(RDNA3_5) || defined(RDNA4)) #define LUCEBOX_RDNA_TILE_DEVICE 1 #else #define LUCEBOX_RDNA_TILE_DEVICE 0 @@ -122,6 +122,9 @@ struct tile_x_sizes { static int get_mmq_x_max_host(const int cc) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return 48; + } #if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) return 64; #else @@ -139,7 +142,9 @@ static int get_mmq_x_max_host(const int cc) { static constexpr __device__ int get_mmq_x_max_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(RDNA3_5) + return 48; +#elif defined(GGML_CUDA_ROCMFPX_MMQ_TILE) return 64; #else return 128; @@ -169,6 +174,9 @@ static constexpr __device__ int get_mmq_x_max_device() { static int get_mmq_y_host(const int cc) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return 64; + } #if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) return 64; #else @@ -189,7 +197,9 @@ static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type static constexpr __device__ int get_mmq_y_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(RDNA3_5) + return 64; +#elif defined(GGML_CUDA_ROCMFPX_MMQ_TILE) return 64; #else return 128; @@ -342,6 +352,9 @@ static constexpr __device__ int mmq_get_granularity_device(const int /*mmq_x*/) #if defined(GGML_USE_HIP) static int mmq_get_nwarps_host(const int cc, const int warp_size) { if (LUCEBOX_RDNA_TILE_HOST(cc)) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return 4; + } #if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) return 4; #else @@ -358,7 +371,9 @@ static int mmq_get_nwarps_host(const int /*cc*/, const int warp_size) { static constexpr __device__ int mmq_get_nwarps_device() { #if LUCEBOX_RDNA_TILE_DEVICE -#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) +#if defined(RDNA3_5) + return 4; +#elif defined(GGML_CUDA_ROCMFPX_MMQ_TILE) return 4; #else return 8; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index d24bb282f..7fb49bcc8 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1813,6 +1813,8 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, DeepSeek4StepTelemetry step_tel; if (timing) step_tel.embed_us = elapsed_us(embed_t0, Clock::now()); + const bool at_snap_boundary = save_snapshot && !snapshot_saved && (pos + n_tok >= snap_pos); + const bool need_logits = (i + n_tok >= n_total) || at_snap_boundary; std::vector logits; bool ok = false; std::vector hc_state; @@ -1848,7 +1850,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, ok = deepseek4_step_layer_range( backend_, cfg_.device.gpu, w_, cache_, hc_state, embed.data(), n_tok, pos, - 0, w_.n_layer, &logits, + 0, w_.n_layer, need_logits ? &logits : nullptr, tokens.data() + i, timing ? &step_tel : nullptr, /*allow_decode_graph_reuse=*/true, hp, @@ -1862,11 +1864,12 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, timing ? &step_tel : nullptr, routing_stats_.get(), hp, - expert_runtime_.compute ? &expert_runtime_ : nullptr); + expert_runtime_.compute ? &expert_runtime_ : nullptr, + need_logits); } else { ok = deepseek4_step_layer_range(backend_, cfg_.device.gpu, w_, cache_, hc_state, embed.data(), n_tok, pos, - 0, w_.n_layer, &logits, + 0, w_.n_layer, need_logits ? &logits : nullptr, tokens.data() + i, timing ? &step_tel : nullptr, cfg_.prefill_mode != PrefillAttentionMode::Sparse, hp); @@ -1893,7 +1896,9 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, add_step_tel(tel_acc, step_tel); steps++; } - last_logits_ = std::move(logits); + if (need_logits) { + last_logits_ = std::move(logits); + } pos += n_tok; last_logits_pos_ = cache_.cur_pos; i += n_tok; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index cf4d4a9a2..d04fbca2a 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -4079,7 +4079,8 @@ static bool deepseek4_step_hybrid( MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats, - MoeExpertComputeRuntime * expert_runtime) { + MoeExpertComputeRuntime * expert_runtime, + bool need_logits = true) { const auto step_t0 = Ds4TimingClock::now(); const int n_embd = w.n_embd; const int n_hc = w.n_hc; @@ -4495,6 +4496,15 @@ static bool deepseek4_step_hybrid( if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); + if (!need_logits) { + out_logits.clear(); + cache.cur_pos = kv_start + n_tokens; + if (telemetry) { + telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + } + return true; + } + // ── Output HC pre → norm → logits ─────────────────────────────────── const auto output_t0 = Ds4TimingClock::now(); std::vector final_embd((size_t)n_embd * (size_t)n_tokens); @@ -4576,7 +4586,8 @@ bool deepseek4_step( DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats, Ds4VerifyHooks * verify_hooks, - MoeExpertComputeRuntime * expert_runtime) { + MoeExpertComputeRuntime * expert_runtime, + bool need_logits) { if (w.moe_hybrid && moe_hybrid != nullptr) { if (!deepseek4_cuda_hc_set_device(device)) { std::fprintf(stderr, @@ -4587,13 +4598,13 @@ bool deepseek4_step( return deepseek4_step_hybrid(backend, w, cache, *moe_hybrid, embed, n_tokens, kv_start, out_logits, token_ids, stream_engine, telemetry, routing_stats, - expert_runtime); + expert_runtime, need_logits); } std::vector hc_state; return deepseek4_step_layer_range( backend, device, w, cache, hc_state, embed, n_tokens, kv_start, - 0, w.n_layer, &out_logits, token_ids, telemetry, + 0, w.n_layer, need_logits ? &out_logits : nullptr, token_ids, telemetry, /*allow_decode_graph_reuse=*/verify_hooks == nullptr, verify_hooks, /*moe_hybrid=*/nullptr, expert_runtime, routing_stats); } @@ -6002,6 +6013,7 @@ struct Ds4LayerMajorGraphCache { PrefillAttentionMode mode = PrefillAttentionMode::Exact; int n_tokens = 0; int kv_start = -1; + bool has_logits = false; bool ready = false; ggml_context * state_ctx = nullptr; ggml_backend_buffer_t state_buf = nullptr; @@ -6010,9 +6022,11 @@ struct Ds4LayerMajorGraphCache { std::vector layers; bool matches(const DeepSeek4Weights & w, ggml_backend_t b, - PrefillAttentionMode m, int tokens, int start) const { + PrefillAttentionMode m, int tokens, int start, + bool logits_needed) const { return ready && owner_ctx == w.ctx && backend == b && mode == m && n_tokens == tokens && kv_start == start && + has_logits == logits_needed && layers.size() == (size_t) w.n_layer; } @@ -6034,6 +6048,7 @@ struct Ds4LayerMajorGraphCache { mode = PrefillAttentionMode::Exact; n_tokens = 0; kv_start = -1; + has_logits = false; ready = false; } }; @@ -6113,7 +6128,7 @@ static int ds4_try_layer_major_prefill( const float * embed, int n_tokens, int kv_start, - std::vector & out_logits, + std::vector * out_logits, const int32_t * token_ids, Ds4VerifyHooks * verify_hooks, DeepSeek4StepTelemetry * telemetry) { @@ -6225,10 +6240,11 @@ static int ds4_try_layer_major_prefill( Ds4LayerMajorGraphCache * graph_cache = nullptr; bool cache_hit = false; bool cache_build = false; + const bool logits_needed = (out_logits != nullptr); if (token_ids) { for (auto & candidate : ds4_layer_major_graph_caches) { if (candidate.matches(w, backend, cache.prefill_mode, - n_tokens, kv_start)) { + n_tokens, kv_start, logits_needed)) { graph_cache = &candidate; cache_hit = true; break; @@ -6239,11 +6255,14 @@ static int ds4_try_layer_major_prefill( // Do not evict a full/larger chunk for an equal-size graph at a // later position or for a short tail. Both execute with the shared // scratch arena below, but only the dominant topology stays cached. + // When logits are needed on a tail/terminal step, execute it + // transiently rather than evicting the dominant no-logits graph. const bool same_owner = candidate.owner_ctx == w.ctx && candidate.backend == backend && candidate.mode == cache.prefill_mode; - if (!candidate.ready || !same_owner || - n_tokens > candidate.n_tokens) { + const bool can_cache_dominant = !logits_needed; + if (can_cache_dominant && (!candidate.ready || !same_owner || + n_tokens > candidate.n_tokens)) { graph_cache = &candidate; graph_cache->destroy(); graph_cache->owner_ctx = w.ctx; @@ -6251,6 +6270,7 @@ static int ds4_try_layer_major_prefill( graph_cache->mode = cache.prefill_mode; graph_cache->n_tokens = n_tokens; graph_cache->kv_start = kv_start; + graph_cache->has_logits = false; graph_cache->layers.resize((size_t) w.n_layer); cache_build = true; } @@ -6392,10 +6412,10 @@ static int ds4_try_layer_major_prefill( compute_t0, Ds4TimingClock::now()); } capture_layer(il, (il & 1) == 0 ? state_b : state_a); - if (layer.logits) { - out_logits.resize((size_t) w.n_vocab); + if (layer.logits && out_logits) { + out_logits->resize((size_t) w.n_vocab); ggml_backend_tensor_get( - layer.logits, out_logits.data(), 0, + layer.logits, out_logits->data(), 0, sizeof(float) * (size_t) w.n_vocab); } @@ -6410,7 +6430,7 @@ static int ds4_try_layer_major_prefill( } } cache.cur_pos = next_pos; - return out_logits.empty() ? -1 : 1; + return (out_logits && out_logits->empty()) ? -1 : 1; } ggml_tensor * state_in = state_a; @@ -6529,7 +6549,7 @@ static int ds4_try_layer_major_prefill( ggml_build_forward_expand(gf, state_copy); ggml_tensor * logits = nullptr; - if (il + 1 == w.n_layer) { + if (out_logits && il + 1 == w.n_layer) { ggml_tensor * last_hc = ggml_view_2d( ctx, hc_next, hc_dim, 1, hc_next->nb[1], (size_t) (n_tokens - 1) * hc_next->nb[1]); @@ -6620,9 +6640,9 @@ static int ds4_try_layer_major_prefill( capture_layer(il, state_out); - if (logits) { - out_logits.resize((size_t) w.n_vocab); - ggml_backend_tensor_get(logits, out_logits.data(), 0, + if (logits && out_logits) { + out_logits->resize((size_t) w.n_vocab); + ggml_backend_tensor_get(logits, out_logits->data(), 0, sizeof(float) * (size_t) w.n_vocab); } @@ -6654,7 +6674,7 @@ static int ds4_try_layer_major_prefill( ggml_free(state_ctx); } cache.cur_pos = next_pos; - return out_logits.empty() ? -1 : 1; + return (out_logits && out_logits->empty()) ? -1 : 1; } static bool ds4_hc_layer_weights_ready(const HcWeightsCpu & weights, @@ -6800,7 +6820,7 @@ bool deepseek4_step_layer_range( !fused_verify_candidate && moe_hybrid && cache.prefill_mode == PrefillAttentionMode::Sparse && n_tokens > 4 && n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && - layer_begin == 0 && is_last_shard && out_logits && + layer_begin == 0 && is_last_shard && ds4_backend_is_gpu(backend); const bool layer_major_hooks_supported = !verify_hooks || @@ -6813,7 +6833,7 @@ bool deepseek4_step_layer_range( const bool standard_layer_major_prefill = !w.moe_hybrid && cache.prefill_mode != PrefillAttentionMode::Exact && n_tokens > 4 && n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && - layer_begin == 0 && is_last_shard && out_logits && + layer_begin == 0 && is_last_shard && ds4_backend_is_gpu(backend) && layer_major_hooks_supported; // These graphs are rebuilt around an owner join on every layer, so tensor // metadata addresses can be recycled for different topologies. Until @@ -7007,7 +7027,7 @@ bool deepseek4_step_layer_range( fused_decode_graph_cache, backend, w, cache, hc_layer_weights_range, hc_output_weights_range, hash_routing_tables_range, scratch.hash_expert_ids, embed, - n_tokens, kv_start, *out_logits, token_ids, verify_hooks, + n_tokens, kv_start, out_logits, token_ids, verify_hooks, telemetry); if (prc < 0) return false; if (prc > 0) { @@ -7148,7 +7168,7 @@ bool deepseek4_step_layer_range( use_backend_decode_hc && !use_backend_decode_hc_direct; const bool use_backend_prefill_hc = heterogeneous_sparse_prefill && - ds4_env_flag("DFLASH_DS4_HYBRID_PREFILL_GPU_HC"); + !ds4_env_flag("DFLASH_DS4_DISABLE_HYBRID_PREFILL_GPU_HC"); ggml_tensor * hc_state_backend = nullptr; if (use_backend_prefill_hc) { if (!ds4_fused_ensure_fn_mirrors( diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index eea667423..6e86ab949 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -359,7 +359,7 @@ int deepseek4_safe_compressor_batch_tokens(const DeepSeek4Weights & w, // Forward: single step (prefill chunk or decode token). // embed: [n_embd, n_tokens] input embeddings (post-embedding lookup). // hc_state: [n_hc * n_embd] persistent HC residual (updated in-place). -// Returns logits for last token. +// Returns logits for last token in out_logits when need_logits is true; out_logits is left untouched/empty when need_logits is false. struct Ds4VerifyHooks; bool deepseek4_step( @@ -377,7 +377,8 @@ bool deepseek4_step( DeepSeek4StepTelemetry * telemetry = nullptr, MoeHybridRoutingStats * routing_stats = nullptr, Ds4VerifyHooks * verify_hooks = nullptr, - MoeExpertComputeRuntime * expert_runtime = nullptr); + MoeExpertComputeRuntime * expert_runtime = nullptr, + bool need_logits = true); // Optional hooks for the DSpark spec-decode batched verify (deepseek4_dspark). // When set on a multi-token deepseek4_step_layer_range call they add: per-layer