Skip to content

fix(mtp): harden transactional cache, logprobs, and guards - #2

Merged
AirRunner merged 1 commit into
AirRunner:feat/mtp-nativefrom
PhilipJohnBasile:fix/mtp-native-transactional-hardening
Sep 1, 2026
Merged

AirRunner merged 1 commit into
AirRunner:feat/mtp-nativefrom
PhilipJohnBasile:fix/mtp-native-transactional-hardening

Conversation

@PhilipJohnBasile

Copy link
Copy Markdown

Carry the selected correctness work from ml-explore#1740 onto AirRunner/mlx-lm:feat/mtp-native at e8ceecc.

This preserves transactional multi-turn prompt-cache alignment and finalization, target/verifier log probabilities for accepted drafts, usable-head detection, fail-closed logits-processor handling, and the focused regression coverage.

Source: PhilipJohnBasile/mlx-lm@bc1d114

@AirRunner AirRunner left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the PR!

Went through the diff, this is in good shape (MTPPromptCacheState boundary mechanism, finalize_plan bookkeeping, fail-closed guards). The verify_lp vs draft_lp fix on accept matches what we discussed.


One thing I want to note rather than just merge past: in qwen3_5.py::sanitize(), the case where mtp_num_hidden_layers > 0 but no MTP weights are present used to raise a hard ValueError. That error exists because of a report from @layer4down early on (a mlx-community quantization had mtp_num_hidden_layers: 1 in its config with no MTP tensors actually present). I see that your version replaces the hard error with a silent del self.mtp / supports_mtp = False fallback.

Actually it's not fully silent, the existing --mtp flag ignored warning in stream_generate still fires when someone actually requests --mtp on such a checkpoint, since _model_supports_mtp reads supports_mtp. So the model isn't unusable anymore, which is actually better than before (used to blocked loading the model at all, even for plain non-MTP).

However what's lost is the specific diagnosis: the old message told the user their config and weights disagree, which is what let @layer4down figure out the checkpoint was mis-packaged. The new warning just says "no MTP head," same message as a checkpoint that never had MTP support to begin with.

So I'm fine with it, but do we just want the warning message to keep the distinction or not? Curious what you think before we merge this.

@AirRunner

Copy link
Copy Markdown
Owner

Note: I just rebased onto upstream, so this PR's diff now shows unrelated noise on top of your change.

I can't push to your branche despite maintainerCanModify=true (probably due to presence of GitHub Actions in mlx-lm or something). So if you want a clean diff you can rebase, you'll get two small conflicts in mlx_lm/generate.py, related to prompt_progress_callback wiring.

Rebase the selected correctness work onto AirRunner's current feat/mtp-native, preserve prompt_progress_callback wiring, and retain a specific warning when checkpoint config declares MTP but the tensors are missing.

Source: bc1d114
Signed-off-by: Philip John Basile <PBasile@Basilecom.com>
@github-actions
github-actions Bot force-pushed the fix/mtp-native-transactional-hardening branch from 476931a to 6630bde Compare August 26, 2026 18:21
@AirRunner
AirRunner merged commit cbb9562 into AirRunner:feat/mtp-native Sep 1, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants