ik_llama_opt/PORTING-OPTIMIZATIONS.md

13 KiB
Raw Blame History

Agent porting guide: re-applying the Ampere+Zen3 optimizations

Use case: upstream ik_llama.cpp has moved (new model arch, refactored file, new kernels) and you must re-apply this branch's optimizations to the new tree. Work top-down: OPT-01 → OPT-08. Each item is self-contained: locate (exact anchor), transform (exact edit), verify.

Target rig (do not re-tune for other hardware): RTX 3090 24GB + RTX 3070 8GB (sm_86, no P2P) + Ryzen 5900XT 16C/32T (znver3, DDR4-3600). PP-bound long-context agentic sessions. Glove-fit non-portable builds are acceptable and expected.

Ground rules

  1. Arch-gate every model-specific change. Gate on model.arch == LLM_ARCH_<X> (loader) or put the code in src/graphs/build_<arch>.cpp (graph). Never change shared paths unconditionally.
  2. New CUDA kernels are opt-in env flags, default off, until A/B proves a win. Pattern: getenv("DELTA_WY_CUDA") → branch → exact fallback. See OPT-06.
  3. One optimization = one commit, message prefix: qwen4exp:, cuda:, tests:, ggml:.
  4. Validate after every item (see "Validation protocol"). If the anchor is gone (upstream refactor), locate the successor by searching the anchor's keywords, adapt the transform, and note the new anchor in the commit message.
  5. Reference implementation: commits e2728c85 (CPU/graph), a47a8a65/f5b494b0/3a7d1016 (harness + WY candidate), a34feeb0 (CUDA kernel), c1a36daa (harness CUDA fix), 7930df51 (MoE microbench), 2e84b994/253cf76f/f9ffb6fb (Bailing rescue + strip, OPT-08). Use git show <sha> -- <file> to see the exact original diff.

OPT-01 — Merged up/gate expert projections (loader)

  • Goal: one [2*n_ff, n_embd] GEMM per expert batch instead of two; halves expert weight-streaming traffic on CUDA mmq_id and CPU IQK paths.
  • Locate: src/llama-load-tensors.cpp, functions create_std_ffn_exps and create_std_ffn_exps_from_meta. Anchor string: ml.merge_up_gate_exps && merge_up_gate_exps(tn, i, 0) (appears in both functions, in the else branch where no fused ug_meta tensor exists in the file).
  • Transform: in both functions, replace the anchor condition with:
    const bool want_merge = ml.merge_up_gate_exps || model.arch == LLM_ARCH_<NEWARCH>;
    merged = flags == 0 && want_merge && merge_up_gate_exps(tn, i, 0);
    
    substituting the new arch enum. merge_up_gate_exps() already falls back gracefully when types/shapes differ — do not bypass it, only widen the opt-in.
  • New arch checklist: if upstream's new arch stores experts as separate up/gate tensors (check its create_*_tensors), add its enum here. If the file already ships a fused up-gate tensor, this item is a no-op.
  • Verify: model loads, llama-bench -p 512 -n 0 PP within noise of unpatched, no ggml_assert on expert shapes.

OPT-02 — Fused PLE / convolution tap accumulation (graph)

  • Goal: remove 2 transposes + 2 conts per tap per layer in persistent-memory convolution loops.
  • Locate: src/graphs/build_<arch>.cpp, the PLE conv helper (here qwen4exp_ple_conv). Anchor pattern: ggml_reshape_1d(...) on the tap weight followed by ggml_mul(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, shifted)), wk) and a re-transpose.
  • Transform:
    wk = ggml_reshape_2d(ctx0, wk, 1, hc_dim);   // [1, hc_dim] broadcasts over rows
    ...
    ggml_tensor * term = ggml_mul(ctx0, shifted, wk); // [n_tokens, hc_dim], no transposes
    
    Keep the existing ggml_cast(..., GGML_TYPE_F32) on wk. Math is identical (row-vector broadcast); confirm output shapes in comments.
  • New arch checklist: any new arch with a tap/conv accumulation loop over a persistent cache gets the same treatment in its build_<arch>.cpp.
  • Verify: perplexity spot-check (llama-perplexity, 12 shards of wikitext) matches unpatched to ≤0.001; PP bench neutral-to-positive.

OPT-03 — k_x_step 64 → 32 (CPU IQK MoE tiling)

  • Goal: Zen3 CCX locality for narrow expert GEMV; 64-wide reuse set thrashes the 2×32MB L3 on 5900XT.
  • Locate: ggml/src/iqk/iqk_mul_mat.cpp, struct MulMat, two sites (mul_mat_NxM and the unary-op variant). Anchor comment: This works best on my Ryzen-7950X. Both are in the #else (non-aarch64) branch of constexpr int k_x_step = 64;.
  • Transform: change both 6432, keep the #ifdef __aarch64__ branch untouched. Carry a comment: 32 for Zen3 CCX locality; Zen4/Intel diff is small.
  • Verify: test-moe-perf (OPT-07) TG/PP neutral-or-better on this rig; full model TG must not regress >1%.

OPT-04 — Expert batching chunk /32/64 (ggml thread scheduling)

  • Goal: halve atomics/barriers on 32-thread Zen3 for narrow experts.
  • Locate: ggml/src/ggml.c, ggml_compute_forward_mul_mat_id (anchor ne01 / 32) and ggml_compute_forward_mul_mat_id_up_gate (anchor nr0_base / 32). Full lines: MAX(1, MIN(nth, (int)(ne01 / 32))) and the _ug equivalent.
  • Transform: / 32/ 64 in both lines. The MIN(nth, ...) cap already protects TG load balance — do not touch it.
  • Verify: TG-heavy bench (-p 128 -n 128) must not regress; PP should improve or hold.

OPT-05 — IQ4_XS AVX2 software prefetch (Zen3 DDR4 latency hiding)

  • Goal: hide DDR4 latency on the 4-bit expert GEMM weight stream during TG.
  • Locate: ggml/src/iqk/iqk_gemm_kquants.cpp, function mul_mat_qX_K_q8_K_T, the for (int i = 0; i < nb; ++i) block-body loop. Anchor: auto all_scales = deq.new_block(i, q8, accd);.
  • Transform: immediately before the anchor line, insert:
    if (i + 4 < nb) {
        __builtin_prefetch(deq.x + i + 4, 0, 3);
    }
    
    Weight stream is ~270 B/block; 4 blocks ≈ 1 KB stays in the L1 streamer window. If upstream renames deq.x/new_block, adapt to the new dequant-struct member holding the next weight block.
  • Verify: TG tok/s neutral-or-better on IQ4_XS expert layers; remove if negative (single-line revert).

OPT-06 — Chunked WY delta-net (recurrent-state prefill) + CUDA kernel

  • Goal: chunked formulation of the GDN recurrent update with per-chunk exp(diff) rescaling; CUDA fast path for long prefills on hybrid archs.
  • Locate (CPU/reference): ggml/src/ggml.c recurrent/delta-net forward op used by the arch. Locate (CUDA): ggml/src/ggml-cuda/delta-net.cu, dispatch function containing the sequential path; anchor: saved_states == nullptr guard region.
  • Transform:
    1. Port tests/test-delta-chunk.cpp first (OPT-07) — it encodes the validated math (chunk sizes 64/32, exp-diff ratios, empty/oversized-chunk edge cases). The candidate must match the sequential reference bit-close on all 42 cases before any kernel work.
    2. Re-apply delta_net_chunked_f32 in delta-net.cu (see git show a34feeb0): shared-memory tiled kernel, tiers for head-dim 64/128 + sequential fallback, chunk constants DELTA_WY_CHUNK 64 / DELTA_WY_CHUNK_SMALL 32.
    3. Dispatch pattern (opt-in, default off):
    const char * wy_env = getenv("DELTA_WY_CUDA");
    bool wy_wanted = wy_env && wy_env[0] == '1';
    if (wy_wanted && n_tokens > DELTA_WY_CHUNK && saved_states == nullptr) {
        int chunk = (head_dim <= 64) ? DELTA_WY_CHUNK : DELTA_WY_CHUNK_SMALL;
        ... chunked path ...
    } else { ... existing sequential path, untouched ... }
    
  • Verify: harness 42/42 on CPU and --cuda; then server A/B (DELTA_WY_CUDA=1 vs unset) on ≥25k prefill. Promote to default-on only with ≥2% repeatable PP win; otherwise keep as opt-in infrastructure (current status: parity, default off).

OPT-07 — Test harnesses (port these, don't skip)

  • tests/test-delta-chunk.cpp (718 lines): correctness sweep (chunks × head-dims × dtypes × edge cases) + --cuda backend mode + chunk-size perf sweep. CUDA backend path must use a no-alloc ggml context (c1a36daa — asserts otherwise). Wire into tests/CMakeLists.txt following the existing test- target pattern (see 7930df51 for the minimal 9-line addition).
  • tests/test-moe-perf.cpp (140 lines): MoE GEMM microbench for A/B without loading a 90GB+ model. Build-only target; run before/after OPT-03/04/05.
  • Rule: any new arch or kernel gets harness coverage before server runs. Harness green (42/42 CPU + CUDA) is the merge gate.

OPT-08 — Bailing/Ling trapped tool-call rescue + strip (server correctness, not perf)

  • Goal: thinking models of the Bailing family (Ling-3.0-Flash) intermittently emit <tool_call> arg_key/arg_value XML inside the thinking block. The PEG layer consumes that region as reasoning_content before the tool stage, so the call is lost (empty content, no structured calls) and agent loops stall into empty-response nudges. Recover well-formed blocks post-parse; strip executed blocks from reasoning (vLLM Ling3 terminator parity: reasoning ends where the call begins).
  • Exempt from the PP/TG gate. Correctness item: validate via HTTP replay below, not benches. Revert is one commit per sub-item (2e84b994 rescue, 253cf76f strip-on-rescue, f9ffb6fb strip-whenever-calls-exist).
  • Locate: examples/server/server-context.cpp, server_slot::update_chat_msg. Anchor: new_msg.set_tool_call_ids(generated_tool_call_ids, gen_tool_call_id); inside if (!new_msg.empty()).
  • Transform:
    1. New file examples/server/parsers/bailing_parser.hpp (namespace bailing): parse_tool_calls(text) (complete <tool_call>...</tool_call> blocks only, bare-identifier names, raw-string args, 8-block cap, never throws), has_complete_block(text), remove_executed_blocks(text) (complete blocks only, truncated fragments survive). Include as <nlohmann/json.hpp>, matching server-common.h, not quoted json.hpp (fails the server TU build).
    2. In update_chat_msg, before the anchor line, insert the rescue: on final (non-partial) parses only, when structured calls are absent and content is empty, promote parsed blocks to new_msg.tool_calls (ids flow through the existing set_tool_call_ids call). Log Bailing rescue: promoted %d tool call(s).
    3. Same place, after the rescue: on any final turn carrying structured calls, strip complete blocks from new_msg.reasoning_content (covers PEG-parsed turns too, which keep block text in reasoning). Log Bailing strip: removed .... Reasoning with no calls is never touched. See git show 2e84b994, 253cf76f, f9ffb6fb for the exact hunks.
  • New arch checklist: any future thinking template whose tool-call markup can land inside its thinking region gets the same treatment: a parsers/<family>_parser.hpp plus the two hooks. If upstream ever ships a dedicated Bailing PEG builder in common_chat_try_specialized_template, re-evaluate whether the post-parse fallback is still needed (keep whichever fires; they are compatible, rescue is dormant when structured calls arrive clean).
  • Verify: standalone harness examples/server/test-bailing-parser.cpp (g++ syntax + run, no full build): well-formed trapped call converts with correct name/args, truncated/garbage/plain-text yield nothing, strip preserves narration and truncated fragments. Live: faithful HTTP replay (44k system prompt + real tool history, 8 rounds) must return structured tool_calls every round with no <tool_call> text in final reasoning_content. Server log must show Bailing rescue/Bailing strip lines on trapped turns only.
  • Known limitation (do not re-litigate without new evidence): streaming partial deltas carry raw XML as generated; the final-message strip cannot retract already-streamed text. Streaming clients that accumulate reasoning from deltas will still display it. Fixed only by client-side display scrub or by the model emitting elsewhere (weights).

Validation protocol (every item)

  1. Build (glove-fit, canonical):
    rm -rf build-ampere-zen3
    cmake -S . -B build-ampere-zen3 \
      -DCMAKE_BUILD_TYPE=Release \
      -DGGML_NATIVE=ON \
      -DGGML_CUDA=ON \
      -DCMAKE_CUDA_ARCHITECTURES="86-real" \
      -DGGML_CUDA_FA_ALL_QUANTS=ON \
      -DGGML_IQK_MUL_MAT=ON
    cmake --build build-ampere-zen3 --config Release -j$(nproc)
    
    Do not add GGML_CPU_ALL_VARIANTS / GGML_CPU_ARM_ARCH (nonexistent), manual -march on top of GGML_NATIVE=ON, or raw -gencode flags.
  2. Harnesses: build/bin/test-delta-chunkPASS 42/42; build/bin/test-delta-chunk --cudaPASS 42/42.
  3. Micro: build/bin/test-moe-perf before/after (≥3 runs, take median).
  4. Server A/B: llama-bench -m <model> -p 4096 -n 0 -ts 4/1 for PP, -p 128 -n 128 for TG; then the real 25k session flags. Accept: PP +2% or TG no-regression; revert anything else.
  5. Correctness: llama-perplexity spot-check ≤0.001 drift after any graph/loader change.

Known non-goals (evaluated, deferred — do not redo without new evidence)

  • QSA sparse gather kernel: needs a new gather kernel; expected gain <0.5% at real n_kv. Revisit only with profiler data showing QSA >5% of PP time.
  • HC fusion: ~0.1% of PP. Skip.
  • --fit + --n-cpu-moe: mutually exclusive in src/llama.cpp (explicit error). Do not "fix" — use --no-mmap + manual placement.
  • Bench -ts syntax: 4/1 (slash) for multi-GPU in llama-bench; comma means separate runs. Server uses comma. Do not "unify".