Skip to content

perf(ds4): elide intermediate prefill logits #633

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
cheese-cakee wants to merge 4 commits into
base: main
Choose a base branch
from
+72 −31
Draft

perf(ds4): elide intermediate prefill logits #633

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8992588
perf(ds4): elide intermediate prefill logits and tune gfx1151 mmq tiles
cheese-cakee Aug 18, 2026
b434977
fix(ds4): address review comments for prefill readout elision and mmq…
cheese-cakee Aug 19, 2026
92366b4
fix(ds4): differentiate layer-major graph cache by logits presence an…
cheese-cakee Aug 19, 2026
a2878db
fix(ds4): preserve dominant no-logits graph cache across transient te…
cheese-cakee Aug 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 20 additions & 5 deletions server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,15 +113,18 @@ 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
#endif

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
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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;
Expand Down
13 changes: 9 additions & 4 deletions server/src/deepseek4/deepseek4_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1813,6 +1813,8 @@ int DeepSeek4Backend::do_prefill(const std::vector<int32_t> & 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<float> logits;
bool ok = false;
std::vector<float> hc_state;
Expand Down Expand Up @@ -1848,7 +1850,7 @@ int DeepSeek4Backend::do_prefill(const std::vector<int32_t> & 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,
Comment thread
cheese-cakee marked this conversation as resolved.
tokens.data() + i,
timing ? &step_tel : nullptr,
/*allow_decode_graph_reuse=*/true, hp,
Expand All @@ -1862,11 +1864,12 @@ int DeepSeek4Backend::do_prefill(const std::vector<int32_t> & 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);
Expand All @@ -1893,7 +1896,9 @@ int DeepSeek4Backend::do_prefill(const std::vector<int32_t> & 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;
Expand Down
62 changes: 41 additions & 21 deletions server/src/deepseek4/deepseek4_graph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<float> final_embd((size_t)n_embd * (size_t)n_tokens);
Expand Down Expand Up @@ -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,
Expand All @@ -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<float> 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);
}
Expand Down Expand Up @@ -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;
Expand All @@ -6010,9 +6022,11 @@ struct Ds4LayerMajorGraphCache {
std::vector<Ds4LayerMajorCachedLayer> 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;
}

Expand All @@ -6034,6 +6048,7 @@ struct Ds4LayerMajorGraphCache {
mode = PrefillAttentionMode::Exact;
n_tokens = 0;
kv_start = -1;
has_logits = false;
ready = false;
}
};
Expand Down Expand Up @@ -6113,7 +6128,7 @@ static int ds4_try_layer_major_prefill(
const float * embed,
int n_tokens,
int kv_start,
std::vector<float> & out_logits,
std::vector<float> * out_logits,
const int32_t * token_ids,
Ds4VerifyHooks * verify_hooks,
DeepSeek4StepTelemetry * telemetry) {
Expand Down Expand Up @@ -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;
Expand All @@ -6239,18 +6255,22 @@ 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;

@Graffioh Graffioh Aug 25, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does this mean single-chunk prompts never populate the cache? eg --chunk 512, a 5–512 token prompt goes straight to the terminal/logits path, with logits_needed always true, so can_cache_dominant is false and every request rebuilds the layer-major graph

repeated single-chunk prompts with the same shape rebuild the layer-major graph on every request...was that case benchmarked, or is the build cost small enough that it does not matter?

extra: A snapshot boundary also sets need_logits to true because restoring that snapshot later requires the logits from that exact position.

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;
graph_cache->backend = backend;
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;
}
Expand Down Expand Up @@ -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);
}

Expand All @@ -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;
Expand Down Expand Up @@ -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) {
Comment thread
cubic-dev-ai[bot] marked this conversation as resolved.
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]);
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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 ||
Expand All @@ -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
Expand Down Expand Up @@ -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) {
Expand Down
3 changes: 2 additions & 1 deletion server/src/deepseek4/deepseek4_internal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading