diff --git a/mtplx/generation.py b/mtplx/generation.py index 5ea8163c2..dfd9cd29f 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -4026,22 +4026,92 @@ def _sample_from_logits( return sample_from_distribution(probs, rng), probs -def _greedy_draft_token_and_top2(logits: mx.array) -> tuple[int, float, float]: - """Materialize one greedy token and its FP32 top-two values together.""" +def _greedy_draft_token_and_top_values( + logits: mx.array, + *, + topk: int, +) -> tuple[int, np.ndarray]: + """Materialize greedy argmax and FP32 top-k values with one synchronization.""" - row = ( - logits[:, -1, :][0] - if logits.ndim == 3 - else logits.reshape(-1) - ).astype(mx.float32) + row = logits[:, -1, :][0] if logits.ndim == 3 else logits.reshape(-1) + # Keep argmax separate because top-k indices need not preserve first-index ties. token_id = mx.argmax(row, axis=-1) - top2_values = mx.topk(row, k=2) - _eval(token_id, top2_values) + k = max(2, min(int(topk), int(row.shape[-1]))) + top_values = mx.topk(row.astype(mx.float32), k=k) + _eval(token_id, top_values) token = int(np.asarray(token_id).reshape(-1)[0]) - top2 = np.asarray(top2_values, dtype=np.float32).reshape(-1) + values = np.asarray(top_values, dtype=np.float32).reshape(-1) + return token, values + + +def _greedy_draft_token_and_top2(logits: mx.array) -> tuple[int, float, float]: + """Materialize one greedy token and its FP32 top-two values together.""" + + token, top2 = _greedy_draft_token_and_top_values(logits, topk=2) return token, float(top2[-1]), float(top2[-2]) +def _confidence_metrics_from_top_values( + top_values: mx.array | np.ndarray, +) -> dict[str, float]: + values = np.sort(np.asarray(top_values, dtype=np.float32).reshape(-1)) + if values.size < 2: + return {"top2_margin": 0.0, "top1_prob_topk": 1.0, "entropy_topk": 0.0} + descending = values[::-1].astype(np.float64) + shifted = descending - float(descending[0]) + exp_values = np.exp(shifted) + probabilities = exp_values / float(np.sum(exp_values)) + entropy = -float( + np.sum(probabilities * np.log(np.maximum(probabilities, 1e-30))) + ) + return { + "top2_margin": float(values[-1] - values[-2]), + "top1_prob_topk": float(probabilities[0]), + "entropy_topk": entropy, + } + + +def _greedy_draft_token_and_metrics( + logits: mx.array, + *, + need_distribution: bool, + topk: int = 8, +) -> tuple[int, SparseDistribution | None, dict[str, float]]: + """Build the greedy proposal and its confidence metrics with one sync.""" + + row = logits[:, -1, :][0] if logits.ndim == 3 else logits.reshape(-1) + token, top_values = _greedy_draft_token_and_top_values(row, topk=topk) + distribution = ( + SparseDistribution.one_hot(token, int(row.shape[-1])) + if need_distribution + else None + ) + return token, distribution, _confidence_metrics_from_top_values(top_values) + + +def _can_combine_greedy_draft_read( + draft_sampler: SamplerConfig, + *, + draft_margin_threshold: float | None, + adaptive_width_policy: Any | None, + target_prefix_route: Any | None, + correction_cache_enabled: bool, + adapter_ensemble_q: bool, + mtp_topk_reranker: Any | None, +) -> bool: + """Limit the combined path to the default greedy draft selector.""" + + return bool( + draft_sampler.temperature <= 0 + and draft_margin_threshold is None + and adaptive_width_policy is None + and target_prefix_route is None + and not correction_cache_enabled + and not adapter_ensemble_q + and mtp_topk_reranker is None + ) + + def _sample_draft_from_logits( logits: mx.array, config: SamplerConfig, @@ -4561,19 +4631,7 @@ def _draft_confidence_metrics(logits: mx.array, *, topk: int = 8) -> dict[str, f k = max(2, min(int(topk), int(logits.shape[-1]))) top_values = mx.topk(logits.astype(mx.float32), k) _eval(top_values) - values = np.sort(np.asarray(top_values, dtype=np.float32).reshape(-1)) - if values.size < 2: - return {"top2_margin": 0.0, "top1_prob_topk": 1.0, "entropy_topk": 0.0} - descending = values[::-1].astype(np.float64) - shifted = descending - float(descending[0]) - exp_values = np.exp(shifted) - probs = exp_values / float(np.sum(exp_values)) - entropy = -float(np.sum(probs * np.log(np.maximum(probs, 1e-30)))) - return { - "top2_margin": float(values[-1] - values[-2]), - "top1_prob_topk": float(probs[0]), - "entropy_topk": entropy, - } + return _confidence_metrics_from_top_values(top_values) def _top2_margin(logits: mx.array) -> float: @@ -7816,6 +7874,17 @@ def emit_new_tokens() -> None: # continuation predictiveness and can cost more to verify than they commit, # while grounded re-emission matches into the prompt (see the PR benchmarks). ccopy_index.sync(prompt_ids) + combine_greedy_draft_read = _can_combine_greedy_draft_read( + draft_sampler, + draft_margin_threshold=draft_margin_threshold, + adaptive_width_policy=adaptive_width_policy, + target_prefix_route=a3b_target_prefix_route, + correction_cache_enabled=( + online_correction_cache or prompt_correction_cache + ), + adapter_ensemble_q=adapter_ensemble_q, + mtp_topk_reranker=mtp_topk_reranker, + ) # Close the pre-first-token setup span here: everything from the # restore/prefill return to this point (prompt-prefix bank commit, # graphbank/policy/sampler construction) is setup wall time that @@ -8650,11 +8719,21 @@ def emit_new_tokens() -> None: wants_policy_metrics = bool( getattr(adaptive_policy, "wants_draft_metrics", False) ) - draft_metrics = ( - _draft_confidence_metrics(draft_logits[:, -1, :][0]) - if draft_margin_threshold is not None or wants_policy_metrics - else {} - ) + prepared_greedy_draft: tuple[int, SparseDistribution | None] | None = None + if wants_policy_metrics and combine_greedy_draft_read: + draft_token, draft_q, draft_metrics = _greedy_draft_token_and_metrics( + draft_logits, + need_distribution=( + sampler.temperature > 0 and not target_prefix_verify + ), + ) + prepared_greedy_draft = (draft_token, draft_q) + else: + draft_metrics = ( + _draft_confidence_metrics(draft_logits[:, -1, :][0]) + if draft_margin_threshold is not None or wants_policy_metrics + else {} + ) margin = draft_metrics.get("top2_margin") if ( draft_margin_threshold is not None @@ -8774,14 +8853,15 @@ def emit_new_tokens() -> None: need_draft_distribution = ( sampler.temperature > 0 and not target_prefix_verify ) - draft_token, draft_q, adaptive_width_stop = ( - cycle_draft_reader( + if prepared_greedy_draft is not None: + draft_token, draft_q = prepared_greedy_draft + else: + draft_token, draft_q, adaptive_width_stop = cycle_draft_reader( draft_logits, depth_index=depth_index, need_distribution=need_draft_distribution, decision_margins=adaptive_width_decision_margins, ) - ) elapsed_draft = time.perf_counter() - started draft_time += elapsed_draft if trace.enabled: diff --git a/tests/test_expected_value_greedy.py b/tests/test_expected_value_greedy.py new file mode 100644 index 000000000..af44498d7 --- /dev/null +++ b/tests/test_expected_value_greedy.py @@ -0,0 +1,273 @@ +"""Greedy draft contracts used by expected-value depth selection.""" + +from __future__ import annotations + +import math + +import mlx.core as mx +import numpy as np +import pytest + +import mtplx.generation as generation +from mtplx.adaptive import ExpectedValueDepthPolicy +from mtplx.generation import ( + _can_combine_greedy_draft_read, + _draft_confidence_metrics, + _greedy_draft_token_and_metrics, + _sample_draft_from_logits, +) +from mtplx.sampling import SamplerConfig, SparseDistribution + + +@pytest.fixture(autouse=True) +def _cpu_device(): + previous = mx.default_device() + mx.set_default_device(mx.cpu) + try: + yield + finally: + mx.set_default_device(previous) + + +@pytest.mark.parametrize("dtype", [mx.float16, mx.float32]) +@pytest.mark.parametrize( + "values", + [ + [1.0, 4.0, 3.25, -2.0, 0.5, 2.0, -1.0, 0.0, 1.5], + [2.0, 2.0, 1.0, -1.0, 0.0, 0.5, 1.5, -3.0, 0.25], + ], +) +def test_confidence_metrics_match_fp32_top8(values, dtype): + logits = mx.array(values, dtype=dtype) + + metrics = _draft_confidence_metrics(logits, topk=8) + + rounded = np.asarray(logits.astype(mx.float32), dtype=np.float32) + top_values = np.sort(rounded)[-8:] + shifted = top_values[::-1].astype(np.float64) - float(top_values[-1]) + exp_values = np.exp(shifted) + probabilities = exp_values / float(np.sum(exp_values)) + expected_entropy = -float( + np.sum(probabilities * np.log(np.maximum(probabilities, 1e-30))) + ) + + assert metrics["top2_margin"] == pytest.approx( + float(top_values[-1] - top_values[-2]) + ) + assert metrics["top1_prob_topk"] == pytest.approx(float(probabilities[0])) + assert metrics["entropy_topk"] == pytest.approx(expected_entropy) + + +@pytest.mark.parametrize("need_distribution", [False, True]) +def test_greedy_draft_reader_preserves_first_maximum(need_distribution): + logits = mx.array([1.0, 4.0, 3.25, -2.0, 4.0], dtype=mx.float16) + + token, distribution = _sample_draft_from_logits( + logits, + SamplerConfig(temperature=0.0, top_p=1.0, top_k=0), + np.random.default_rng(7), + need_distribution=need_distribution, + ) + + assert token == 1 + if need_distribution: + assert isinstance(distribution, SparseDistribution) + assert distribution.vocab_size == 5 + np.testing.assert_array_equal(distribution.token_ids, np.array([1])) + np.testing.assert_array_equal(distribution.probs, np.array([1.0])) + else: + assert distribution is None + + +def test_confidence_metrics_preserve_expected_value_decision(): + metrics = _draft_confidence_metrics( + mx.array([0.0, 1.0, 3.5, 2.0, -1.0, 0.25, 0.5, 1.5]), + topk=8, + ) + policy = ExpectedValueDepthPolicy( + max_depth=3, + base_depth=2, + warmup_full_depth_cycles=0, + exploration_interval=0, + ) + policy._attempt_counts[2] = 1 + policy._cycles_observed = 1 + + decision = policy.should_continue_after_draft( + drafted_depth=2, + max_depth=3, + draft_metrics=metrics, + ) + + assert math.isfinite(float(decision["confidence_factor"])) + assert decision["drafted_depth"] == 2 + assert decision["next_depth"] == 3 + assert decision["reason"] in {"ev_pass", "ev_fail"} + + +@pytest.mark.parametrize("need_distribution", [False, True]) +def test_combined_greedy_read_matches_separate_operations(need_distribution): + logits = mx.array( + [1.0, 4.0, 3.25, -2.0, 4.0, 2.0, -1.0, 0.0, 1.5], + dtype=mx.float16, + ) + sampler = SamplerConfig(temperature=0.0, top_p=1.0, top_k=0) + expected_token, expected_distribution = _sample_draft_from_logits( + logits, + sampler, + np.random.default_rng(7), + need_distribution=need_distribution, + ) + expected_metrics = _draft_confidence_metrics(logits) + + token, distribution, metrics = _greedy_draft_token_and_metrics( + logits.reshape(1, 1, -1), + need_distribution=need_distribution, + ) + + assert token == expected_token == 1 + assert metrics == expected_metrics + if need_distribution: + assert isinstance(distribution, SparseDistribution) + np.testing.assert_array_equal( + distribution.token_ids, + expected_distribution.token_ids, + ) + np.testing.assert_array_equal( + distribution.probs, + expected_distribution.probs, + ) + else: + assert distribution is expected_distribution is None + + +def test_combined_greedy_read_uses_one_sync(monkeypatch): + original_eval = generation._eval + evaluations = [] + + def audited_eval(*values, **kwargs): + evaluations.append(values) + return original_eval(*values, **kwargs) + + monkeypatch.setattr(generation, "_eval", audited_eval) + + token, _, _ = _greedy_draft_token_and_metrics( + mx.array([[[1.0, 4.0, 3.25, -2.0]]], dtype=mx.float16), + need_distribution=False, + ) + + assert token == 1 + assert len(evaluations) == 1 + assert len(evaluations[0]) == 2 + assert evaluations[0][0].ndim == 0 + assert tuple(evaluations[0][1].shape) == (4,) + assert evaluations[0][1].dtype == mx.float32 + + +def _combined_read_kwargs() -> dict: + return { + "draft_margin_threshold": None, + "adaptive_width_policy": None, + "target_prefix_route": None, + "correction_cache_enabled": False, + "adapter_ensemble_q": False, + "mtp_topk_reranker": None, + } + + +def test_combined_greedy_read_accepts_default_expected_value_path(): + assert _can_combine_greedy_draft_read( + SamplerConfig(temperature=0.0, top_p=1.0, top_k=0), + **_combined_read_kwargs(), + ) + + +@pytest.mark.parametrize( + ("option", "value"), + [ + ("draft_margin_threshold", 0.5), + ("adaptive_width_policy", object()), + ("target_prefix_route", object()), + ("correction_cache_enabled", True), + ("adapter_ensemble_q", True), + ("mtp_topk_reranker", object()), + ], +) +def test_combined_greedy_read_rejects_alternate_selectors(option, value): + options = _combined_read_kwargs() + options[option] = value + + assert not _can_combine_greedy_draft_read( + SamplerConfig(temperature=0.0, top_p=1.0, top_k=0), + **options, + ) + + +def test_combined_greedy_read_requires_greedy_draft_sampler(): + assert not _can_combine_greedy_draft_read( + SamplerConfig(temperature=0.7, top_p=0.9, top_k=20), + **_combined_read_kwargs(), + ) + + +def test_expected_value_generation_preserves_control_trace(monkeypatch): + from test_generation_sustained import AcceptingTinyMTPModel, _runtime + + monkeypatch.setenv("MTPLX_COMPILED_VERIFY", "off") + monkeypatch.setenv("MTPLX_CONTEXT_COPY", "0") + + def run_once(): + return generation.generate_mtpk( + _runtime(AcceptingTinyMTPModel()), + [0], + max_tokens=8, + sampler=SamplerConfig(temperature=0.0, top_p=1.0, top_k=0), + speculative_depth=3, + verify_strategy="batched", + mtp_history_policy="cycle", + stop_token_ids=set(), + adaptive_policy=ExpectedValueDepthPolicy(max_depth=3, base_depth=2), + ) + + original_combined_read = generation._greedy_draft_token_and_metrics + combined_calls = 0 + + def audited_combined_read(*args, **kwargs): + nonlocal combined_calls + combined_calls += 1 + return original_combined_read(*args, **kwargs) + + monkeypatch.setattr( + generation, + "_greedy_draft_token_and_metrics", + audited_combined_read, + ) + combined = run_once() + assert combined_calls > 0 + + monkeypatch.setattr( + generation, + "_can_combine_greedy_draft_read", + lambda *_args, **_kwargs: False, + ) + control = run_once() + + def policy_trace(output): + return [ + { + "depth": draft["depth"], + "token": draft["token"], + "top2_margin": draft.get("top2_margin"), + "top1_prob_topk": draft.get("top1_prob_topk"), + "entropy_topk": draft.get("entropy_topk"), + "policy_continue": draft.get("policy_continue"), + } + for event in output.stats.events + for draft in event.get("drafts", []) + ] + + assert combined.tokens == control.tokens + assert combined.stats.drafted_by_depth == control.stats.drafted_by_depth + assert combined.stats.accepted_by_depth == control.stats.accepted_by_depth + assert combined.stats.rejected_drafts == control.stats.rejected_drafts + assert policy_trace(combined) == policy_trace(control)