Prompt scoring: top-K per chunk without a full-vocabulary log-softmax

[jvmenen]

Summary

Prompt scoring (/v1/completions with echo: true, max_tokens: 0, logprobs: K) picks its top-K per chunk without building a full-vocabulary float32 log-softmax. Results are bitwise identical; scoring is faster.

Motivation

For every 256-row chunk, score_prompt_logprobs materialised a float32 log-softmax over all 248,320 logits (256 × 248,320 × 4 bytes ≈ 254 MB), a second tensor for the subtraction, and then partitioned the whole vocabulary to find the top-K. Only K values per row and the scored token are ever returned.

Change

• One float32 logsumexp per row.
• Top-K ids through a block-max prefilter: the vocabulary is viewed as 64-logit blocks (248,320 = 3,880 × 64, no padding), the K blocks with the largest maxima are kept (the true top-K must lie in them), and the exact top-K is selected among those candidates.
• float32(logit) - lse only at the K survivors and at the scored token.
• Small helpers in mtplx/generation.py: _row_logsumexp_f32, _logprobs_at, _exact_top_k_ids, _argpartition_top_k, _sorted_top_k, _TOP_K_PREFILTER_BLOCK.
• Values are bitwise those of the full log-softmax. Ids are the same except where the K-th value is an exact tie; exact ties now list by ascending token id instead of an unspecified order.

Evidence

Synthetic, one 256 × 248,320 bf16 chunk, M5 Pro: 47 ms → 3.4–5.2 ms on the GPU (K 1 to 128), 420 ms → 150–170 ms on the CPU.

Real model: M5 Pro 64 GB, Qwen3.6-35B-A3B MTPLX Optimized-Balance, profile turbo, depth 2, fan mode default, logprobs: 20, fresh server per run, 2026-09-26, against 2.12.0 (1de2b1c). 240 prompts of 250–840 tokens plus ~2k, ~4k and ~8k token prompts:

	2.12.0	this PR
Scored positions bitwise equal		121,263 / 121,263 (largest logprob difference 0)
p50, prompts < 512 tokens (n=188)	536 ms	432 ms
p50, 512–1,023 tokens (n=52)	630 ms	521 ms
~2k tokens	2.13 s	1.82 s
~4k tokens	4.28 s	3.66 s
~8k tokens	8.84 s	7.55 s
Errors	0	0

Tests

• New tests/test_prompt_scoring_topk.py (25 tests): parity with the full log-softmax on synthetic bf16/f16 logits (normal, peaked, Gumbel-like), exact ties, -inf, K = 0, 1 and 128, CPU and GPU.
• Scoring, server and test_no_mlx_imports tests: 539 passed, 1 skipped.
• python -m build and scripts/fresh_venv_smoke.sh pass. Ruff: no new findings.

🤖 Generated with Claude Code