feat(mtp): Phases 2.2 + 3.5 + 5.5 — MoE, gated attn, KV cache + softmax scan; bugs fixed → 0% → ~30% accept rate

[kekzl]

Summary

Completes the MTP transformer block (Phase 2.2 = MoE + Attn-with-KV-cache + gated-V-output), wires auto-invoke telemetry (Phase 3.5), and validates end-to-end (Phase 5.5). TWO root-cause bugs found and fixed — accept rate jumped from 0/4-at-0% to 4/4-with-signal (avg 29.5%).

Phase 5.5 validation (Qwen3.6-NVFP4, --mtp-spec-decode 1, max_tokens=128)

class	matches	total	rate
factual	22	67	32.8%
verbose-think	23	126	18.3%
code	47	126	37.3%
instruction	37	126	29.4%
average			29.5%

All 4 classes show real signal. Below the ≥ 60%-on-3/4 default-on threshold but far above noise. RoPE on Q/K is the next quality lever.

Root-cause bugs (latest commit 166f3fa)

Bug 1: RMSNorm 1D-shape early-return

imp::rmsnorm reads x.shape[0] as rows and x.shape[1] as d_model. mtp_forward.cu (since PR #172 Phase 2.1) passed 1D tensors [hidden_dim] → kernel saw rows=hidden_dim, d_model=0 and early-returned without writing output. The MTP forward's RMSNorm outputs were uninitialized FP16 buffer contents (often saturated ~22000). LM-head argmax locked deterministically to token 6178 ('awn') regardless of input.

Fix: 4 sites in mtp_forward.cu changed from [hidden_dim] → [1, hidden_dim].

Bug 2: Missing arch_norm_offset on MTP norms

Qwen3.5/3.6 SafeTensors stores RMSNorm gammas as deltas W where actual gamma = 1 + W. Main-model loader applies the +1 via ctx.arch_norm_offset. upload_mtp_weights() used upload_unquantized_weight() which doesn't expose the offset → MTP norms ran with scale ≈ 0 (raw W with mean near zero).

Fix: dispatch the 7 norm tensors (pre_fc_norm_{embedding,hidden}, input_layernorm, post_attention_layernorm, q_norm, k_norm, final_norm) through upload_weight(..., weight_offset=ctx.arch_norm_offset).

Per-phase shipped

Sub-phase	Status
2.2.MoE — 256-expert top-8 + shared expert + sigmoid gate	✅
2.2.Attn MVP — attn_output_gate=True gated V-broadcast (no KV)	✅
2.2.Attn+KV — per-session KV cache + softmax attention scan (NEW)	✅
3.5 — auto-invoke + accuracy telemetry	✅
5.5 — validation harness + finding	✅
Bug fixes — RMSNorm shape + arch_norm_offset (NEW)	✅

Phase 2.2.Attn+KV details

• Per-session K and V cache (MtpDraftWorkspace::d_k_cache + d_v_cache) up to 16K context — 16 MiB each for Qwen3.6 dims (max_seq_len × num_kv_heads × head_dim × 2 bytes).
• mtp_kv_append_kernel: appends current step's k/v at position mtp_pos.
• mtp_attn_kv_scan_kernel: one CTA per Q-head, softmax over [0, mtp_pos+1) with shared-mem max-reduce, GQA broadcast from kv_h, scaled by 1/√head_dim.
• mtp_gate_attn_out_kernel: applies silu(gate) elementwise to attention output.
• Engine::mtp_accuracy_reset() resets mtp_pos on imp_context_reset for clean new-session state.
• No RoPE yet — content-only attention. Adding RoPE (imp::qknorm_rope_fused already exists for main model with partial-rope + mrope) should close roughly half the gap to DeepSeek-V3 paper's ~85% expectations.

What's still placeholder

• RoPE on Q/K: Qwen3.6 uses partial-rope 0.25 (rope_dim=64) + mrope sections [11, 11, 10] for multimodal. Adding this should bump factual/instruction acceptance toward the 60% default-on threshold.
• K=2+ MTP chaining: each draft step is independent right now. Multi-step chaining is a trivial extension once telemetry shows good single-step quality.
• Phase 3.5 batched-verify proper: with the verify forward running [prev_token, draft_0, ..., draft_{K-1}] as a batched prefill and accept-prefix logic. Only worth implementing after RoPE pushes acceptance to default-on territory.

Validation

• MtpForwardTest.DraftStepProducesValidToken: PASS (full MoE + Attn+KV path engaged).
• make verify-fast: green (decode +1.89× graph speedup, smoke 'Paris' check passes).
• Manual: imp-cli --mtp-spec-decode 1 produces identical tokens with/without MTP (telemetry remains non-behavioral).
• scripts/mtp_accuracy_bench.sh: 4/4 classes with signal (above).

Files changed

src/runtime/mtp_forward.{h,cu}    + KV cache (workspace fields, alloc/free, append/scan/gate kernels)
                                  + arch fix (1D → 2D rmsnorm shapes)
src/runtime/engine.{h,cpp}        + accuracy telemetry + async-graph-loop gate
                                  + mtp_accuracy_reset() resets KV pos
src/api/imp_api.cpp               + imp_context_reset clears MTP state
src/model/weight_upload.cu        + arch_norm_offset on MTP norm uploads (root-cause fix #2)
src/model/hf_config_loader.cpp    + read shared_expert_intermediate_size for Qwen3.5/3.6
tools/imp-cli/main.cpp            + mtp accept-rate at end-of-generation log line
scripts/mtp_accuracy_bench.sh     + Phase 5.5 harness with 3-way verdict (default-on / try-batched-verify / blocked)
tests/test_mtp_forward.cpp        + Phase 2.2 integration test
CMakeLists.txt                    + register test_mtp_forward
docs/superpowers/plans/2026-05-14-mtp-phase2-onwards.md
                                  + Phase 2.2 + 3.5 + 5.5 task-by-task plan

Memory: mtp_phase5_validation_2026_05_14 — re-run recipe + next-step recommendations.

🤖 Generated with Claude Code