ik_llama_opt/PORTING-OPTIMIZATIONS.md

9.6 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-07. 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). 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.

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".