fix: Skippy Nemotron MTP loading

[ndizazzo]

Summary

• load optional NextN/MTP tensors for final-stage Skippy model loads and native MTP sidecars
• select the MTP source explicitly: disabled, integrated, or external
• use a plain KV cache for Nemotron-H MoE’s appended dense-attention MTP layer

Root cause

llama_model_params.load_mtp defaults to false, so Skippy omitted Nemotron’s NextN tensors before creating its native MTP context. Once enabled, the hybrid cache path also excluded the appended MTP layer, causing MTP execution to access an invalid KV-cache index.

A terminal stage’s include_output flag only identifies output ownership; it does not select an MTP mode. Inferring MTP from that flag loaded NextN tensors and created an integrated context even when native MTP was disabled. It also blocked an explicitly configured external MTP sidecar from attaching.

Impact

NVIDIA Nemotron 3.5 Lightning can load and execute native MTP speculation through Skippy. Disabled MTP now retains no optional MTP tensors or KV context, while an external sidecar can attach to a target that contains integrated MTP tensors.

Skippy ABI change

ABI version: 0.1.37 → 0.1.38.

• Added enum skippy_mtp_source { DISABLED, INTEGRATED, EXTERNAL }.
• Added enum skippy_mtp_source mtp_source to struct skippy_runtime_config.
• Updated the Rust FFI mirror and regenerated the native API reference. Older native runtimes are intentionally rejected by ABI negotiation.

Validation

• fresh pinned llama.cpp checkout replayed patches 0001–0018 through scripts/prepare-llama.sh pinned
• clean static CPU native build via scripts/build-llama.sh
• independent C11 and C++17 public-header compilation
• cargo check -p skippy-server
• cargo test -p skippy-runtime --lib (71 tests)
• cargo test -p skippy-server --lib (398 tests; model-backed GLM MTP fixtures remain opt-in)
• regenerated and checked the Skippy API reference
• just build
• just-built Metal runtime previously loaded NVIDIA-Nemotron-3.5-Lightning-30B-A3B-Q4_K_M and completed a native-MTP request: 12 drafted, 3 accepted, 0 rejected; 5 generated tokens at 42.38 tok/s

Summary by CodeRabbit

• New Features
  • Added configurable MTP sources: disabled, integrated, or external.
  • Native and externally attached draft models now use the selected MTP source.
  • Added MTP tensor loading for supported model configurations.
  • MTP is disabled by default for existing runtime configurations.
• Bug Fixes
  • Prevented incompatible MTP configurations with external draft models.
• Documentation
  • Updated public API documentation and ABI details for MTP source configuration.

[github-actions]

This pull request is currently a draft. Reviews will not take place until the PR is marked as ready for review.

[coderabbitai]

📝 Walkthrough

Walkthrough

The patch adds disabled, integrated, and external MTP source selection to the Skippy ABI and runtime configuration. Server loading paths propagate the selection, llama.cpp enforces it during model loading, and tests cover all three modes.

Changes

MTP source selection

Layer / File(s)	Summary
ABI and runtime configuration crates/skippy-ffi/src/lib.rs, crates/skippy-runtime/src/..., website/src/docs/pages/skippy-api.md	Adds MtpSource, extends RuntimeConfig, bumps the ABI patch version to 38, defines conversions and defaults, and updates API documentation.
Server source resolution and propagation crates/skippy-server/src/binary_transport/..., crates/skippy-server/src/runtime_state.rs, crates/skippy-server/src/embedded.rs, crates/skippy-server/src/frontend/generation/draft_runner.rs	Resolves MTP mode from stage and draft-model settings, then passes it through runtime launch overrides.
Native and external MTP loading third_party/llama.cpp/patches/0016-*.patch, third_party/llama.cpp/patches/0017-*.patch, third_party/llama.cpp/patches/0018-*.patch	Loads MTP tensors, supports Nemotron-H MoE handling, selects integrated or external sources, and validates external draft attachment.
Validation and explicit runtime configuration crates/skippy-server/src/runtime_state.rs, crates/skippy-bench/src/..., crates/skippy-correctness/src/..., crates/skippy-prompt/src/..., crates/skippy-quantize/src/..., crates/llama-spec-bench/src/main.rs	Adds MTP integration coverage and explicitly configures MTP in benchmark, correctness, prompt, and quantization paths.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant BinaryStage
  participant RuntimeLoader
  participant NativeModelLoader
  participant DraftAttacher
  BinaryStage->>BinaryStage: resolve MTP source
  BinaryStage->>RuntimeLoader: load runtime with override
  RuntimeLoader->>NativeModelLoader: apply selected MTP source
  NativeModelLoader->>DraftAttacher: load or attach MTP tensors

Loading

Possibly related PRs

• Mesh-LLM/mesh-llm#899: Modifies native MTP sidecar creation and loading.
• Mesh-LLM/mesh-llm#1034: Modifies the Skippy runtime configuration ABI.
• Mesh-LLM/mesh-llm#1217: Updates the generated Skippy API declarations.

Suggested labels: experimental

Suggested reviewers: michaelneale, i386

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 39.29% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the primary change: fixing Skippy Nemotron MTP loading and configuration.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch codex/fix-nemotron-mtp-load

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[i386]

Sol high-effort review found one actionable P1. The patch queue applies cleanly to the pinned llama.cpp revision, and the submitted integrated Nemotron MTP smoke test does not cover the configuration paths below.

[coderabbitai]

🧹 Nitpick comments (1)

crates/skippy-quantize/src/mtp_attach.rs (1)

199-205: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Cover the new MTP source contract in the test.

runtime_config now accepts and stores mtp_source, but this test does not inspect config.mtp_source. A regression that drops the new argument would still pass. Assert MtpSource::Disabled here, and add a focused case for MtpSource::External, which is used by the attached draft path on Lines 80-86.

Suggested assertion

         assert!(config.include_output);
+        assert!(matches!(config.mtp_source, MtpSource::Disabled));

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@crates/skippy-quantize/src/mtp_attach.rs` around lines 199 - 205, Extend the
test around runtime_config to assert that config.mtp_source preserves
MtpSource::Disabled. Add a focused test case covering MtpSource::External,
including the attached draft configuration path referenced by the existing test
setup, and verify the stored source value.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@crates/skippy-quantize/src/mtp_attach.rs`:
- Around line 199-205: Extend the test around runtime_config to assert that
config.mtp_source preserves MtpSource::Disabled. Add a focused test case
covering MtpSource::External, including the attached draft configuration path
referenced by the existing test setup, and verify the stored source value.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: b60fff82-0e73-4b5f-b46d-034ab01f0adc

📥 Commits

Reviewing files that changed from the base of the PR and between 0510f1b and 29a95b3.

📒 Files selected for processing (4)

• crates/skippy-prompt/src/prompt_cli/binary_repl.rs
• crates/skippy-prompt/src/prompt_cli/draft.rs
• crates/skippy-prompt/src/prompt_cli/mod.rs
• crates/skippy-quantize/src/mtp_attach.rs