feat(mtp): Phases 1.B-5 end-to-end MTP scaffolding#172
Merged
Conversation
…lding Builds on Phase 1.A (PR #171 detection) to ship a complete MTP scaffolding stack: tensors load to GPU, reduced forward kernel runs, engine API exists, CLI flag works, model produces output without crashes on Qwen3.6-NVFP4. ## Per-phase deliverables **Phase 1.B — Tensor loading** - `MtpHead` struct expanded from metadata-only to 19 named Tensor fields (pre_fc_norm_*, fc, input/post_attention_layernorm, q/k/v/o_proj, q/k_norm, router, experts_gate_up_packed, experts_down_packed, shared_expert_*, final_norm) - `safetensors_loader::load_safetensors` runs a separate load pass on `model_mtp.safetensors` after main load, dispatches the 19 tensors to MtpHead fields by name, retains the mmap via `Model::split_mmaps_` - Translation NOT applied to MTP names (literal `mtp.*` preserved) **Phase 1.C — Storage decision** - BF16 retained on disk, converted to FP16 on GPU upload (matches main weights path). NVFP4 quant deferred — 1.6 GB FP16 cost is acceptable on a 32 GB GPU running a 35 B model. **Phase 2.1 — Reduced forward kernel** (`src/runtime/mtp_forward.{cu,h}`) - `mtp_draft_step()`: emb → pre_fc_norm × 2 → concat → fc → final_norm → lm_head → argmax - Workspace alloc/free helpers (`MtpDraftWorkspace`) - **Phase 2.1 limitation**: transformer block (attention + 256-expert MoE) is SKIPPED. Compute is a passthrough of fc_out. Production correctness requires the full block (Phase 2.2 — genuinely multi-week to write the MoE forward from scratch). Acceptance rate will be far below trained optimum until Phase 2.2 lands. **Phase 3 — Engine API** - `Engine::enable_mtp_spec_decode(int k)` + `Engine::mtp_draft_one(...)` - `Engine::mtp_ws_storage_` field (type-erased to avoid header include) - Workspace allocated on enable, freed on destroy - Decode-loop auto-invocation deferred to Phase 3.5/Phase 4 production work **Phase 4 — CLI + C API** - `--mtp-spec-decode K` CLI flag - `imp_enable_mtp_spec_decode(ctx, k)` C API entry point - main.cpp calls C API after context creation if flag set **Phase 5 — End-to-end smoke test** (validated, this PR) - Qwen3.6-NVFP4 + `--mtp-spec-decode 2`: - MTP head loads (1.57 GiB, 19 tensors, BF16) - GPU upload succeeds (19 allocations) - Spec-decode enabled (k=2, hidden=2048, vocab=248320, workspace allocated) - Model produces output (125 tok/s decode, no crashes) ## Production gaps documented - Phase 2.2: full transformer block (multi-week) - Phase 3.5: decode-loop auto-invocation - Phase 5.5: acceptance rate measurement (needs 2.2 + 3.5) ## Validation - verify-fast green (decode -0.28%, prefill +1.64%) - Qwen3.6-NVFP4 smoke: end-to-end works without crashing - No production behavior change without explicit `--mtp-spec-decode` flag ## Design spec docs/superpowers/specs/2026-05-14-mtp-wiring-design.md (Phase 1.A PR #171)
github-actions Bot
pushed a commit
that referenced
this pull request
May 14, 2026
…CONFIG setenv (#173) Pre-existing DISABLED_GreedyDeterminism was a known gotcha not an imp bug: cuBLAS on Blackwell sm_120 picks different GEMM algorithms across calls within the same process unless CUBLAS_WORKSPACE_CONFIG=:4096:8 is set BEFORE the cuBLAS handle is created. Greedy decode (temp=0) then diverges due to accumulated FP16 rounding from different algorithms. Fix: set the env var in SetUpTestSuite() before the test fixture creates the engine. Renamed DISABLED_GreedyDeterminism → GreedyDeterminism. Validates that imp itself IS deterministic when given a deterministic GEMM dispatch. ## Secondary DISABLED_DispatchManual (test_attention_fmha_sm120.cu): investigated but could not root-cause in a single session. The kernel produces NaN when called via a direct manual setup but works fine via run_test() with identical Tensor shapes + data patterns. Likely a CUDA-stream / initial- state interaction specific to top-level TEST_F invocation. Kept DISABLED with expanded comment documenting the observed behavior + repro recipe (gtest_also_run_disabled_tests + Q nonzero / O NaN debug prints) for the next debug session. ## Validation - imp-tests --gtest_filter='*GreedyDeterminism*' → PASSES - verify-fast green (decode -2.00%, prefill -0.80%, graphs 1.39×) - All other DegenerationTest cases still pass ## Bug-audit summary (no fix actionable today) - Qwen3.5-27B MXFP4 IMA: model not local, can't reproduce - Mistral-3.2-NVFP4 long-context: model not local, can't reproduce - Gemma-4 Q4_K_M degeneration: deep Q4_K precision issue, Q8_0 workaround documented - DISABLED_BasicHD256 (MXFP4 FMHA): architectural smem limit, kernel optimization required - DISABLED_DispatchManual: investigated, deferred (see comment above) - NVFP4 dequant Stage-2 cuBLAS replacement: multi-day work - Spec-decode self-spec on stock models: conceptual issue, MTP scaffold shipped (PR #172)
github-actions Bot
pushed a commit
that referenced
this pull request
May 14, 2026
* docs(mtp): plan for remaining Phase 2.2 + 3.5 + 5.5 work after PR #172 PR #172 shipped end-to-end MTP scaffolding (load + reduced FC-only forward + engine API + CLI). Three open work items remain for "MTP fully": Phase 2.2 — full transformer block in mtp_forward.cu (currently a no-op passthrough at line 186-190). Design fork documented: Path A (TransformerLayer view-adapter, reuse existing run_attention + run_moe_ffn) vs Path B (from-scratch fused kernels). Path A recommended. Phase 3.5 — auto-invoke mtp_draft_one + verify forward + accept-prefix from the decode loop. Currently mtp_draft_one exists but nothing in step_decode calls it. Phase 5.5 — A/B matrix to decide default-on/off. Task-by-task breakdown for each phase. Cross-references the memory entry mtp_phase2_open_2026_05_14 capturing what's shipped vs open. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * feat(mtp): Phase 2.2 MoE block — 256-expert top-8 + shared expert + sigmoid gate Replaces the no-op Step 5 placeholder in mtp_forward.cu:186-190 with the full MoE branch of the MTP transformer block: Step 5.B.1 post_attention_layernorm(fc_out) → d_post_norm Step 5.B.2 moe_gate_topk_fused: router @ post_norm, softmax, top-k=8 Step 5.B.3 D2H sync of routing indices+weights for host-side dispatch Step 5.B.4 Per chosen expert (k ∈ [0, 8)): gate_up = experts_gate_up_packed[idx] @ post_norm → [1024] act = silu(gate) * up → [512] down = experts_down_packed[idx] @ act → [2048] store into d_expert_outputs[k * hidden] Step 5.B.5 moe_weighted_sum_residual: fc_out += Σ w[k] * out[k] Step 5.B.6 shared expert: silu(gate_proj·x) * (up_proj·x) → down_proj scaled by sigmoid(shared_expert_gate_inp · x), added to fc_out All compute reuses existing imp primitives: - imp::rmsnorm - imp::moe_gate_topk_fused (fused gate-GEMV + softmax + top-k for M=1) - imp::gemm (M=1 GEMV for per-expert weights and shared expert projections) - imp::swiglu (silu(gate) * up) - imp::moe_weighted_sum_residual (Σ + residual) - imp::shared_expert_gate_scale (sigmoid scalar gate in-place) + one tiny new kernel: mtp_add_shared_kernel to fold shared_out into fc_out Per-expert weight handling: experts_gate_up_packed is [256, 1024, 2048] and experts_down_packed is [256, 2048, 512] FP16. For each chosen expert, we build a 2D Tensor view at the expert's slice offset (no extra copies). The 3D packed layout sticks with the shipped MtpHead design. Workspace gains MoE scratch buffers (post_norm, gate_up scratch, act, per-expert outputs, moe_out, shared_*) plus a MoeRoutingBuffers pool and pinned host buffers for the routing D2H. mtp_workspace_allocate gains n_experts / top_k / expert_d_ff / shared_d_ff params so the Engine sizes correctly. The 2-arg form is retained for back-compat. Engine threads model config (256 / 8 / 512 / 512 for Qwen3.6) into the workspace allocator. Also fixes hf_config_loader to read Qwen3.5/3.6's shared_expert_intermediate_size (previously only read DeepSeek's moe_shared_expert_intermediate_size) so expert_shared_d_ff = 512 lands on the config for Qwen3.6-NVFP4. Without this, the MTP shared expert block silently disabled itself. Attention block remains a passthrough (Step 5.A) — Qwen3.6 MTP has unusual attention shapes (q_proj [8192,2048] but o_proj input is 4096) that need upstream-reference investigation. Documented in the header. Smoke test on Qwen3.6-NVFP4 with --mtp-spec-decode 2: workspace allocates cleanly (d_ff_shared=512), main-model decode produces coherent output ("The capital of France is Paris"), verify-fast green (decode +3.23%, prefill +2.31%, graphs 1.72×). The MoE block only RUNS when mtp_draft_one is invoked, which is still manual (Phase 3.5 auto-invoke not yet wired). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * test(mtp): integration test for Phase 2.2 MoE block MtpForwardTest.DraftStepProducesValidToken: - Loads Qwen3.6-NVFP4 + MTP sidecar end-to-end - Allocates MTP workspace with full MoE config (256 experts / top-8 / expert_d_ff=512 / shared_d_ff=512) - Calls mtp_draft_step with a random FP16 hidden state + arbitrary token id - Asserts out_token_id ∈ [0, vocab_size) PASSES on RTX 5090 (14.4s including 1.57 GiB MTP upload), exercising: - router GEMV + top-8 selection - per-expert gate_up + swiglu + down (8 experts dispatched) - moe_weighted_sum_residual - shared expert gate_proj/up_proj/down_proj - sigmoid scalar gate This is the first test that actually invokes the MoE block; existing E2E paths don't auto-call mtp_draft_one (Phase 3.5 deferred). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Builds on PR #171 (Phase 1.A detection) to ship complete MTP scaffolding. Tensors load to GPU, reduced forward kernel runs, engine API exists, CLI flag works, model produces output without crashes on Qwen3.6-NVFP4. Hits all 5 phases of the design spec in scaffolding form.
Per-phase deliverables
Phase 5 validation output
```
imp-cli --model /models/Qwen3.6-35B-A3B-NVFP4 --mtp-spec-decode 2 --bench ...
[INFO] MTP head loaded: ... (1.57 GiB, 19 tensors, BF16)
[INFO] MTP head: uploaded to GPU (19 allocations, 1.57 GiB BF16→FP16)
[INFO] MTP spec-decode enabled (k=2, hidden=2048, vocab=248320, workspace allocated)
pp 64 tokens avg 25.56 ms (2503 tok/s)
tg 4 tokens avg 32.12 ms ( 124 tok/s)
```
Production gaps (documented in spec)
Files changed
```
src/model/mtp_head.h ⬆ expand to 19 named Tensor fields
src/model/safetensors_loader.cpp ⬆ separate MTP load pass + dispatch
src/model/weight_upload.cu ⬆ upload_mtp_weights helper
src/model/model.h ⬆ mtp_info_ → mtp_ (full MtpHead)
src/runtime/mtp_forward.{cu,h} ✨ new — reduced forward kernel
src/runtime/engine.{cpp,h} ⬆ enable_mtp_spec_decode + mtp_draft_one
src/api/imp_api.cpp ⬆ imp_enable_mtp_spec_decode C API
include/imp/imp.h ⬆ public C API decl
tools/imp-cli/{args.h,args.cpp,main.cpp} ⬆ --mtp-spec-decode K flag
CMakeLists.txt ⬆ +mtp_forward.cu
```
Validation
14 files changed, 639 insertions(+), 62 deletions(-)
🤖 Generated with Claude Code