Split mode graph for dense Gemma4 assistant by ikawrakow · Pull Request #2022 · ikawrakow/ik_llama.cpp

ikawrakow · 2026-06-24T08:49:41Z

The Gemma4 assistant models reuse the KV cache of the main model. The implementation on the main branch forces split mode layer for the assistant model. As a result, when split mode graph is used for the main mode, one needs to combine the KV cache split over the participating GPUs into a single KV cache, which is a fairly costly operation.

This PR adds the ability to use split mode graph for Gemma4 assistant models. The split mode used for the assistant is forced to be the same as the split mode of the main model.

We get in the range of 10% better TG for Gemma4-31B compared to the main branch.

The PR is bigger than what one might have expected because I had to extend the mechanism for managing CUDA contexts to allow for more than one set of contexts, which is necessary for correctly performing AllReduce operations for the main and the assistant models.

Something is not right - acceptance drops to nearly zero.

Still not working!?

The issue was that I was not correctly calculating the number of KV heads for the split KV cache.

…et with <2 devices When a GEMMA4_ASSISTANT draft model is loaded with fewer devices than the target GEMMA4 model (e.g. via -devd 0), the split_mode inheritance at llm_load_tensors blindly copied the target's split_mode (GRAPH/ATTN) without validating the draft model's device count. This caused a crash: 'It is not possible to use split mode graph with less than 2 devices'. Now after inheriting split_mode, we check device count and fall back to LAYER if insufficient devices are available. PR ikawrakow#2022 introduced the inheritance, PR ikawrakow#2027 extended it to Qwen35 MTP.

… target with <2 devices When a GEMMA4_ASSISTANT draft model is loaded with fewer devices than the target GEMMA4 model (e.g. via -devd 0), the split_mode inheritance at llm_load_tensors blindly copied the target's split_mode (GRAPH/ATTN) without validating the draft model's device count. This caused a crash: 'It is not possible to use split mode graph with less than 2 devices'. Now after inheriting split_mode, we check device count and fall back to LAYER if insufficient devices are available. PR ikawrakow#2022 introduced the inheritance, PR ikawrakow#2027 extended it to Qwen35 MTP.

MTP: asynchronous hidden state capture — move sync out of feature view path The llama_spec_prepare_hidden_feature_view function called llama_synchronize to ensure the async GPU→CPU copy of output embeddings (ctx->embd) completed before reading the data. This sync was redundant — the embd copy was submitted earlier inside llama_decode_internal via ggml_backend_tensor_get_async. Move the ggml_backend_sched_synchronize into llama_decode_internal right after the embd get_async submission (for MTP modes only). This ensures ctx->embd is always ready by the time any feature view function accesses it, removing the need for a synchronize in the speculative capture path. The sync happens microseconds earlier (right after copy submission), so total wait time is unchanged. However this decouples data readiness from the capture API — callers can read ctx->embd without blocking. llama: persist backend sampling topk tensor for CUDA graph compatibility The TOPK node for MTP draft backend sampling was allocated from a throwaway ggml_init context every decode, causing: 1. A memory leak (the context was never freed after use) 2. The tensor's data address changed on every call, defeating CUDA graph replay (is_cuda_graph_update_required detected the address mismatch and forced a full graph re-capture each decode) Fix: store a persistent ggml_context + ggml_tensor in llama_context. On each decode, reuse the cached tensor if K and nrows match; reallocate only when these shapes change (rare). The persistent allocation keeps the tensor's data address stable, so the CUDA graph update check finds matching properties and skips the re-capture. Cleanup added to llama_context::~llama_context. topk: initialize stale cross-warp reduction lanes in CUDA kernel When n_warps < 32, lanes [n_warps..31] in warp 0 retain their stale intra-warp max from the per-thread scan. This was safe in practice (the stale value is warp 0's own max, which is <= the true global max) but fragile: NaN in the stale value or an edge-case ordering could produce wrong results. Fix: initialize maxval = -FLT_MAX, argmax = -1 for all lanes >= n_warps before the final shuffle reduction. topk: fix reader index arithmetic to match corrected 3D tensor strides The ggml_topk output has shape [2, K, nrows] with C-order strides. Value for rank k at row r resides at float offset 2*k + 2*K*r, and the index resides at offset 2*k + 2*K*r + 1. Update the decomposition loop in llama_decode_internal to use these offsets, matching the writer fixes in the previous commit. topk: fix output indexing to use correct 3D C-order strides for [2, K, nrows] tensor The ggml_topk output tensor has shape ne = [2, K, nrows]. Element (plane p, rank k, row r) must be stored at float offset p + 2*k + 2*K*r per ggml's 3D C-order layout (nb[0]=4, nb[1]=8, nb[2]=8*K). The writer (CPU forward + CUDA kernel) was using k*nrows + r and (K+k)*nrows + r, which does not match the stride layout. Both writer and reader happened to use the same wrong convention, so data flowed correctly inside the system, but the raw buffer violated ggml tensor metadata. Fix the writer to use the correct stride math. Also fix the K=1 fast path in the CPU forward (out[r] -> out[2*r], out[1*nrows + r] -> out[2*r + 1]). ggml: replace O(N*K^2/K) topk CPU forward with O(N*log K) min-heap The previous CPU topk implementation used a K-iteration approach: for each of K selections, it scanned all N columns and skipped already-found indices via a malloc/free bitmap. This was O(N*K) per row with additional allocation overhead (calloc/free per row). Replace with a single-pass min-heap: 1. Build a min-heap of size K from the first K elements (O(K)) 2. Scan the remaining N-K elements: if element > heap root (minimum), replace root and sift down (O((N-K)*log K)) 3. Extract from the heap in reverse order for descending output (O(K*log K)) Total: O(N*log K) per row, no heap allocations. Memory: two fixed-size stack arrays (float val[256], int32_t idx[256]) for the heap — the K <= 256 bound is now asserted at the ggml_topk() API entry. This eliminates the calloc/free per row (fix ikawrakow#3) and the O(N*K) scan (fix ikawrakow#2). ggml: add K<=256 assertion to ggml_topk() API function The CPU forward path uses a stack-allocated array int32_t selected[256], so K > 256 causes a silent buffer overflow. The assertion was only present in the CUDA kernel (K <= 4096) but missing at the API entry point. Move the bound check to ggml_topk() so both CPU and CUDA paths are protected from the start. fix: align backend_sampling default to true in common_params_speculative The default value in common_params_speculative::backend_sampling was false, but the comment claimed 'default: on' and llama_context_default_params() sets it to true. Therefore, CLI users had to explicitly pass -sdbs to enable GPU backend sampling for MTP drafts, even though the API default was on. Changed the default to true so --spec-draft-backend-sampling is 'on' by default, matching both the comment and the llama.h API contract. cuda: rename blocks_dim/blocks_num to block_dim/grid_dim in topk_f32 launch The variable names were misleading: - blocks_dim was the per-block thread count (blockDim), not a 'blocks' dimension - blocks_num was the grid dimension (gridDim), not a plural of 'number' Renamed to block_dim and grid_dim to match CUDA naming conventions used in other files (ggml-cuda.cu). The kernel launch <<<grid_dim, block_dim>>> now clearly separates grid (number of blocks) from block (threads per block). topk: skip bitmap init/access when K=1 (argmax fast path) K=1 needs no bitmap at all: only one pass through the data, no previously-selected indices to skip. Guard the bitmap init loop, the bitmap skip check in the scan, and the bitmap-set at result write with if (K > 1). The compiler's branch predictor perfectly predicts this static per-launch condition. CPU forward: add a dedicated K==1 branch that avoids malloc/ free of the bitmap altogether, just a single linear scan. spec: add --spec-draft-max-candidates (-sdmc) CLI flag Wire through common_params_speculative, llama_context_params, llama_cparams, and both speculative init paths (load_draft_model and prepare_mtp_runtime). Default 0 = argmax (K=1). rename max_candidates -> spec_max_candidates in cparams and graph building Avoid confusion with sampler's --max-candidates / -mc. topk: replace O(k) skip loop with shared-memory bitmap Replace the linear scan over already-found indices (O(k) per column) with a shared-memory bitmap (O(1) per column). Bump K limit from 256 to 4096 since the skip is no longer a bottleneck. Both CUDA and CPU forward paths updated. backend sampling: ggml_topk F32 values+indices, wire p_min>0 path ggml_topk output changes from I32 [K, nrows] to F32 [2, K, nrows]: plane 0 = top-K logit values (float32) plane 1 = top-K indices stored as float32 (cast to int32_t on CPU) This enables the MTP draft caller to obtain both token IDs and their raw logit values from the GPU in a single tiny D2H transfer, without materializing the full n_vocab logits buffer (~800 KB per token). Changes by layer: ggml: - ggml_topk() creates F32 [2, k, ne[1]] tensor instead of I32 [k, ne[1]] - CPU forward writes max_val + max_idx as float pairs - CUDA kernel takes nrows parameter, stores value at dst[k*nrows + row] and index at dst[(K+k)*nrows + row] llama.cpp graph building: - Always uses ggml_topk for backend_sampling (replaces ggml_argmax), with K = max(1, max_candidates). The values+indices output works for both greedy (K=1, only top-1 needed) and p_min>0 (K>=2). llama.cpp extraction: - Reads F32 tensor, decomposes into logits_argmax (int32 indices) and new logits_argmax_values (float values). Tracks K in logits_argmax_K. Clears state when no tensor present. Public API: - llama_get_logits_argmax_ith() updated to use K-aware offset - New llama_get_logits_topk_ith(ctx, i, k, &val) returns token ID and optional logit value for rank k at position i. Returns LLAMA_TOKEN_NULL if k >= K or out of range. speculative.cpp (mtp_speculative_gen_draft): - skip_d2h = backend_sampling (extends D2H skip to p_min>0 too) - New p_min>0 branch: reads top-K values from backend, computes approximate softmax probability 1/(1 + sum exp(v_k - v_0)) over the K candidates. Falls back to full logits sampler if backend returns NULL (defensive only). - With K=1 (default max_candidates=0), the sum has only the top-1 term, so prob = 1.0 (no early termination). Users should set --spec-draft-max-candidates >= 2 for meaningful p_min checks. - The approximate probability overestimates true probability (only sums K of n_vocab terms), making p_min early termination slightly conservative. This is safe (we generate slightly more draft tokens that get verified) and correct for peaked MTP head distributions. llama: skip full logits D2H when backend_sampling and p_min==0 Add llama_set_skip_logits_d2h() API and skip_logits_d2h flag in llama_cparams. When set during MTP draft decode with backend_sampling, llama_decode_internal reads only the GPU-side argmax/topk result and suppresses the full n_vocab logits D2H copy. Wire this into mtp_speculative_gen_draft: set skip_logits_d2h before each llama_decode call when backend_sampling && p_min == 0.0f, then reset after. This avoids transferring ~800KB per draft token, keeping only the 4-160 byte argmax/topk D2H. llama: wire max_candidates into backend sampling graph When max_candidates > 1, use GGML_OP_TOPK instead of GGML_OP_ARGMAX in the MTP draft graph. This enables the -mc / --max-candidates flag to limit logit transfer from GPU to CPU when p_min > 0. - Add max_candidates field to llama_cparams - Graph builder selects ggml_topk vs ggml_argmax based on max_candidates - Use ggml_nelements for output size (handles both 1D argmax and 2D topk) fix: move max_candidates pre-filter after llama_prep_adaptive_p The --max-candidates / -mc top-k pre-filter at the start of llama_sampling_sample_impl ran before llama_prep_adaptive_p, reducing the candidate array from n_vocab to max_candidates. llama_prep_adaptive_p requires the full vocabulary (checks candidates->size != orig_prob.size()) and aborted with: sampler initialized with vocabulary of 151552, but called with 256 candidates The same pre-filter already exists inside sampler_queue() which is called for the adaptive_p and temperature paths. Removing the standalone copy in llama_sampling_sample_impl fixes the ordering: adaptive_p prep sees the full vocabulary, and sampler_queue applies the reduction afterward. sampling: add --max-candidates / -mc prefilter option - Add max_candidates field to common_params_sampling (default: 0 = disabled) - Add CLI args: --max-candidates / -mc to specify max candidates to keep - Apply prefilter in llama_sampling_sample_impl before all sampling paths - Also apply in sampler_queue for temperature sampling path Use -mc 1 for deterministic inference (always picks highest logit). Use -mc 256 for speed optimization (prefilter before top_p, etc.). cuda: add GPU top-K kernel (ggml_cuda_topk) New CUDA kernel topk_f32 that finds top-K indices per row in a single kernel launch. Uses K-iteration approach: each iteration finds the argmax among elements not yet selected, broadcasting selected indices through shared memory. - One block per row, K up to 256 - Warp shuffle + cross-warp shared-memory reduction - Registered in CUDA backend dispatch and supports_op ggml: add GGML_OP_TOPK operation Add a new ggml operation for top-K selection along rows. - GGML_OP_TOPK enum entry, ggml_topk() API - output shape: [K, nrows] int32, K stored in op_params[0] - CPU forward compute: simple O(N*K) selection per row - op name, symbol, dispatch, task count, backward stub Arg -sdbs: backend sampling for MTP draft path: add GPU argmax to draft graph Add an optional GGML_OP_ARGMAX node to the MTP draft computation graph, replacing the CPU-side argmax over the full n_vocab logits with a GPU-side reduction. The result (1 int32 per draft token) is read back via a new API llama_get_logits_argmax_ith(), avoiding the full D2H logits transfer. - llama_context_params, llama_cparams, common_params_speculative: add backend_sampling bool (default: true) - llama_decode_internal: when mtp_op_type == DRAFT_GEN && backend_sampling, insert ggml_argmax(result_output) into the graph and read the result into lctx.logits_argmax after compute - llama_get_logits_argmax_ith(ctx, i): public API to retrieve the GPU-computed argmax token ID for draft position i; returns LLAMA_TOKEN_NULL when unavailable - mtp_speculative_gen_draft: when backend_sampling && p_min == 0.0 (greedy), use llama_get_logits_argmax_ith instead of common_sampler_sample_speculative; for p_min > 0 fall back to the full logits path (needs softmax denominator) - --spec-draft-backend-sampling / --no-spec-draft-backend-sampling CLI flags - Wire backend_sampling from common_params_speculative into cparams_dft in both common_speculative_load_draft_model and common_speculative_prepare_mtp_runtime delta_net: remove redundant state contiguity check and assertion The state tensor from the KV cache has shape [S_v, S_v, H_v, 1] with a strided nb[3] (full KV row stride). After reshape to [S_v, S_v*H_v, 1, 1], ne[2]=ne[3]=1 causes ggml_is_contiguous to skip checking dims 2 and 3, so the ggml_cont_4d block was always dead code. Remove the assertion in ggml_delta_net for the same reason — with n_seqs=1 (the only caller path), the kernel index calc uses batch_idx=0, making the last-dim stride irrelevant. Verified CPU and CUDA kernels handle the strided input correctly. qwen35: capture post-norm hidden state for MTP The MTP hidden state should be captured after output_norm (post-norm) rather than from the raw residual stream (pre-norm inpL). Qwen3.5's MTP tail expects the normalized hidden state, matching what is fed into the lm_head. Using pre-norm directly degrades the acceptance rate. enable inp_out_ids for MTP models in non-MTP graph path Remove !lctx.cparams.mtp guard from inp_out_ids in build_glm4_moe(), build_qwen35moe(), and build_qwen35(). This lets batched non-MTP inference on MTP models skip the lm_head for non-logit positions. The MTP tail helpers (build_glm4_moe_mtp, build_qwen35moe_mtp, build_qwen35_mtp) already used inp_out_ids unconditionally, so MTP update batches were already optimized. fix: generalize result_output search in decode graph When backend sampling adds a result_argmax node as the last graph node, the existing fallback search for result_output only triggered for DFlash (lctx.dflash.draft_tokens_tensor != nullptr). MTP draft gen with backend_sampling hit the same issue but lacked the search, causing a GGML_ASSERT failure. Fix: search for result_output whenever the last node is not named result_output, removing the DFlash-specific guard. MTP: Fix GEMMA4_ASSISTANT crash when inheriting split_mode GRAPH from target with <2 devices When a GEMMA4_ASSISTANT draft model is loaded with fewer devices than the target GEMMA4 model (e.g. via -devd 0), the split_mode inheritance at llm_load_tensors blindly copied the target's split_mode (GRAPH/ATTN) without validating the draft model's device count. This caused a crash: 'It is not possible to use split mode graph with less than 2 devices'. Now after inheriting split_mode, we check device count and fall back to LAYER if insufficient devices are available. PR ikawrakow#2022 introduced the inheritance, PR ikawrakow#2027 extended it to Qwen35 MTP. MTP: fix dead code in mtp_model_uses_recurrent_conditioning The function had an unconditional 'return true;' before the actual architecture-specific logic, causing every MTP model to be treated as using recurrent conditioning (shifted hidden rows). This forced unnecessary hidden-state shifting for non-recurrent architectures. Removed the early return so the function correctly checks: 1. Is the model recurrent (has recurrent KV)? 2. Is the architecture qwen35 or qwen35moe? Only architectures that pass both checks get the shifted hidden-row treatment. MTP: Fix model loading when extra MTP output tensor is present When a model has been quantized with the extra MTP output tensor, the split metadata tensor count (LLM_KV_SPLIT_TENSORS_COUNT) may not reflect the additional tensor on some models downloadable on HF. Changed the strict equality check to allow extra tensors: missing tensors still trigger an error (corrupted model), but extra tensors only log a warning and allow loading to proceed. Fixes (example I had with some Qwen3.5/6 or Gemma 4 model, don't remember) : 'corrupted model: 866 tensors expected but 867 found'. The edit made it work. MTP: pre-reserve draft input hidden state vector capacity llama_set_draft_input_hidden_state_copy is called multiple times per decode step: before each draft decode (up to n_draft), during mtp_accept_batch (batch + last embd), and during prompt warmup. The std::vector::assign would reallocate on first call, then reuse capacity on subsequent calls. Pre-reserving capacity ensures the first call also avoids reallocation, and makes the allocation intent explicit. The size is stable (n_embd floats) after initialization, so the vector never reallocates after the first reservation.

ikawrakow added 5 commits June 24, 2026 07:10

WIP: Split mode graph for Gemma4 assistant

3530b65

Something is not right - acceptance drops to nearly zero.

Per model CUDA contexts

75a5f6d

Still not working!?

This works

118c82d

The issue was that I was not correctly calculating the number of KV heads for the split KV cache.

Compiler warnings

de6c2df

It is better to use llama_context pointers as keys

1f5828e

ikawrakow merged commit d5507e3 into main Jun 24, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Split mode graph for dense Gemma4 assistant #2022

Split mode graph for dense Gemma4 assistant #2022
ikawrakow merged 5 commits into
mainfrom
ik/g4_assistant_smgraph

ikawrakow commented Jun 24, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

ikawrakow commented Jun 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

ikawrakow commented Jun 24, 2026 •

edited

Loading