Skip to content

[codex] Add a reserve-only Gemma4 MTP reserve context#23

Draft
nycdubliner wants to merge 2 commits into
AtomicBot-ai:feature/turboquant-kv-cachefrom
nycdubliner:fix-gemma4-mtp-reserve-context
Draft

[codex] Add a reserve-only Gemma4 MTP reserve context#23
nycdubliner wants to merge 2 commits into
AtomicBot-ai:feature/turboquant-kv-cachefrom
nycdubliner:fix-gemma4-mtp-reserve-context

Conversation

@nycdubliner

Copy link
Copy Markdown

What changed

  • add llama_kv_cache_iswa::init_mtp_reserve(llama_ubatch) for shape-only Gemma4 MTP graph reservation
  • switch llama_context::ensure_sched_mtp() to use the reserve-only helper
  • keep real MTP decode on init_mtp(seq_id, ubatch) unchanged
  • update the Gemma4 multislot MTP crash worklog with the follow-up hardening

Why

The original crash fix stopped using memory->init_full() during MTP graph reservation, but it still routed reserve through init_mtp(0, ub). That worked in practice, yet it implicitly treated reserve as if it were decoding real seq_id 0.

This change removes that semantic risk by introducing a dedicated reserve-only memory context with the exact single-stream, single-index topology the MTP graph needs, without depending on any user sequence id or KV state.

Root cause

With -np > 1, reserving the MTP graph through a full KV context built dummy slot info spanning multiple streams. Gemma4 MTP expects a single-stream topology in its reserve and draft paths, so the mismatched reserve shape tripped the ggml_reshape_3d() assert.

Impact

  • preserves the Gemma4 A4B MTP multislot crash fix
  • makes the reserve path explicit and shape-only
  • does not change Qwen MTP / NextN behavior
  • does not change real MTP decode semantics

Validation

  • cmake --build build-hip-rocwmma --target llama-server -j "$(nproc)"
  • CTX=4096 PARALLEL=1 BATCH=128 UBATCH=64 SPLIT_MODE=layer KV_K=turbo4 KV_V=turbo4 REASONING_BUDGET=1024 ENABLE_MTP=1 PORT=8084 NO_WARMUP=1 ~/scripts/local-opencode-llama/scripts/run-gemma4-26b-a4b-mtp.sh
  • CTX=4096 PARALLEL=2 BATCH=128 UBATCH=64 SPLIT_MODE=layer KV_K=turbo4 KV_V=turbo4 REASONING_BUDGET=1024 ENABLE_MTP=0 PORT=8084 NO_WARMUP=1 ~/scripts/local-opencode-llama/scripts/run-gemma4-26b-a4b-mtp.sh
  • CTX=4096 PARALLEL=2 BATCH=128 UBATCH=64 SPLIT_MODE=layer KV_K=turbo4 KV_V=turbo4 REASONING_BUDGET=1024 ENABLE_MTP=1 PORT=8084 NO_WARMUP=1 ~/scripts/local-opencode-llama/scripts/run-gemma4-26b-a4b-mtp.sh
  • tiny /v1/messages request: {"model":"gemma4-26b-a4b-mtp","max_tokens":8,"messages":[{"role":"user","content":"hi"}]}
  • Claude-style /v1/messages request with system + messages
  • /metrics still exposes:
    • llamacpp:speculative_drafts_generated_total{spec_type="mtp"}
    • llamacpp:speculative_draft_tokens_generated_total{spec_type="mtp"}

@github-actions github-actions Bot added documentation Improvements or additions to documentation examples server labels Jun 1, 2026
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…cBot-ai#23

Codex review found:
1. Stale duplicate code in dequantize_turbo3_0_t4 (compile would fail)
2. thread static is risky/non-portable in MSL

Fixed: removed thread static caching, using plain thread locals.
Speed unchanged (2.4 tok/s) — the static caching wasn't actually working
on Metal. True optimization needs architectural change in flash attention
kernel to dequantize once per block, not per chunk.

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…tomicBot-ai#23

Root cause analysis: 8-32× redundant full-block dequantize per block
from flash attention template. Four approaches documented with expected
speedups and risk levels.

Plan: D (reduce overhead) → A/B (eliminate redundant calls)
Target: 2.4 tok/s → 20-40 tok/s

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…icBot-ai#23

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…tomicBot-ai#23

No-op dequant test: even returning all zeros from dequantize, turbo3
runs at 2.4 tok/s (same as with full WHT rotation). The bottleneck is
NOT in the attention dequantize path.

New hypothesis: the SET_ROWS (quantize) path is the bottleneck. The
Metal quantize_turbo3_0 function does 3 WHT rotations per KV write,
totaling ~3200 ops per block × 224 blocks per token.

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…Bot-ai#23

CRITICAL BUG: The #include "turbo-wht.h" caused Metal JIT compilation
to fail at runtime. The model silently fell back to CPU for ALL ops.
ALL previous benchmarks (2.4 tok/s) were measuring CPU, not Metal GPU.

After inlining the header:
- MoE gen: 2.4 → 10.7 tok/s (4.5× improvement, now actually on Metal)
- MoE prompt: 4.2 → 60.9 tok/s (14.5× improvement)

Remaining gap vs q8_0: 85 → 10.7 tok/s (8× slower, down from 35×)

This is the SAME bug we hit with turbo-matrices.h earlier.
Rule: NEVER use #include in ggml-metal.metal — always inline.

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…cBot-ai#23

Previous 2.4 tok/s was CPU fallback. Real Metal numbers:
MoE: 10.7 tok/s gen (8× slower than q8_0, was thought to be 35×)
Qwopus: 5.3 tok/s gen (3.3× slower than q8_0)

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…tomicBot-ai#27

Full investigation log with all tests, results, and the root cause.
Upstream TurboQuant activity tracked in AtomicBot-ai#27.

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…in) AtomicBot-ai#23

Removing WHT rotation from dequant (quality broken, speed test only):
  gen: 10.7 → 49.1 tok/s (4.6× improvement, 57% of q8_0)
  prompt: 67.3 → 162.6 tok/s

Confirms pre-rotate-queries would deliver ~49 tok/s.
Remaining gap (49 vs 85) is block size + QJL overhead.

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…cBot-ai#23

Instead of inverse-rotating every K during dequant, rotate Q once
before attention. Math: <q, R^T*c[idx]> = <R*q, c[idx]>.

Changes:
- Store rotation matrix (R^T) in KV cache, filled after buffer clear
- Apply ggml_mul_mat(R_T, q) in build_attn_mha after permute
- Strip turbo_rotate_inverse from Metal dequant
- Dynamic cast to access rotation from mctx

Results:
- MoE gen: 10.7 → 51.4 tok/s (4.8× speedup)
- MoE prompt: 67.3 → 160.3 tok/s (2.4× speedup)
- Now at 60% of q8_0 speed with 4.9× compression
- Model produces coherent output

Codex review: fixed buffer clear ordering (was zeroing rotation after init).
Verified: rotation point is correct (after 4d reshape + permute, ne[0]=128).

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026
…tomicBot-ai#23

Full investigation log documenting every test, every dead end, and every
breakthrough. 21× total improvement from CPU fallback to pre-rotate-queries.

Key lessons: no #include in Metal, no-op testing, pre-rotate-queries,
buffer clear ordering, codex+roast catch real bugs.

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fukuro-kun pushed a commit to fukuro-kun/fukuro-llama-cpp-turboquant that referenced this pull request Jul 5, 2026


Validated on real Qwen3 KV tensors: cosine sim 0.9508 → 0.9831 (+3.2%)
MSE-only better on 99.3% of vectors including p1 tails.

3-bit index split: lower 2 bits in qs[], upper 1 bit in signs[].
No QJL stage in quantize or dequant.

Results:
- MoE gen: 51.4 → 62.2 tok/s (73% of q8_0, was 60%)
- MoE prompt: 160 → 200 tok/s (90% of q8_0)
- Qwopus gen: 14.6 → 15.5 tok/s (88% of q8_0, was 83%)
- Qwopus prompt: 67 → 83 tok/s (100% of q8_0!)

Codex verified: bit packing correct, quantize/dequant consistent.

Co-Authored-By: tturney@psyguard.ai
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation examples server

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant