119 lines
9.6 KiB
Markdown
119 lines
9.6 KiB
Markdown
# 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:
|
||
```cpp
|
||
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:**
|
||
```cpp
|
||
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`, 1–2 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 `64` → `32`, 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:
|
||
```cpp
|
||
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):
|
||
```cpp
|
||
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-chunk` → `PASS 42/42`; `build/bin/test-delta-chunk --cuda` → `PASS 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".
|